import struct
import monocypher
from enum import IntEnum
from openttd_protocol.wire.tcp import TCPProtocol
from openttd_protocol.wire.read import read_uint8, read_string, read_uint16, read_uint32
from openttd_protocol.wire.exceptions import SocketClosed

class PacketGameType(IntEnum):
    ServerFull = 0
    ServerBanned = 1
    ClientJoin = 2
    ServerError = 3
    ClientUnused = 4
    ServerUnused = 5
    ServerGameInfo = 6
    ClientGameInfo = 7
    ServerNewGame = 8
    ServerShutdown = 9
    ServerGameInfoExtended = 10
    ServerAuthenticationRequest = 11
    ClientAuthenticationResponse = 12
    ServerEnableEncryption = 13
    ClientIdentify = 14
    ServerCheckNewGRFs = 15
    ClientNewGRFsChecked = 16
    ServerNeedCompanyPassword = 17
    ClientCompanyPassword = 18
    ClientSettingsPassword = 19
    ServerSettingsAccess = 20
    ServerWelcome = 21
    ServerClientInfo = 22
    ClientGetMap = 23
    ServerWaitForMap = 24
    ServerMapBegin = 25
    ServerMapSize = 26
    ServerMapData = 27
    ServerMapDone = 28
    ClientMapOk = 29
    ServerClientJoined = 30
    ServerFrame = 31
    ClientAck = 32
    ServerSync = 33
    ClientCommand = 34
    ServerCommand = 35
    ClientChat = 36
    ServerChat = 37
    ServerExternalChat = 38
    ClientQuit = 47
    ServerCompanyUpdate = 45
    PACKET_END = 100

class PacketAdminType(IntEnum):
    AdminJoin = 0
    AdminQuit = 1
    AdminUpdateFrequency = 2
    AdminPoll = 3
    AdminChat = 4
    AdminRcon = 5
    AdminGamescript = 6
    AdminPing = 7
    AdminExternalChat = 8
    AdminJoinSecure = 9
    AdminAuthResponse = 10

    ServerFull = 100
    ServerBanned = 101
    ServerError = 102
    ServerProtocol = 103
    ServerWelcome = 104
    ServerNewGame = 105
    ServerShutdown = 106

    ServerDate = 107
    ServerClientJoin = 108
    ServerClientInfo = 109
    ServerClientUpdate = 110
    ServerClientQuit = 111
    ServerClientError = 112
    ServerCompanyNew = 113
    ServerCompanyInfo = 114
    ServerCompanyUpdate = 115
    ServerCompanyRemove = 116
    ServerCompanyEconomy = 117
    ServerCompanyStats = 118
    ServerChat = 119
    ServerRcon = 120
    ServerConsole = 121
    ServerCmdNames = 122
    ServerCmdLoggingOld = 123
    ServerGamescript = 124
    ServerRconEnd = 125
    ServerPong = 126
    ServerCmdLogging = 127
    ServerAuthRequest = 128
    ServerEnableEncryption = 129

    PACKET_END = 130

class AdminUpdateFrequency(IntEnum):
    Poll = 1 << 0
    Daily = 1 << 1
    Weekly = 1 << 2
    Monthly = 1 << 3
    Quarterly = 1 << 4
    Annually = 1 << 5
    Automatic = 1 << 6

class AdminUpdateType(IntEnum):
    Date = 0
    ClientInfo = 1
    CompanyInfo = 2
    CompanyEconomy = 3
    CompanyStats = 4
    Chat = 5
    Console = 6
    CmdNames = 7
    CmdLogging = 8
    Gamescript = 9
    End = 10

class NetworkAuthenticationMethod(IntEnum):
    X25519_KeyExchangeOnly = 0
    X25519_PAKE = 1
    X25519_AuthorizedKey = 2

