import socket
import threading
from typing import Dict, Callable
from protocol import Protocol


class TCPClient:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.socket = None
        self.connected = False
        self.callbacks: Dict[int, Callable] = {}
        self.receive_thread = None
        self.packet_id = 0
        self.debug_callback = None

    def connect(self):
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            print(f"host={self.host},port={self.port}\n")
            self.socket.connect((self.host, self.port))
            self.connected = True
            self._start_receive_thread()
            return True
        except Exception:
            return False

    def disconnect(self):
        self.connected = False
        if self.socket:
            self.socket.close()
            self.socket = None

    def register_callback(self, cmd, callback):
        self.callbacks[cmd] = callback
        print(f"注册回调: cmd=0x{cmd:02x}, callback={callback.__name__}")

    def send_packet(self, cmd, data=b''):
        if not self.connected:
            return False
        try:
            packet = Protocol.build_packet(cmd, data)

            # 记录发送消息（如果有调试回调）
            if self.debug_callback:
                self.debug_callback(cmd, data)

            self.socket.send(packet)
            return True
        except Exception:
            self.connected = False
            return False

    def _start_receive_thread(self):
        self.receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
        self.receive_thread.start()

    def _receive_loop(self):
        #buffer = b''
        while self.connected:
            try:
                data = self.socket.recv(4096)
                if not data:
                    break
                self._process_buffer(data)
            except Exception:
                break

    def _process_buffer(self, buffer):
        while True:
            result = Protocol.parse_packet(buffer)
            if not result:
                break
            header, body = result
            packet_length = len(buffer)

            if header.cmd in self.callbacks:
                print(f"调用回调: cmd=0x{header.cmd:02x}")
                self.callbacks[header.cmd](header, body)
            else:
                print(f"未找到回调: cmd=0x{header.cmd:02x}")
            buffer = buffer[packet_length:]

            packet_length = len(buffer)