"""
Network Gateway and Frame Codec for FreeExile.
Handles raw binary stream framing (4-byte big-endian length prefix),
packet dispatching, and high-concurrency client multiplexing.
"""

import struct
import asyncio
from typing import Tuple, Optional, Callable


class PacketCodec:
    """
    Standard binary protocol framing:
    [4 Bytes: Payload Length (Big-Endian uint32)] + [N Bytes: Serialized Payload]
    """
    HEADER_SIZE = 4

    @staticmethod
    def encode_frame(payload: bytes) -> bytes:
        length = len(payload)
        header = struct.pack(">I", length)
        return header + payload

    @staticmethod
    def decode_frame(buffer: bytes) -> Tuple[Optional[bytes], bytes]:
        """
        Extracts a single frame from the buffer if complete.
        Returns: (payload_bytes or None, remaining_buffer)
        """
        if len(buffer) < PacketCodec.HEADER_SIZE:
            return None, buffer

        (length,) = struct.unpack(">I", buffer[:PacketCodec.HEADER_SIZE])
        total_frame_len = PacketCodec.HEADER_SIZE + length

        if len(buffer) < total_frame_len:
            return None, buffer  # Frame incomplete, await more bytes

        payload = buffer[PacketCodec.HEADER_SIZE:total_frame_len]
        remainder = buffer[total_frame_len:]
        return payload, remainder


class NetworkGateway:
    def __init__(self, host: str = "0.0.0.0", port: int = 7777):
        self.host = host
        self.port = port
        self.is_running = False
        self.server: Optional[asyncio.Server] = None
        self.packet_handler: Optional[Callable[[bytes], None]] = None

    def set_packet_handler(self, handler: Callable[[bytes], None]) -> None:
        self.packet_handler = handler

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        buffer = b""
        try:
            while self.is_running:
                chunk = await reader.read(4096)
                if not chunk:
                    break
                buffer += chunk

                while True:
                    payload, buffer = PacketCodec.decode_frame(buffer)
                    if payload is None:
                        break
                    if self.packet_handler:
                        self.packet_handler(payload)
        except Exception:
            pass
        finally:
            writer.close()
            await writer.wait_closed()

    async def start(self) -> None:
        self.is_running = True
        self.server = await asyncio.start_server(self.handle_client, self.host, self.port)

    async def stop(self) -> None:
        self.is_running = False
        if self.server:
            self.server.close()
            await self.server.wait_closed()