class OpenTTDProtocol(TCPProtocol):
    """Low-level OpenTTD TCP protocol handler with encryption support."""
    PacketType = PacketGameType
    PACKET_END = PacketGameType.PACKET_END

    def __init__(self, handler):
        super().__init__(handler)
        self.handler = handler

    def receive_packet(self, source, data):
        try:
            if self.handler.encryption_enabled:
                if not self.handler._recv_aead:
                    self.handler._recv_aead = monocypher.IncrementalAuthenticatedEncryption(self.handler._session_key_recv, self.handler._encryption_nonce)
                length, rest = read_uint16(data)
                payload = self.handler._recv_aead.unlock(bytes(rest[:16]), bytes(rest[16:]))
                if payload is None:
                    raise SocketClosed("Decryption failed")
                data = memoryview(struct.pack("<H", len(payload) + 2) + payload)
            
            return super().receive_packet(source, data)
        except Exception:
            return PacketGameType.ServerUnused, {}

    async def send_packet(self, data):
        if self.handler.encryption_enabled:
            if not self.handler._send_aead:
                self.handler._send_aead = monocypher.IncrementalAuthenticatedEncryption(self.handler._session_key_send, self.handler._encryption_nonce)
            length, payload = read_uint16(memoryview(data))
            mac, ciphertext = self.handler._send_aead.lock(payload.tobytes())
            data = struct.pack("<H", 18 + len(ciphertext)) + mac + ciphertext
        
        await self._can_write.wait()
        if self.transport.is_closing():
            raise SocketClosed
        self.transport.write(data)
        return len(data)

    # --- Static Parsers ---
    @staticmethod
    def receive_ServerGameInfo(source, data):
        from openttd_protocol.protocol.game import GameProtocol
        return GameProtocol.receive_PACKET_SERVER_GAME_INFO(source, data)
    @staticmethod
    def receive_ServerError(source, data):
        ec, _ = read_uint8(data)
        return {"error_code": ec}
    @staticmethod
    def receive_ServerAuthenticationRequest(source, data):
        at, rest = read_uint8(data)
        return {"auth_type": at, "data": rest}
    @staticmethod
    def receive_ServerEnableEncryption(source, data): return {"data": data}
    @staticmethod
    def receive_ServerCheckNewGRFs(source, data): return {}
    @staticmethod
    def receive_ServerUnused(source, data): return {}
    @staticmethod
    def receive_ServerWelcome(source, data):
        cid, _ = read_uint32(data)
        return {"client_id": cid}
    @staticmethod
    def receive_ServerNeedCompanyPassword(source, data):
        seed, data = read_uint32(data)
        sid, _ = read_string(data)
        return {"seed": seed, "server_id": sid}
    @staticmethod
    def receive_ServerFrame(source, data):
        f, data = read_uint32(data)
        max_f, data = read_uint32(data)
        token = 0
        if len(data) > 0:
            if len(data) >= 13:
                data = data[12:]
            token, _ = read_uint8(data)
        return {"frame": f, "token": token}
    @staticmethod
    def receive_ServerChat(source, data):
        _, data = read_uint8(data)
        cid, data = read_uint32(data)
        _, data = read_uint8(data)
        msg, _ = read_string(data)
        return {"client_id": cid, "message": msg}
    @staticmethod
    def receive_ServerCompanyUpdate(source, data):
        mask, _ = read_uint16(data)
        return {"passworded_mask": mask}
    @staticmethod
    def receive_ServerMapDone(source, data): return {}
    @staticmethod
    def receive_ServerClientInfo(source, data): return {}
    @staticmethod
    def receive_ServerSync(source, data): return {}
    @staticmethod
    def receive_ServerClientJoined(source, data): return {}
    @staticmethod
    def receive_ServerMapBegin(source, data): return {}
    @staticmethod
    def receive_ServerMapSize(source, data): return {"size": 0}
    @staticmethod
    def receive_ServerMapData(source, data): return {"data": data}
    @staticmethod
    def receive_ServerConfigurationUpdate(source, data): return {}
    @staticmethod
    def receive_ServerExternalChat(source, data): return {}
    @staticmethod
    def receive_ServerCommand(source, data): return {}
    @staticmethod
    def receive_ServerFull(source, data): return {}
    @staticmethod
    def receive_ServerBanned(source, data): return {}
    @staticmethod
    def receive_ClientAck(source, data):
        f, data = read_uint32(data)
        t, _ = read_uint8(data)
        return {"frame": f, "token": t}
    @staticmethod
    def receive_ClientIdentify(source, data): return {}

