import asyncio
from nazo_rand import randrange
from game_server.net.kcp import get_conv
from game_server.net.packet import NetOperation
from game_server.net.session import PlayerSession
from utils.logger import Info,Warn,Debug
from database.mongodb import get_database


class KCPGateway(asyncio.DatagramProtocol):
    def __init__(self, db):
        self.id_counter = 0
        self.sessions : dict[int, PlayerSession] = {}
        self.running = True
        self.shutdown_event = asyncio.Event()
        self.db = db
        self.timeout_check_interval = 5
        self.save_duration = 60

    async def check_sessions_timeout(self):
        while self.running:
            for conv_id, session in list(self.sessions.items()):
                if session.is_timeout():
                    self.drop_kcp_session(conv_id)
            await asyncio.sleep(self.timeout_check_interval)

    async def periodic_save_sessions(self):
        while self.running:
            start = asyncio.get_event_loop().time()

            for session in self.sessions.values():
                session.player.save_all()

            elapsed = asyncio.get_event_loop().time() - start
            Info(f"Database saved in {elapsed:.2f}s")

            await asyncio.sleep(self.save_duration)
        
    def connection_made(self, transport):
        self.transport = transport
        Info(f"Listening on {self.transport.get_extra_info('sockname')}")

        asyncio.create_task(self.check_sessions_timeout())
        asyncio.create_task(self.periodic_save_sessions())

    def datagram_received(self, data, addr):
        data_len = len(data)

        Debug(f"Received {data_len} bytes from {addr}: {[b for b in data]}")

        if data_len == 20:
            self.process_net_operation(NetOperation.from_bytes(data), addr)
        elif data_len >= 28:
            self.process_kcp_payload(data, addr)
        else:
            Warn("Unknown data length received")

    def process_net_operation(self, op: NetOperation, addr):
        if (op.head, op.tail) == (0xFF, 0xFFFFFFFF):
            self.establish_kcp_session(op.data, addr)
        elif (op.head, op.tail) == (0x194, 0x19419494):
            self.drop_kcp_session(op.conv_id)
        else:
            Warn(f"Unknown magic pair: {op.head}-{op.tail}")

    def establish_kcp_session(self, data, addr):
        conv_id, session_token = self.next_conv_pair()
        session_id = conv_id << 32 | session_token

        Info(f"New connection: {addr} with conv_id: {conv_id}")

        self.sessions[conv_id] = PlayerSession(
            self.transport, session_id, addr, self.db
        )

        net_op = NetOperation(
            head=0x145,
            conv_id=conv_id,
            session_token=session_token,
            data=data,
            tail=0x14514545,
        ).to_bytes()

        self.transport.sendto(net_op, addr)

    def drop_kcp_session(self, conv_id):
        if conv_id in self.sessions:
            self.sessions[conv_id].player.save_all()
            del self.sessions[conv_id]
            Info(f"Dropped KCP session with conv_id: {conv_id}")
        else:
            Warn(
                f"Attempted to drop non-existent KCP session with conv_id: {conv_id}"
            )

    def process_kcp_payload(self, data, addr):
        conv_id = get_conv(data)
        session = self.sessions.get(conv_id)

        if session:
            session.update_last_received()
            asyncio.create_task(session.consume(data))
        else:
            Warn(f"Received data for unknown session conv_id: {conv_id}")

    def next_conv_pair(self):
        self.id_counter += 1
        return self.id_counter, randrange(0, 0xFF)

    def connection_lost(self, exc):
        Info("UDP connection lost, shutting down...")

        for session in self.sessions.values():
            session.player.save_all()

        self.running = False
        self.shutdown_event.set()

    def shutdown(self):
        Info("Shutting down server...")

        for session in self.sessions.values():
            session.player.save_all()

        self.running = False
        self.shutdown_event.set()

    @staticmethod
    async def new(host, port):
        loop = asyncio.get_running_loop()
        db = get_database()

        transport, protocol = await loop.create_datagram_endpoint(
            lambda: KCPGateway(db), local_addr=(host, port)
        )

        try:
            await protocol.shutdown_event.wait()
        except asyncio.CancelledError:
            Info("Server tasks cancelled.")
        finally:
            transport.close()
            Info("Server stopped.")