from json import load
from dataclasses import dataclass, field
from typing import Literal, Optional, Union
from pprint import pprint


@dataclass(slots=True)
class Condition:
    a: Union["Condition", str, int, float]
    operator: Literal["gt", "lt", "ge", "le", "eq", "and", "or", "xor"]
    b: Union["Condition", str, int, float]


@dataclass(slots=True)
class Field:
    type: Literal[
        "bool",
        "byte",
        "ubyte",
        "short",
        "ushort",
        "int",
        "uint",
        "long",
        "ulong",
        "float",
        "double",
        "itemstack",
        "optional_itemstack",
        "itemstack_nbt",
        "optional_itemstack_nbt",
        "compoundtag",
        "synchedentitydata",
        "list",  # AAABBBCCC
        "column_major_list",  # ABCABCABC
        "string",
        "utfstring",
        "bytes",
        "bytes_compressed",
        "struct",
    ]

    maxsize: Optional[int] = None
    size: Optional[int] = None
    sizetype: Optional[
        Literal["byte", "ubyte", "short", "ushort", "int", "uint"]
    ] = None
    condition: Optional[Condition] = None
    method: Optional[Literal["zlib", "gzip"]] = None
    fields: Optional[list["NamedField"]] = None
    item: Optional[Union["Field", "NamedField"]] = None

    @staticmethod
    def _add_size(kwargs: dict, data: dict, sizetype: str = "short", maxsize: int = 256) -> dict:
        if (size := data.get("size")) is not None:
            kwargs['size'] = size
        else:
            kwargs['sizetype'] = data.get("sizetype", sizetype)
            kwargs['maxsize'] = int(data.get("maxsize", maxsize))
        return kwargs


    @classmethod
    def _get_init_args(cls, data: dict) -> tuple[tuple, dict]:
        args = (data['type'],)
        kwargs = {}
        kwargs['condition'] = data.get("condition")

        match data["type"]:
            case "bool":
                pass
            case "byte" | "ubyte":
                pass
            case "short" | "ushort":
                pass
            case "int" | "uint":
                pass
            case "long" | "ulong":
                pass
            case "float":
                pass
            case "double":
                pass
            case "string" | "utfstring" | "bytes":
                kwargs = cls._add_size(kwargs, data)
            case "bytes_compressed":
                kwargs['method'] = data.get("method")
                assert kwargs['method'] in ("gzip", "zlib")
                kwargs = cls._add_size(kwargs, data, maxsize=16777216)
            case "list":
                kwargs = cls._add_size(kwargs, data, maxsize=1024)
                kwargs['item'] = data.get("item")
                raise NotImplementedError(f"can't do list yet")
            case "column_major_list":
                kwargs = cls._add_size(kwargs, data)
                kwargs['fields'] = list(map(NamedField.fromdict, data["fields"]))
            case "optional_itemstack":
                print("TODO: optional_itemstack")
            case "optional_itemstack_nbt":
                print("TODO: optional_itemstack_nbt")
            case "itemstack":
                print("TODO: itemstack")
            case "itemstack_nbt":
                print("TODO: itemstack_nbt")
            case "synchedentitydata":
                print("TODO: synchedentitydata")
            case "compoundtag":
                print("TODO: compoundtag")
            case name:
                raise NotImplementedError(f"unknown field type {name}")
        return args, kwargs

    @classmethod
    def fromdict(cls, data: dict) -> "Field":
        args, kwargs = cls._get_init_args(data)
        return cls(*args, **kwargs)


@dataclass(slots=True)
class NamedField(Field):
    name: str = ""

    @classmethod
    def _get_init_args(cls, data: dict) -> tuple[tuple, dict]:
        args, kwargs = Field._get_init_args(data)
        kwargs['name'] = data['name']
        return args, kwargs

    @classmethod
    def fromdict(cls, data: dict) -> "NamedField":
        args, kwargs = cls._get_init_args(data)
        return cls(*args, **kwargs)


@dataclass(slots=True)
class PacketSchema:
    id: int
    name: str
    server: bool = False
    client: bool = False
    fields: list[NamedField] = field(default_factory=list)

    @classmethod
    def fromdict(cls, data: dict) -> "PacketSchema":
        assert data.get("server", False) or data.get("client", False)
        return cls(
            data["id"],
            data["name"],
            data.get("server", False),
            data.get("client", True),
            list(map(NamedField.fromdict, data["fields"])),
        )


@dataclass(slots=True)
class ProtocolSchema:
    protocol_version: int
    game_version: str
    packets: list[PacketSchema]

    @classmethod
    def fromdict(cls, data: dict) -> "ProtocolSchema":
        return cls(
            protocol_version=data["protocol_version"],
            game_version=data["game_version"],
            packets=list(map(PacketSchema.fromdict, data["packets"])),
        )


with open("../packets.json", "r") as f:
    data = load(f)

    schema = ProtocolSchema.fromdict(data)