from dataclasses import dataclass
from random import randint
from socket import socket
import struct

MAX_OUTGOING_PAYLOAD_LENGTH = 1446
MAX_INCOMING_PAYLOAD_LENGTH = 4096
INVALID_PACKET_REQUEST_ID = 32767

# can be overridden to make send_command echo all commands before sending them; useful for troubleshooting
ECHO_COMMANDS = False 

# packet types
PKT_LOGIN = 3
PKT_LOGIN_RESPONSE = 2
PKT_COMMAND = 2
PKT_COMMAND_RESPONSE = 0

@dataclass
class Packet:
    """A single RCON packet, either sent or received."""
    packet_size: int
    request_id: int
    packet_type: int
    payload: bytes

    @staticmethod
    def recv(sock: socket) -> Packet:
        """Read and parse the next full packet from sock, blocking until complete."""
        # first four bytes (little-endian) contain the size of the remainder of the packet
        [packet_size] = struct.unpack('<i', sock.recv(4))
        # then, 4 bytes each for ID & type, arbitrary-length string payload, and two null bytes
        request_id, packet_type, payload, null = struct.unpack(f'<ii{packet_size-4-4-2}sh',
                                                               Packet._recvall(sock, packet_size))
        
        return Packet(packet_size, request_id, packet_type, payload)

    @staticmethod
    def _recvall(sock: socket, n: int) -> bytes:
        """Block until exactly n bytes have been read from sock."""
        buf = b''
        while len(buf) < n:
            chunk = sock.recv(n - len(buf))
            if not chunk:
                raise ConnectionError("socket closed before full packet received")
            buf += chunk
        return buf

    @staticmethod
    def send(sock: socket, packet_type: int, payload: str, request_id: int | None = None):
        """Build a packet from payload and transmit it on sock.

        Raises ValueError if the encoded payload exceeds MAX_OUTGOING_PAYLOAD_LENGTH.
        """
        payload_bytes = bytes(payload, encoding='utf-8')

        if len(payload_bytes) > MAX_OUTGOING_PAYLOAD_LENGTH:
            raise ValueError(f"attempted to send payload of length {len(payload_bytes)} bytes, which exceeds the maximum of {MAX_OUTGOING_PAYLOAD_LENGTH} bytes")

        # 4 bytes each for ID & type, arbitrary-length string payload, and two null bytes
        packet_size = 4 + 4 + len(payload_bytes) + 2

        if request_id == None:
            request_id = randint(1, 32766)

        data = struct.pack(f'<iii{len(payload_bytes)}sh', packet_size, request_id, packet_type, payload_bytes, 0)
        sock.sendall(data)


def login(sock: socket, password: str) -> bool:
    """Authenticate on an already-connected socket. Returns a bool representing whether login succeeded."""
    Packet.send(sock, PKT_LOGIN, password)
    response = Packet.recv(sock)

    if response.request_id == -1:
        return False # auth failure
    return True

def send_command(sock: socket, command: str) -> str:
    """Run command on an already-authenticated socket and return the decoded response."""
    # Some commands can have responses greater than the max response length of 4096, so the server will send the response fragmented across several packets. However, the Notchian server doesn't properly implement that fragmentation logic, meaning there can be ambiguity about when the actual last packet is. 
    # To handle this, we send two packets with distinct IDs. The first packet has the command we want to run; the second has an invalid packet type. Since the server responds to packets in the order they were received, we read from the socket until we get a response that matches the second packet; everything before was part of the first packet's response. 
    if ECHO_COMMANDS: print(command)
    Packet.send(sock, PKT_COMMAND, command)
    Packet.send(sock, 200, "", request_id=INVALID_PACKET_REQUEST_ID)

    response = b''
    while True:
        received = Packet.recv(sock)
        if received.request_id == INVALID_PACKET_REQUEST_ID: break # we got a response from our second packet
        response += received.payload

    return response.decode('utf-8', errors='replace')