import socket
import threading
import sys

class TyNetworkManager:
    def __init__(self, on_packet_received):
        self.sock = None
        self.on_packet_received = on_packet_received
        self.is_running = False

    def connect(self, host: str, port: int) -> bool:
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((host, port))
            self.is_running = True
            threading.Thread(target=self._read_loop, daemon=True).start()
            return True
        except Exception:
            return False

    def send_packet(self, packet: bytes) -> bool:
        if not self.sock or not self.is_running:
            return False
        try:
            self.sock.sendall(packet)
            return True
        except Exception:
            self.disconnect()
            return False

    def register_id(self, my_id: int):
        header = bytes([0x01, 0x03]) + my_id.to_bytes(3, 'big') + (0).to_bytes(3, 'big')
        meta = b'\x00' * 12 + (0).to_bytes(4, 'big')
        self.send_packet(header + meta)

    def disconnect(self):
        self.is_running = False
        if self.sock:
            try:
                self.sock.close()
            except Exception:
                pass
            self.sock = None

    def _read_loop(self):
        while self.is_running:
            try:
                header = self._recv_exact(8)
                if not header:
                    break
                
                meta = self._recv_exact(16)
                if not meta:
                    break
                
                iv = meta[0:12]
                payload_len = int.from_bytes(meta[12:16], 'big')
                
                payload = b""
                if payload_len > 0:
                    payload = self._recv_exact(payload_len)
                    if not payload:
                        break
                
                self.on_packet_received(header, iv, payload)
                
            except Exception:
                break
        self.disconnect()

    def _recv_exact(self, length: int) -> bytes:
        data = b""
        while len(data) < length:
            chunk = self.sock.recv(length - len(data))
            if not chunk:
                return b""
            data += chunk
        return data