class OpenTTDAdminProtocol(OpenTTDProtocol):
    """Low-level OpenTTD Admin TCP protocol handler."""
    PacketType = PacketAdminType
    PACKET_END = PacketAdminType.PACKET_END

    @staticmethod
    def receive_ServerProtocol(source, data):
        v, data = read_uint8(data)
        updates = {}
        while True:
            has_more, data = read_uint8(data)
            if not has_more:
                break
            ut, data = read_uint16(data)
            freqs, data = read_uint16(data)
            updates[ut] = freqs
        return {"version": v, "updates": updates}

    @staticmethod
    def receive_ServerWelcome(source, data):
        name, data = read_string(data)
        ver, data = read_string(data)
        dedi, data = read_uint8(data)
        map_name, data = read_string(data)
        seed, data = read_uint32(data)
        land, data = read_uint8(data)
        date, data = read_uint32(data)
        width, data = read_uint16(data)
        height, _ = read_uint16(data)
        return {
            "server_name": name, "openttd_version": ver, "dedicated": bool(dedi),
            "map_name": map_name, "generation_seed": seed, "landscape": land,
            "start_date": date, "map_width": width, "map_height": height
        }

    @staticmethod
    def receive_ServerDate(source, data):
        d, _ = read_uint32(data)
        return {"date": d}

    @staticmethod
    def receive_ServerChat(source, data):
        action, data = read_uint8(data)
        dest_type, data = read_uint8(data)
        cid, data = read_uint32(data)
        msg, data = read_string(data)
        money = 0
        if len(data) >= 8:
            import struct
            money = struct.unpack("<Q", data[:8])[0]
        return {"action": action, "dest_type": dest_type, "client_id": cid, "message": msg, "money": money}

    @staticmethod
    def receive_ServerConsole(source, data):
        origin, data = read_string(data)
        text, _ = read_string(data)
        return {"origin": origin, "text": text}

    @staticmethod
    def receive_ServerRcon(source, data):
        color, data = read_uint16(data)
        text, _ = read_string(data)
        return {"color": color, "text": text}

    @staticmethod
    def receive_ServerRconEnd(source, data):
        cmd, _ = read_string(data)
        return {"command": cmd}

    @staticmethod
    def receive_ServerAuthRequest(source, data):
        at, rest = read_uint8(data)
        return {"auth_type": at, "data": rest}

    @staticmethod
    def receive_ServerEnableEncryption(source, data): return {"data": data}
    @staticmethod
    def receive_ServerError(source, data):
        ec, _ = read_uint8(data)
        return {"error_code": ec}
    @staticmethod
    def receive_ServerFull(source, data): return {}
    @staticmethod
    def receive_ServerBanned(source, data): return {}
    @staticmethod
    def receive_ServerShutdown(source, data): return {}
    @staticmethod
    def receive_ServerNewGame(source, data): return {}
    @staticmethod
    def receive_ServerClientJoin(source, data):
        cid, _ = read_uint32(data)
        return {"client_id": cid}

    @staticmethod
    def receive_ServerClientInfo(source, data):
        cid, data = read_uint32(data)
        addr, data = read_string(data)
        name, data = read_string(data)
        lang, data = read_uint8(data)
        date, data = read_uint32(data)
        playas, _ = read_uint8(data)
        return {
            "client_id": cid, "network_address": addr, "name": name,
            "language": lang, "join_date": date, "play_as": playas
        }

    @staticmethod
    def receive_ServerClientUpdate(source, data):
        cid, data = read_uint32(data)
        name, data = read_string(data)
        playas, _ = read_uint8(data)
        return {"client_id": cid, "name": name, "play_as": playas}

    @staticmethod
    def receive_ServerClientQuit(source, data):
        cid, _ = read_uint32(data)
        return {"client_id": cid}

    @staticmethod
    def receive_ServerClientError(source, data):
        cid, data = read_uint32(data)
        error, _ = read_uint8(data)
        return {"client_id": cid, "error_code": error}

    @staticmethod
    def receive_ServerCompanyNew(source, data):
        cid, _ = read_uint8(data)
        return {"company_id": cid}

    @staticmethod
    def receive_ServerCompanyInfo(source, data):
        cid, data = read_uint8(data)
        name, data = read_string(data)
        manager, data = read_string(data)
        color, data = read_uint8(data)
        protected, data = read_uint8(data)
        year, data = read_uint32(data)
        is_ai, _ = read_uint8(data)
        return {
            "company_id": cid, "name": name, "manager_name": manager,
            "color": color, "password_protected": bool(protected),
            "inaugurated_year": year, "is_ai": bool(is_ai)
        }

    @staticmethod
    def receive_ServerCompanyUpdate(source, data):
        cid, data = read_uint8(data)
        name, data = read_string(data)
        manager, data = read_string(data)
        color, data = read_uint8(data)
        protected, data = read_uint8(data)
        bankrupt, data = read_uint8(data)
        s1, data = read_uint8(data)
        s2, data = read_uint8(data)
        s3, data = read_uint8(data)
        s4, _ = read_uint8(data)
        return {
            "company_id": cid, "name": name, "manager_name": manager,
            "color": color, "password_protected": bool(protected),
            "quarters_of_bankruptcy": bankrupt, "share_owners": [s1, s2, s3, s4]
        }

    @staticmethod
    def receive_ServerCompanyRemove(source, data):
        cid, data = read_uint8(data)
        reason, _ = read_uint8(data)
        return {"company_id": cid, "reason": reason}

    @staticmethod
    def receive_ServerCompanyEconomy(source, data):
        import struct
        cid, data = read_uint8(data)
        money = struct.unpack("<Q", data[:8])[0]
        data = data[8:]
        loan = struct.unpack("<Q", data[:8])[0]
        data = data[8:]
        income = struct.unpack("<q", data[:8])[0]
        data = data[8:]
        delivered, data = read_uint16(data)
        val_lq = struct.unpack("<Q", data[:8])[0]
        data = data[8:]
        perf_lq, data = read_uint16(data)
        del_lq, data = read_uint16(data)
        val_pq = struct.unpack("<Q", data[:8])[0]
        data = data[8:]
        perf_pq, data = read_uint16(data)
        del_pq, _ = read_uint16(data)
        return {
            "company_id": cid, "money": money, "loan": loan, "income": income,
            "delivered_cargo": delivered, "value_last_quarter": val_lq,
            "performance_last_quarter": perf_lq, "delivered_cargo_last_quarter": del_lq,
            "value_previous_quarter": val_pq, "performance_previous_quarter": perf_pq,
            "delivered_cargo_previous_quarter": del_pq
        }

    @staticmethod
    def receive_ServerCompanyStats(source, data):
        cid, data = read_uint8(data)
        trains, data = read_uint16(data)
        lorries, data = read_uint16(data)
        buses, data = read_uint16(data)
        planes, data = read_uint16(data)
        ships, data = read_uint16(data)
        t_stations, data = read_uint16(data)
        l_stations, data = read_uint16(data)
        b_stops, data = read_uint16(data)
        airports, data = read_uint16(data)
        harbours, _ = read_uint16(data)
        return {
            "company_id": cid,
            "vehicles": {"trains": trains, "lorries": lorries, "buses": buses, "planes": planes, "ships": ships},
            "stations": {"train": t_stations, "lorry": l_stations, "bus": b_stops, "airport": airports, "harbour": harbours}
        }

    @staticmethod
    def receive_ServerGamescript(source, data):
        json_str, _ = read_string(data)
        import json
        try:
            return {"data": json.loads(json_str)}
        except Exception:
            return {"raw_data": json_str}

    @staticmethod
    def receive_ServerPong(source, data):
        payload, _ = read_uint32(data)
        return {"payload": payload}