import asyncio
import logging
import struct

logger = logging.getLogger(__name__)

RCON_TYPE_COMMAND = 2
RCON_TYPE_AUTH = 3
RCON_TYPE_RESPONSE = 0


class McRconClient:
    def __init__(self, host: str, port: int, password: str):
        self.host = host
        self.port = port
        self.password = password
        self._reader: asyncio.StreamReader | None = None
        self._writer: asyncio.StreamWriter | None = None
        self._request_id = 0

    async def connect(self) -> bool:
        try:
            self._reader, self._writer = await asyncio.wait_for(
                asyncio.open_connection(self.host, self.port), timeout=10
            )
        except Exception:
            logger.exception("RCON connection failed to %s:%d", self.host, self.port)
            return False

        self._request_id += 1
        success = await self._send_packet(RCON_TYPE_AUTH, self.password)
        if not success:
            logger.error("RCON authentication failed")
            return False
        return True

    async def execute(self, command: str) -> str:
        if not self._writer or self._writer.is_closing():
            if not await self.connect():
                return "RCON connection failed"
        self._request_id += 1
        return await self._send_packet(RCON_TYPE_COMMAND, command)

    async def _send_packet(self, ptype: int, payload: str) -> str:
        req_id = self._request_id
        packet = struct.pack("<iii", len(payload) + 10, req_id, ptype) + payload.encode("utf-8") + b"\x00\x00"
        self._writer.write(packet)
        await self._writer.drain()

        response_parts = []
        while True:
            length_data = await self._reader.readexactly(4)
            length = struct.unpack("<i", length_data)[0]
            remaining = await self._reader.readexactly(length)
            resp_id = struct.unpack("<i", remaining[:4])[0]
            body = remaining[8:-2].decode("utf-8", errors="replace")

            if resp_id != req_id:
                continue

            response_parts.append(body)

            if ptype == RCON_TYPE_AUTH:
                return resp_id != -1

            if len(body) < 4096:
                break

        return "".join(response_parts)

    async def close(self):
        if self._writer:
            self._writer.close()
            await self._writer.wait_closed()