"""
mysql_prepared_statement.py
预处理语句的完整实现
"""
import socket
import struct


def encode_packet(payload: bytes, seq: int) -> bytes:
    return struct.pack('<I', len(payload))[:3] + struct.pack('B', seq) + payload


def read_packet(sock: socket.socket) -> tuple:
    """从 socket 读取一个数据包"""
    header = b''
    while len(header) < 4:
        header += sock.recv(4 - len(header))

    pkt_len = struct.unpack('<I', header[0:3] + b'\x00')[0]
    seq = header[3]

    payload = b''
    while len(payload) < pkt_len:
        payload += sock.recv(pkt_len - len(payload))

    return payload, seq


class PreparedStatement:
    """预处理语句"""
    def __init__(self, statement_id, num_columns, num_params, warnings):
        self.statement_id = statement_id
        self.num_columns = num_columns
        self.num_params = num_params
        self.warnings = warnings
        self.param_definitions = []
        self.column_definitions = []


def send_prepare(sock: socket.socket, sql: str) -> PreparedStatement:
    """发送 COM_STMT_PREPARE 命令"""
    payload = b'\x16' + sql.encode('utf-8')
    sock.send(encode_packet(payload, 0))

    # 读取 Prepare OK
    response, seq = read_packet(sock)

    if response[0] == 0xFF:
        # ERR
        error_code = struct.unpack('<H', response[1:3])[0]
        message = response[9:].decode('utf-8', errors='replace')
        raise Exception(f"Prepare 失败: [{error_code}] {message}")

    statement_id = struct.unpack('<I', response[1:5])[0]
    num_columns = struct.unpack('<H', response[5:7])[0]
    num_params = struct.unpack('<H', response[7:9])[0]
    warnings = struct.unpack('<H', response[10:12])[0]

    stmt = PreparedStatement(statement_id, num_columns, num_params, warnings)
    print(f"[PREPARE] statement_id={statement_id}, columns={num_columns}, params={num_params}")

    # 读取参数定义
    if num_params > 0:
        for i in range(num_params):
            param_data, _ = read_packet(sock)
            # 解析参数定义（简化版）
            stmt.param_definitions.append(parse_param_def(param_data))
        # EOF
        eof, _ = read_packet(sock)

    # 读取列定义
    if num_columns > 0:
        for i in range(num_columns):
            col_data, _ = read_packet(sock)
            stmt.column_definitions.append(parse_col_def(col_data))
        # EOF
        eof, _ = read_packet(sock)

    return stmt


def parse_param_def(data: bytes) -> dict:
    """解析参数定义"""
    # Param Definition 格式与 Column Definition 相同
    # 简化解析
    return {
        'type': data[-5] if len(data) >= 5 else 0,
        'flags': struct.unpack('<H', data[-4:-2])[0] if len(data) >= 5 else 0,
    }


def parse_col_def(data: bytes) -> dict:
    """解析列定义（简化版）"""
    return {'raw': data[:30].hex()}


