import asyncio
import socket
import struct
import re
import os
from typing import Tuple, Optional, Dict, Any

class GoldSrcRCON:
    """GoldSrc (CS 1.6 / Half-Life 1) UDP RCON Client."""
    def __init__(self, host: str, port: int, password: str, timeout: float = 3.0):
        self.host = host
        self.port = int(port)
        self.password = password
        self.timeout = timeout

    async def execute(self, command: str) -> Tuple[bool, str]:
        loop = asyncio.get_running_loop()
        try:
            return await loop.run_in_executor(None, self._sync_execute, command)
        except Exception as e:
            return False, f"RCON Error: {str(e)}"

    def _sync_execute(self, command: str) -> Tuple[bool, str]:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(self.timeout)
        try:
            # 1. Get Challenge
            sock.sendto(b"\xFF\xFF\xFF\xFFchallenge rcon\n", (self.host, self.port))
            data, _ = sock.recvfrom(4096)
            
            challenge_match = re.search(r"challenge rcon (\d+)", data.decode("latin-1", errors="ignore"))
            challenge = challenge_match.group(1) if challenge_match else ""
            
            # 2. Send RCON command
            if challenge:
                cmd_packet = f'\xFF\xFF\xFF\xFFrcon {challenge} "{self.password}" {command}\n'.encode("latin-1")
            else:
                cmd_packet = f'\xFF\xFF\xFF\xFFrcon "{self.password}" {command}\n'.encode("latin-1")
                
            sock.sendto(cmd_packet, (self.host, self.port))
            
            # 3. Receive output
            response_chunks = []
            while True:
                try:
                    resp_data, _ = sock.recvfrom(4096)
                    # Strip GoldSrc packet header (0xFF 0xFF 0xFF 0xFF and prefix like 'l' or 'print')
                    if resp_data.startswith(b"\xFF\xFF\xFF\xFF"):
                        cleaned = resp_data[4:]
                        if cleaned.startswith(b"l") or cleaned.startswith(b"n"):
                            cleaned = cleaned[1:]
                        response_chunks.append(cleaned.decode("latin-1", errors="ignore"))
                    else:
                        response_chunks.append(resp_data.decode("latin-1", errors="ignore"))
                except socket.timeout:
                    break

            result_str = "".join(response_chunks).strip()
            if "Bad rcon_password" in result_str or "Invalid challenge" in result_str:
                return False, result_str
            return True, result_str if result_str else "Command sent successfully (no output)"
        except socket.timeout:
            return False, "Server timed out (no response via UDP RCON)"
        except Exception as e:
            return False, str(e)
        finally:
            sock.close()


class SourceRCON:
    """Source / CS2 TCP RCON Client."""
    SERVERDATA_AUTH = 3
    SERVERDATA_EXECCOMMAND = 2
    SERVERDATA_AUTH_RESPONSE = 2
    SERVERDATA_RESPONSE_VALUE = 0

    def __init__(self, host: str, port: int, password: str, timeout: float = 4.0):
        self.host = host
        self.port = int(port)
        self.password = password
        self.timeout = timeout

    def _pack_packet(self, packet_id: int, packet_type: int, body: str) -> bytes:
        encoded_body = body.encode("utf-8") + b"\x00\x00"
        size = len(encoded_body) + 8
        return struct.pack("<iii", size, packet_id, packet_type) + encoded_body

    async def execute(self, command: str) -> Tuple[bool, str]:
        loop = asyncio.get_running_loop()
        try:
            return await loop.run_in_executor(None, self._sync_execute, command)
        except Exception as e:
            return False, f"Source RCON Error: {str(e)}"

    def _sync_execute(self, command: str) -> Tuple[bool, str]:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(self.timeout)
        try:
            sock.connect((self.host, self.port))

            # 1. Authenticate
            auth_id = 999
            auth_packet = self._pack_packet(auth_id, self.SERVERDATA_AUTH, self.password)
            sock.sendall(auth_packet)

            # Receive Auth Response
            header = sock.recv(12)
            if len(header) < 12:
                return False, "Invalid response from Source server during Auth"
            
            size, resp_id, resp_type = struct.unpack("<iii", header)
            payload = sock.recv(size - 8) if size > 8 else b""

            # Some servers send an empty RESPONSE_VALUE before AUTH_RESPONSE
            if resp_type == self.SERVERDATA_RESPONSE_VALUE:
                header = sock.recv(12)
                size, resp_id, resp_type = struct.unpack("<iii", header)
                payload = sock.recv(size - 8) if size > 8 else b""

            if resp_id == -1 or resp_id != auth_id:
                return False, "Authentication failed (Invalid RCON password)"

            # 2. Send Command
            cmd_id = 1001
            cmd_packet = self._pack_packet(cmd_id, self.SERVERDATA_EXECCOMMAND, command)
            sock.sendall(cmd_packet)

            # 3. Read output
            header = sock.recv(12)
            if len(header) < 12:
                return True, "Command executed (empty response)"
            
            size, resp_id, resp_type = struct.unpack("<iii", header)
            body = sock.recv(size - 8)
            result = body.decode("utf-8", errors="ignore").rstrip("\x00")
            return True, result if result else "OK"
        except socket.timeout:
            return False, "Source server connection timed out"
        except Exception as e:
            return False, str(e)
        finally:
            sock.close()


async def execute_rcon(host: str, port: int, password: str, command: str, game: str = "cs16") -> Tuple[bool, str]:
    """
    Executes a real RCON command against CS 1.6 (GoldSrc) or CS2 (Source).
    """
    if not password:
        return False, "RCON password not configured for this server"
    
    if game == "cs2" or "csgo" in game:
        client = SourceRCON(host, port, password)
    else:
        client = GoldSrcRCON(host, port, password)
        
    return await client.execute(command)