#!/usr/bin/env python3
"""MC RCON 客户端"""

import socket
import struct

class RCONClient:
    MAX_RETRIES = 3
    RETRY_DELAY = 2
    CONNECT_TIMEOUT = 5
    SOCKET_TIMEOUT = 10

    def __init__(self, host, port, password):
        self.host = host
        self.port = port
        self.password = password
        self.sock = None
        self.packet_id = 1
        self.authenticated = False

    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.settimeout(self.CONNECT_TIMEOUT)
        self.sock.connect((self.host, self.port))
        self.authenticate()

    def authenticate(self):
        packet = self._build_packet(self.packet_id, 3, self.password)
        self.sock.sendall(packet)
        response = self._read_response()
        if not response or response.get("id") == -1:
            raise Exception("RCON authentication failed")
        self.authenticated = True
        self.packet_id += 1

    def is_connected(self):
        if not self.sock or not self.authenticated:
            return False
        try:
            self.sock.send(b"")
            return True
        except Exception:
            return False

    def send_command(self, command, max_retries=None):
        if max_retries is None:
            max_retries = self.MAX_RETRIES
        last_error = None
        for attempt in range(max_retries):
            try:
                if not self.is_connected():
                    self.close()
                    self.connect()
                packet = self._build_packet(self.packet_id, 2, command)
                self.sock.sendall(packet)
                self.packet_id += 1
                response = self._read_response()
                return response
            except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError, OSError) as e:
                last_error = e
                self.close()
                if attempt < max_retries - 1:
                    import time
                    time.sleep(self.RETRY_DELAY)
            except Exception as e:
                last_error = e
                self.close()
                if attempt < max_retries - 1:
                    import time
                    time.sleep(self.RETRY_DELAY)
        return None

    def _build_packet(self, packet_id, packet_type, body_str):
        body = body_str.encode("utf-8") + b"\x00\x00"
        payload = struct.pack("<ii", packet_id, packet_type) + body
        length = struct.pack("<i", len(payload))
        return length + payload

    def _read_response(self):
        try:
            self.sock.settimeout(self.SOCKET_TIMEOUT)
            length_data = self.sock.recv(4)
        except (ConnectionResetError, ConnectionAbortedError, OSError, socket.timeout):
            return None
        if not length_data:
            return None
        length = struct.unpack("<i", length_data)[0]
        if length <= 0:
            return None
        remaining = b""
        while len(remaining) < length:
            try:
                chunk = self.sock.recv(length - len(remaining))
                if not chunk:
                    break
                remaining += chunk
            except (ConnectionResetError, ConnectionAbortedError, OSError, socket.timeout):
                break
        if len(remaining) < 8:
            return None
        packet_id, packet_type = struct.unpack("<ii", remaining[:8])
        body = remaining[8:-2] if len(remaining) > 8 else b""
        return {"id": packet_id, "type": packet_type, "body": body.decode("utf-8", errors="replace")}

    def close(self):
        if self.sock:
            self.sock.close()
        self.sock = None
        self.authenticated = False

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()