def send_execute(sock: socket.socket, stmt: PreparedStatement,
                 params: list = None) -> bytes:
    """发送 COM_STMT_EXECUTE 命令"""
    payload = bytearray()

    # 命令类型
    payload.append(0x17)

    # statement_id
    payload.extend(struct.pack('<I', stmt.statement_id))

    # flags (no cursor)
    payload.append(0x00)

    # iteration_count
    payload.extend(struct.pack('<I', 1))

    num_params = stmt.num_params
    if num_params > 0 and params:
        # null_bitmap
        bitmap_size = (num_params + 7) // 8
        null_bitmap = bytearray(bitmap_size)
        for i, param in enumerate(params):
            if param is None:
                null_bitmap[i // 8] |= (1 << (i % 8))
        payload.extend(null_bitmap)

        # new_params_bound_flag
        payload.append(0x01)

        # 参数类型和值
        for param in params:
            if param is None:
                payload.append(0x06)  # MYSQL_TYPE_NULL
                payload.append(0x00)
            elif isinstance(param, int):
                payload.append(0x08)  # MYSQL_TYPE_LONGLONG
                payload.append(0x00)
            elif isinstance(param, str):
                payload.append(0xFD)  # MYSQL_TYPE_VAR_STRING
                payload.append(0x00)
            elif isinstance(param, float):
                payload.append(0x05)  # MYSQL_TYPE_DOUBLE
                payload.append(0x00)
            else:
                payload.append(0xFD)
                payload.append(0x00)

        # 参数值
        for param in params:
            if param is None:
                pass
            elif isinstance(param, int):
                payload.extend(struct.pack('<q', param))
            elif isinstance(param, str):
                encoded = param.encode('utf-8')
                payload.extend(encode_length(len(encoded)))
                payload.extend(encoded)
            elif isinstance(param, float):
                payload.extend(struct.pack('<d', param))
    else:
        payload.extend(b'\x00')  # 空 bitmap
        payload.append(0x00)     # 无新参数

    sock.send(encode_packet(bytes(payload), 0))
    response, seq = read_packet(sock)
    return response


def send_close(sock: socket.socket, statement_id: int):
    """发送 COM_STMT_CLOSE 命令"""
    payload = b'\x19' + struct.pack('<I', statement_id)
    sock.send(encode_packet(payload, 0))
    # COM_STMT_CLOSE 不返回响应


def encode_length(length: int) -> bytes:
    if length < 0xFB:
        return struct.pack('B', length)
    elif length < 0x10000:
        return b'\xFC' + struct.pack('<H', length)
    elif length < 0x1000000:
        return b'\xFD' + struct.pack('<I', length)[:3]
    else:
        return b'\xFE' + struct.pack('<Q', length)

def send_fetch(sock: socket.socket, statement_id: int, num_rows: int) -> bytes:
    """发送 COM_STMT_FETCH 命令"""
    payload = b'\x1C' + struct.pack('<I', statement_id) + struct.pack('<I', num_rows)
    sock.send(encode_packet(payload, 0))
    response, seq = read_packet(sock)
    return response

def send_long_data(sock, statement_id, param_index, data, chunk_size=65535):
    """分批发送大数据参数"""
    offset = 0
    while offset < len(data):
        chunk = data[offset:offset + chunk_size]
        payload = (b'\x18' +
                   struct.pack('<I', statement_id) +
                   struct.pack('<H', param_index) +
                   chunk)
        sock.send(encode_packet(payload, 0))
        offset += chunk_size

def send_stmt_reset(sock, statement_id):
    """重置预处理语句"""
    payload = b'\x1A' + struct.pack('<I', statement_id)
    sock.send(encode_packet(payload, 0))
    response, _ = read_packet(sock)
    return response  # OK 或 ERR

"""
mysql_prepared_crud.py
使用预处理语句执行 CRUD 操作
"""


def prepared_crud_example(sock):
    """预处理语句 CRUD 示例"""

    # === INSERT ===
    stmt = send_prepare(sock,
        "INSERT INTO users (name, age, email) VALUES (?, ?, ?)")
    print(f"INSERT: stmt_id={stmt.statement_id}")

    # 批量插入
    users = [
        ("Alice", 28, "alice@example.com"),
        ("Bob", 35, "bob@example.com"),
        ("Charlie", 22, "charlie@example.com"),
    ]

    for name, age, email in users:
        response = send_execute(sock, stmt, [name, age, email])
        print(f"  插入 {name}: OK")

    send_close(sock, stmt.statement_id)

    # === SELECT ===
    stmt = send_prepare(sock,
        "SELECT id, name, email FROM users WHERE age > ?")
    print(f"\nSELECT: stmt_id={stmt.statement_id}")

    response = send_execute(sock, stmt, [25])
    print(f"  查询结果: {response[:50].hex()}")

    send_close(sock, stmt.statement_id)

    # === UPDATE ===
    stmt = send_prepare(sock,
        "UPDATE users SET email = ? WHERE id = ?")
    print(f"\nUPDATE: stmt_id={stmt.statement_id}")

    response = send_execute(sock, stmt, ["newalice@example.com", 1])
    print(f"  更新结果: OK")

    send_close(sock, stmt.statement_id)

    # === DELETE ===
    stmt = send_prepare(sock,
        "DELETE FROM users WHERE id = ?")
    print(f"\nDELETE: stmt_id={stmt.statement_id}")

    response = send_execute(sock, stmt, [3])
    print(f"  删除结果: OK")

    send_close(sock, stmt.statement_id)

def fetch_large_resultset(sock, sql, params=None, batch_size=1000):
    """使用游标分批获取大结果集"""
    # 预处理
    stmt = send_prepare(sock, sql)
    print(f"预处理完成: {stmt.num_columns} 列, {stmt.num_params} 参数")

    # 启用只读游标执行
    # flags = 0x01 (CURSOR_TYPE_READ_ONLY)
    # 需要修改 send_execute 来支持 flags 参数
    response = send_execute_with_cursor(sock, stmt, params, cursor_type=0x01)

    all_rows = []
    total_fetched = 0

    while True:
        # 检查是否有 SERVER_STATUS_CURSOR_EXISTS (0x0040)
        status = struct.unpack('<H', response[-5:-3])[0] if len(response) >= 5 else 0

        if status & 0x0040:
            # 游标打开中，继续获取
            rows = parse_binary_rows(response, stmt.column_definitions)
            all_rows.extend(rows)
            total_fetched += len(rows)
            print(f"  已获取 {total_fetched} 行...")

            # 获取下一批
            response = send_fetch(sock, stmt.statement_id, batch_size)
        else:
            # 游标关闭（EOF 或 最后一批）
            if len(response) > 5:
                rows = parse_binary_rows(response, stmt.column_definitions)
                all_rows.extend(rows)
                total_fetched += len(rows)
            break

    print(f"总共获取 {total_fetched} 行")

    # 关闭语句
    send_close(sock, stmt.statement_id)

    return all_rows


def send_execute_with_cursor(sock, stmt, params, cursor_type=0):
    """支持游标的 COM_STMT_EXECUTE"""
    # 类似 send_execute，但 flags 参数使用 cursor_type
    pass  # 实现略，参考 send_execute

# SQLAlchemy 示例
from sqlalchemy import text
stmt = text("SELECT * FROM users WHERE id = :id")
result = conn.execute(stmt, {"id": 42})

def batch_insert(conn, table, data, batch_size=1000):
    """使用预处理语句批量插入"""
    placeholders = ", ".join(["?"] * len(data[0]))
    sql = f"INSERT INTO {table} VALUES ({placeholders})"
    stmt = conn.prepare(sql)

    for i in range(0, len(data), batch_size):
        batch = data[i:i + batch_size]
        for row in batch:
            stmt.execute(row)
        conn.commit()

    stmt.close()

-- 存储过程
CREATE PROCEDURE get_user_count(OUT count INT)
BEGIN
    SELECT COUNT(*) INTO count FROM users;
END;