ForcePilot/backend/package/yuxi/channel/extensions/minecraft/gateway.py

import asyncio
import hashlib
import json
import logging
import random
import struct
from collections.abc import Callable, Awaitable

from yuxi.channel.extensions.minecraft.client import McClient
from yuxi.channel.extensions.minecraft.protocol import (
    write_varint,
    write_string,
    write_long,
    read_string,
    read_varint,
    read_uuid,
)
from yuxi.channel.extensions.minecraft.version_adapter import get_adapter, VersionAdapter, DEFAULT_PROTOCOL
from yuxi.channel.extensions.minecraft.types import ConnectionState, ResolvedMcAccount, McPacket
from yuxi.channel.extensions.minecraft.keepalive import KeepAliveManager
from yuxi.channel.extensions.minecraft.monitor import mc_packet_to_unified

logger = logging.getLogger(__name__)

MAX_RECONNECT_DELAY = 120.0
BASE_RECONNECT_DELAY = 1.0


class McGateway:
    def __init__(self, account: ResolvedMcAccount):
        self.account = account
        self.client = McClient(account.host, account.port)
        self.adapter: VersionAdapter | None = None
        self.keepalive: KeepAliveManager | None = None
        self.on_message: Callable[[object], Awaitable[None]] | None = None
        self._reconnect_task: asyncio.Task | None = None
        self._stop_event = asyncio.Event()
        self._entity_id: int = 0
        self._player_x: float = 0.0
        self._player_y: float = 0.0
        self._player_z: float = 0.0
        self._player_name_cache: dict[str, str] = {}
        self._player_cache: dict[str, dict] = {}

    @property
    def is_connected(self) -> bool:
        return self.client.is_connected and self.client.state == ConnectionState.PLAY

    async def connect(self) -> None:
        self._stop_event.clear()
        await self._do_connect_cycle()
        if self._reconnect_task is None or self._reconnect_task.done():
            self._reconnect_task = asyncio.create_task(self._reconnect_monitor())

    async def _do_connect_cycle(self) -> None:
        await self.client.connect()
        await self._do_handshake()
        await self._do_status_probe()
        await self._do_login()
        await self._enter_play()

    async def disconnect(self) -> None:
        self._stop_event.set()
        if self._reconnect_task:
            self._reconnect_task.cancel()
            self._reconnect_task = None
        if self.keepalive:
            await self.keepalive.stop()
        await self.client.disconnect()

    async def _reconnect_monitor(self) -> None:
        while not self._stop_event.is_set():
            try:
                await asyncio.wait_for(self._stop_event.wait(), timeout=5.0)
                break
            except TimeoutError:
                pass
            if self._stop_event.is_set():
                break
            if not self.client.is_connected:
                logger.info("Minecraft connection lost, attempting reconnect...")
                try:
                    await self.reconnect()
                except Exception:
                    logger.exception("Minecraft reconnect failed")

    async def reconnect(self) -> None:
        delay = BASE_RECONNECT_DELAY
        for attempt in range(self.account.reconnect_max_retries):
            jitter = random.uniform(0, delay * 0.3)
            wait = delay + jitter
            logger.info(
                "Minecraft reconnect attempt %d/%d in %.1fs",
                attempt + 1,
                self.account.reconnect_max_retries,
                wait,
            )
            await asyncio.sleep(wait)
            try:
                await self.connect()
                logger.info("Minecraft reconnected successfully (attempt %d)", attempt + 1)
                return
            except Exception:
                logger.exception("Minecraft reconnect attempt %d failed", attempt + 1)
                delay = min(delay * 2, MAX_RECONNECT_DELAY)

        logger.error("Minecraft reconnect failed after %d attempts", self.account.reconnect_max_retries)

    async def _do_handshake(self) -> None:
        protocol_version = DEFAULT_PROTOCOL
        if self.account.version not in ("auto", ""):
            try:
                protocol_version = int(self.account.version)
            except ValueError:
                protocol_version = DEFAULT_PROTOCOL

        packet_id_bytes = write_varint(0x00)
        data = (
            write_varint(protocol_version)
            + write_string(self.account.host)
            + struct.pack(">H", self.account.port)
            + write_varint(2)
        )
        payload_length = len(packet_id_bytes) + len(data)
        await self.client.send_raw(write_varint(payload_length) + packet_id_bytes + data)
        self.client.state = ConnectionState.LOGIN
        logger.info("Handshake sent (protocol %d, next_state=Login)", protocol_version)

    async def _do_status_probe(self) -> None:
        try:
            temp_client = McClient(self.account.host, self.account.port)
            await temp_client.connect()

            # Handshake for Status state
            handshake_data = (
                write_varint(DEFAULT_PROTOCOL)
                + write_string(self.account.host)
                + struct.pack(">H", self.account.port)
                + write_varint(1)
            )
            h_pid = write_varint(0x00)
            await temp_client.send_raw(write_varint(len(h_pid) + len(handshake_data)) + h_pid + handshake_data)
            temp_client.state = ConnectionState.STATUS

            # Status Request (0x00)
            await temp_client.send_raw(write_varint(1) + write_varint(0x00))

            packet = await temp_client.recv_packet()
            if packet.packet_id != 0x00:
                logger.warning("Expected Status Response (0x00), got 0x%02X", packet.packet_id)
                await temp_client.disconnect()
                return

            json_str, _ = read_string(packet.data, 0)
            status = json.loads(json_str)
            version_info = status.get("version", {})
            protocol_version = version_info.get("protocol", DEFAULT_PROTOCOL)
            version_name = version_info.get("name", "unknown")
            players_info = status.get("players", {})
            online = players_info.get("online", 0)
            max_players = players_info.get("max", 0)

            self.adapter = get_adapter(protocol_version)
            logger.info(
                "Server: %s (protocol %d), Players: %d/%d",
                version_name,
                protocol_version,
                online,
                max_players,
            )

            ping_data = write_long(int(asyncio.get_event_loop().time() * 1000))
            await temp_client.send_packet(0x01, ping_data)

            pong = await temp_client.recv_packet()
            if pong.packet_id == 0x01:
                logger.debug("Ping/Pong OK, latency normal")

            await temp_client.disconnect()
        except Exception as e:
            logger.warning("Status probe failed: %s (continuing with default protocol)", e)
            self.adapter = get_adapter("auto")

    async def _do_login(self) -> None:
        username = self.account.username

        if self.account.auth_mode == "microsoft":
            await self._do_microsoft_login(username)
            return

        offline_uuid = hashlib.md5(f"OfflinePlayer:{username}".encode()).hexdigest()
        uuid_bytes = bytes.fromhex(offline_uuid)

        await self._send_login_start(username, uuid_bytes)
        await self._handle_login_packets()

    async def _do_microsoft_login(self, username: str) -> None:
        from yuxi.channel.extensions.minecraft.auth import MinecraftAuth

        auth = MinecraftAuth()
        try:
            result = await auth.authenticate()
            username = result["username"]
            uuid_str = result["uuid"]
            self._ms_access_token = result.get("access_token")
        finally:
            await auth.close()

        uuid_bytes = __import__("uuid").UUID(uuid_str).bytes
        await self._send_login_start(username, uuid_bytes)
        await self._handle_login_packets()

    def _get_protocol_version(self) -> int:
        if self.adapter:
            return self.adapter.version
        return DEFAULT_PROTOCOL

    async def _send_login_start(self, username: str, uuid_bytes: bytes) -> None:
        data = bytearray()
        data += write_string(username)
        if self._get_protocol_version() >= 764:
            data += uuid_bytes
        await self.client.send_packet(0x00, bytes(data))
        logger.info("Login Start sent: username=%s", username)

    async def _handle_login_packets(self) -> None:
        while True:
            packet = await self.client.recv_packet()
            if packet.packet_id == 0x01:
                logger.error("Server requested encryption — online-mode server rejects offline login")
                raise RuntimeError("Server requires online-mode authentication (Microsoft OAuth not yet supported)")
            elif packet.packet_id == 0x02:
                offset = 0
                uuid_str, consumed = read_uuid(packet.data, offset)
                offset += consumed
                server_username, _ = read_string(packet.data, offset)
                logger.info("Login Success: uuid=%s username=%s", uuid_str, server_username)
                self.client.state = ConnectionState.CONFIGURATION
                break
            elif packet.packet_id == 0x03:
                try:
                    reason, _ = read_string(packet.data, 0)
                    logger.error("Login Disconnect: %s", reason)
                except Exception:
                    logger.error("Login Disconnect (could not parse reason)")
                raise RuntimeError("Login rejected by server")
            elif packet.packet_id == 0x04:
                await self._handle_login_plugin_request(packet.data)
            else:
                logger.error("Unexpected Login response: 0x%02X", packet.packet_id)
                raise RuntimeError(f"Unexpected packet during login: 0x{packet.packet_id:02X}")

    async def _handle_login_plugin_request(self, data: bytes) -> None:
        message_id, offset = read_varint(data, 0)
        channel, offset = read_string(data, offset)
        logger.debug("Login Plugin Request: channel=%s", channel)
        response = bytearray()
        response += write_varint(message_id)
        response += b"\x00"
        await self.client.send_packet(0x02, bytes(response))

    async def _enter_play(self) -> None:
        self.client.state = ConnectionState.CONFIGURATION
        logger.debug("Entered Configuration state")

        while True:
            packet = await self.client.recv_packet()

            if packet.packet_id == 0x03:
                threshold, _ = read_varint(packet.data, 0)
                self.client.compression_threshold = threshold
                logger.info("Compression enabled (threshold=%d)", threshold)

            elif packet.packet_id == 0x02:
                logger.debug("Configuration finished")
                await self._send_client_information()
                await self.client.send_packet(0x02)
                break

            elif packet.packet_id == 0x04:
                keep_alive_id = struct.unpack(">q", packet.data[:8])[0]
                response = write_long(keep_alive_id)
                await self.client.send_packet(0x04, response)
                logger.debug("Configuration Keep-Alive acknowledged")

            elif packet.packet_id == 0x05:
                ping_id = struct.unpack(">q", packet.data[:8])[0]
                response = write_long(ping_id)
                await self.client.send_packet(0x05, response)
                logger.debug("Configuration Ping/Pong acknowledged")

            elif packet.packet_id == 0x01:
                channel, offset = read_string(packet.data, 0)
                logger.debug("Configuration Plugin Message: channel=%s", channel)
                if channel == "minecraft:brand":
                    response = bytearray()
                    response += write_string("minecraft:brand")
                    response += write_string("yuxi")
                    await self.client.send_packet(0x01, bytes(response))

            elif packet.packet_id == 0x07:
                logger.debug("Registry Data received")

            elif packet.packet_id == 0x09:
                logger.debug("Resource Pack Push received, accepting")
                response = bytearray()
                response += write_varint(0)
                await self.client.send_packet(0x09, bytes(response))

            else:
                logger.debug("Unhandled Configuration packet: 0x%02X", packet.packet_id)

        packet = await self.client.recv_packet()
        login_play_id = self.adapter.cb("login_play") if self.adapter else 0x2B
        if login_play_id is None:
            login_play_id = 0x2B
        if packet.packet_id == login_play_id:
            self.client.state = ConnectionState.PLAY
            logger.info("Entered Play state")

            self.keepalive = KeepAliveManager(self.client)
            await self.keepalive.start()

            self.client.set_packet_handler(self._handle_packet)
            asyncio.create_task(self.client.run_recv_loop())
        else:
            logger.error("Expected Login (Play, 0x%02X), got 0x%02X", login_play_id, packet.packet_id)

    async def _handle_packet(self, packet: McPacket) -> None:
        pid = packet.packet_id

        keep_alive_cb = self.adapter.cb("keep_alive") if self.adapter else None
        if keep_alive_cb is None:
            keep_alive_cb = 0x26
        if pid == keep_alive_cb:
            keep_alive_id = struct.unpack(">q", packet.data[:8])[0]
            keep_alive_sb = self.adapter.sb("keep_alive") if self.adapter else 0x18
            if keep_alive_sb is None:
                keep_alive_sb = 0x18
            response = write_long(keep_alive_id)
            await self.client.send_packet(keep_alive_sb, response)
            if self.keepalive:
                self.keepalive.record_keep_alive()
            logger.debug("Keep-Alive: id=%d acknowledged", keep_alive_id)
            return

        player_chat_cb = self.adapter.cb("player_chat_message") if self.adapter else None
        if player_chat_cb and pid == player_chat_cb:
            if self.on_message:
                unified = mc_packet_to_unified(packet, self.account)
                if unified:
                    await self.on_message(unified)
            return

        system_chat_cb = self.adapter.cb("system_chat_message") if self.adapter else None
        if system_chat_cb and pid == system_chat_cb:
            if self.on_message:
                unified = mc_packet_to_unified(packet, self.account)
                if unified:
                    await self.on_message(unified)
            return

        player_info_cb = self.adapter.cb("player_info_update") if self.adapter else None
        if player_info_cb and pid == player_info_cb:
            await self._handle_player_info(packet.data)
            return

        respawn_cb = self.adapter.cb("respawn") if self.adapter else None
        if respawn_cb and pid == respawn_cb:
            logger.debug("Respawn packet received, re-entering dimension")
            if self.keepalive:
                await self.keepalive.stop()
                await self.keepalive.start()
            return

        disconnect_cb = self.adapter.cb("disconnect") if self.adapter else None
        if disconnect_cb is None:
            disconnect_cb = 0x1D
        if pid == disconnect_cb:
            reason = "Unknown"
            try:
                reason, _ = read_string(packet.data, 0)
            except Exception:
                pass
            logger.warning("Disconnected by server: %s", reason)
            if self.on_message:
                unified = mc_packet_to_unified(packet, self.account)
                if unified:
                    await self.on_message(unified)
            await self.client.disconnect()
            if self.keepalive:
                await self.keepalive.stop()
            return

    async def _handle_player_info(self, data: bytes) -> None:
        import uuid as uuid_mod

        if not data:
            return
        action = data[0]
        offset = 1
        count, offset = read_varint(data, offset)

        for _ in range(count):
            if offset + 16 > len(data):
                break
            uuid_bytes = data[offset : offset + 16]
            offset += 16
            uuid_str = str(uuid_mod.UUID(bytes=uuid_bytes))

            if action == 0:
                name, offset = read_string(data, offset)
                prop_count, offset = read_varint(data, offset)
                for __ in range(prop_count):
                    _, offset = read_string(data, offset)
                    _, offset = read_string(data, offset)
                    if offset < len(data) and data[offset]:
                        offset += 1
                        sig_len, offset = read_varint(data, offset)
                        offset += sig_len
                    else:
                        offset += 1
                if offset + 1 > len(data):
                    break
                gamemode, offset = read_varint(data, offset)
                ping, offset = read_varint(data, offset)
                if offset < len(data) and data[offset]:
                    offset += 1
                    _, offset = read_string(data, offset)
                else:
                    offset += 1
                if offset < len(data) and data[offset]:
                    offset += 1
                    _, offset = read_string(data, offset)
                else:
                    offset += 1
                self._player_cache[uuid_str] = {"name": name, "ping": ping, "gamemode": gamemode}
                logger.debug("Player joined: %s (uuid=%s)", name, uuid_str)

            elif action == 4:
                removed = self._player_cache.pop(uuid_str, None)
                if removed:
                    logger.debug("Player left: %s", removed["name"])

            else:
                break

    async def _send_client_information(self) -> None:
        data = bytearray()
        data += write_string("en_US")
        data += b"\x08"
        data += b"\x01"
        data += b"\x7f"
        data += b"\x00"
        data += b"\x01"
        data += b"\x01"
        data += write_varint(0)
        await self.client.send_packet(0x00, bytes(data))
        logger.debug("Client Information sent")