# pyright: reportInvalidTypeForm=false
from __future__ import annotations

import os
import re
import io
import enum
import math
import tqdm
import argparse
import pydantic
import subprocess

from ctypes import c_uint16, c_int16, c_uint64, c_uint32, c_int32, c_uint8, c_float, sizeof as c_sizeof
from typing import Protocol, BinaryIO, Type, Iterable, cast
from pathlib import Path

from cstruct import CStructure, ctypes_types, CTypeType, c_addr_ptr, c_str, print_carr


def vram2offset(vram: int):
    return vram - 0xFF000


sceVu0FVECTOR = c_float * 4
sceVu0IVECTOR = c_int32 * 4
sceVu0FMATRIX = sceVu0FVECTOR * 4

qword = c_int32 * 4


class ANI_CODE(c_uint16):
    def __str__(self):
        return f"{self.value}"


class GS_REG(c_uint64):
    def __str__(self):
        return f"0x{self.value:x}"


class HEX_ADDR(c_uint32):
    def __str__(self):
        return f"0x{self.value:x}"


# CTypeTypeEX = CTypeType | type[sceVu0FVECTOR] | type[sceVu0IVECTOR]


###########################################################################

# struct _MNG_STATE_ST {
#     // Members
#     int state; // offset 0x0, size 0x4
#     u_int flag; // offset 0x4, size 0x4
#     u_int calc_total_time; // offset 0x8, size 0x4
#     u_int total_time; // offset 0xC, size 0x4
#     u_int get_time; // offset 0x10, size 0x4
#     u_int dmy[3]; // offset 0x14, size 0xC
# };
class MNG_STATE_ST(CStructure):
    state: c_int32
    flag: c_uint32
    calc_total_time: c_uint32
    total_time: c_uint32
    get_time: c_uint32
    flag: c_uint32
    dmy: c_uint32 * 3


# struct _TMP_MENU_MNG_ST {
#     // Members
#     struct _MNG_STATE_ST mng_state; // offset 0x0, size 0x20
#     int DrawState; // offset 0x20, size 0x4
#     u_char** pTmpMenuFileBuf; // offset 0x24, size 0x4
#     float draw_up_rate; // offset 0x28, size 0x4
#     u_char dmy[12]; // offset 0x2C, size 0xC
# };
class TMP_MENU_MNG_ST(CStructure):
    mng_state: MNG_STATE_ST
    DrawState: c_int32
    pTmpMenuFileBuf: c_addr_ptr
    draw_up_rate: c_float
    dmy: c_uint8 * 12

###########################################################################

elf_names: dict[str, str] = {
    "us": "SLUS_203.88",
    "eu": "SLES_508.21",
}


class VRamElf(BinaryIO):
    def __init__(self, elf: BinaryIO):
        self.elf = elf

    def read(self, n: int = -1):
        return self.elf.read(n)

    def seek(self, offset: int, whence: int = 0):
        return self.elf.seek(vram2offset(offset), whence)


# matches "var_name = 0x12345678; // attr1:val1 attr2:val2 ... attrN:valN"
line_pattern = re.compile(r"^([^\s=]+)\s*=\s*(0x[0-9a-fA-F]+);\s*\/\/((?:\s*[0-9a-zA-Z_]+(:?:[0-9a-zA-Z_\*]+)?)+)\s*$")
# matches multiple attributes in the form of attr:val or attr
attr_pattern = re.compile(r"([0-9a-zA-Z_]+)(?::([0-9a-zA-Z_\*]+))?")
# match address lines: "0xaabbccdd = &var"
addr_pattern = re.compile(r"^\s*(0x[0-9a-fA-F]+)\s*=\s*(.*)\s*;\s*(\/\/.*)?$")


class DataVar(pydantic.BaseModel):
    model_config = pydantic.ConfigDict(extra="forbid")

    _elf: VRamElf = pydantic.PrivateAttr()

    address: int
    name: str
    type: Type[CStructure] | CTypeType | type[sceVu0FVECTOR] | type[sceVu0IVECTOR] | type[qword] | type[c_str]
    numel: int | list[int] = 0
    nosize: bool = False
    static: bool = False

    num_ptr: int = 0

    @pydantic.model_validator(mode="after")
    def store_num_ptr(self):
        # transfer class attribute to instance
        if DataVar.num_ptr > 0:
            self.num_ptr = DataVar.num_ptr
        return self

    @pydantic.field_validator("type", mode="before")
    @classmethod
    def type_from_str(  # pyright: ignore
        cls, v: str | Type[CStructure] | CTypeType
    ) -> Type[CStructure] | CTypeType | sceVu0FVECTOR | sceVu0IVECTOR | qword | ANI_CODE | GS_REG | HEX_ADDR:
        if not isinstance(v, str):
            return v
        cls.num_ptr = v.count("*")  # temporary store num_ptr in class attribute
        v = v.rstrip("*")  # remove pointer(s)
        if v in ctypes_types:
            return ctypes_types[v]
        if v == "sceVu0FVECTOR":
            return sceVu0FVECTOR
        if v == "sceVu0IVECTOR":
            return sceVu0IVECTOR
        if v == "qword":
            return qword
        if v == "ANI_CODE":
            return ANI_CODE
        if v == "GS_REG":
            return GS_REG
        if v == "HEX_ADDR":
            return HEX_ADDR
        if v == "c_str":
            return c_str
        class_type = globals()[v]
        if issubclass(class_type, CStructure):
            return class_type
        raise ValueError(f"{v} is unknown/not a valid type")

    def data_var_dumps(self, addresses: dict[int, str]):
        self._elf.seek(self.address, os.SEEK_SET)

        if isinstance(self.numel, list):
            numel = math.prod(self.numel)
            ndims = len(self.numel)
        else:
            numel = max(1, self.numel)
            ndims = 1
        if self.num_ptr > 0:
            var_data = (c_addr_ptr * numel).from_buffer_copy(self._elf.read(numel * c_sizeof(c_addr_ptr)))
            if self.type == ANI_CODE:  # pyright: ignore
                type_str = "ANI_CODE" + ("*" * self.num_ptr)
            if self.type == GS_REG:  # pyright: ignore
                type_str = "u_long" + ("*" * self.num_ptr)
            if self.type == HEX_ADDR:  # pyright: ignore
                type_str = "u_int" + ("*" * self.num_ptr)
            elif issubclass(self.type, CStructure):  # pyright: ignore
                type_str = self.type.__name__ + ("*" * self.num_ptr)
            elif self.type == sceVu0FVECTOR:  # pyright: ignore
                type_str = "sceVu0FVECTOR" + ("*" * self.num_ptr)
            elif self.type == sceVu0IVECTOR:  # pyright: ignore
                type_str = "sceVu0IVECTOR" + ("*" * self.num_ptr)
            else:
                type_str = next(
                    k
                    for k, v in ctypes_types.items()
                    if getattr(v, "_type_") == getattr(self.type, "_type_")  # pyright: ignore
                ) + ("*" * self.num_ptr)
            stream = io.StringIO()
            if self.static:
                stream.write("static ")
            stream.write(f"{type_str} {self.name}")
            if numel > 1:
                numel = f"{self.numel}" if not self.nosize else ""
                var_str = f"{{ {', '.join(str(var) for var in cast(Iterable[c_addr_ptr], var_data))}, }}"  # pyright: ignore
                stream.write(f"[{numel}]")
            else:
                # should have been parsed as c_float_Array_4_Array_1
                assert len(var_data) == 1 and var_data[0].__class__ is c_addr_ptr
                var_str = str(var_data[0])
            stream.write(f" = {var_str};")
        elif issubclass(self.type, CStructure):  # pyright: ignore
            data = self._elf.read(numel * self.type.sizeof())

            return self.type.dumps(
                self.name,
                data,
                numel=self.numel,
                static=self.static,
                nosize=self.nosize,
                noarray=self.numel == 0,
            )
        elif self.type == sceVu0FVECTOR:  # pyright: ignore[reportUnknownMemberType]
            var_data = (sceVu0FVECTOR * numel).from_buffer_copy(self._elf.read(numel * c_sizeof(sceVu0FVECTOR)))
            type_str = "sceVu0FVECTOR"
            stream = io.StringIO()
            if self.static:
                stream.write("static ")
            stream.write(f"{type_str} {self.name}")

            def sceVu0FVECTOR_to_str(_v: sceVu0FVECTOR):  # pyright: ignore
                return f"{{ {', '.join(f'{flt}f' for flt in cast(Iterable[c_float], _v))} }}"

            if numel > 1:
                numel = f"{self.numel}" if not self.nosize else ""
                var_str = (
                    f"{{ {', '.join(sceVu0FVECTOR_to_str(var) for var in cast(Iterable[sceVu0FVECTOR], var_data))}, }}"  # pyright: ignore
                )
                stream.write(f"[{numel}]")
            else:
                # should have been parsed as c_float_Array_4_Array_1
                assert len(var_data) == 1 and var_data[0].__class__ is sceVu0FVECTOR
                var_str = sceVu0FVECTOR_to_str(var_data[0])
            stream.write(f" = {var_str};")
        elif self.type == sceVu0IVECTOR:  # pyright: ignore[reportUnknownMemberType]
            var_data = (sceVu0IVECTOR * numel).from_buffer_copy(self._elf.read(numel * c_sizeof(sceVu0IVECTOR)))
            type_str = "sceVu0IVECTOR"
            stream = io.StringIO()
            if self.static:
                stream.write("static ")
            stream.write(f"{type_str} {self.name}")

            def sceVu0IVECTOR_to_str(_v: sceVu0IVECTOR):  # pyright: ignore
                return f"{{ {', '.join(f'{int32}' for int32 in cast(Iterable[c_int32], _v))} }}"

            if numel > 1:
                numel = f"{self.numel}" if not self.nosize else ""
                var_str = (
                    f"{{ {', '.join(sceVu0IVECTOR_to_str(var) for var in cast(Iterable[sceVu0IVECTOR], var_data))}, }}"  # pyright: ignore
                )
                stream.write(f"[{numel}]")
            else:
                # should have been parsed as c_int_Array_4_Array_1
                assert len(var_data) == 1 and var_data[0].__class__ is sceVu0IVECTOR
                var_str = sceVu0IVECTOR_to_str(var_data[0])
            stream.write(f" = {var_str};")
        elif self.type == c_str:  # pyright: ignore
            assert numel > 1, "we only handle string arrays here, simple strings should be embedded in the source"
            var_data = (self.type * numel).from_buffer_copy(self._elf.read(numel * c_sizeof(self.type)))  # pyright: ignore
            stream = io.StringIO()
            if self.static:
                stream.write("static ")
            numel = f"{self.numel}" if not self.nosize else ""
            stream.write(f"char *{self.name}[{numel}] = {{")
            for var in cast(Iterable[c_str], var_data):
                stream.write(f"{var.to_str(self._elf)},")
            stream.write("};")
        else:
            typ = self.type  # pyright: ignore
            if ndims == 1:
                typ *= numel  # pyright: ignore
            elif isinstance(self.numel, list):
                for n in reversed(self.numel):
                    typ *= n  # pyright: ignore
            var_data = (typ).from_buffer_copy(self._elf.read(numel * c_sizeof(self.type)))  # pyright: ignore
            if self.type.__name__ == "ANI_CODE":  # pyright: ignore
                type_str = "ANI_CODE"
            elif self.type.__name__ == "GS_REG":  # pyright: ignore
                type_str = "u_long"
            elif self.type.__name__ == "HEX_ADDR":  # pyright: ignore
                type_str = "u_int"
            else:
                type_str = next(
                    k
                    for k, v in ctypes_types.items()
                    if getattr(v, "_type_") == getattr(self.type, "_type_")  # pyright: ignore
                )
            stream = io.StringIO()
            if self.static:
                stream.write("static ")
            stream.write(f"{type_str} {self.name}")
            if ndims == 1:
                numel = f"{self.numel}" if not self.nosize else ""
                var_str = f"{{ {', '.join(f'{var}' for var in cast(Iterable[CTypeType], var_data))} }}"
                stream.write(f"[{numel}]")
            elif ndims > 1:
                assert isinstance(self.numel, list)
                numel = ""
                for n, num in enumerate(self.numel):
                    numel += "[]" if n == 0 and self.nosize else f"[{num}]"
                stream.write(numel)
                tmpbuf = io.StringIO()
                tmpbuf.write("{")
                print_carr(var_data, tmpbuf)  # pyright: ignore
                tmpbuf.write("}")
                var_str = tmpbuf.getvalue()
            else:
                var_str = f"{var_data.value}"  # pyright: ignore
            stream.write(f" = {var_str};")
        return stream.getvalue()


def parse_data_vars(data_vars_txt: Path, strict: bool = True):
    data_vars: list[DataVar] = []
    addr_vals: dict[int, str] = {}

    with open(data_vars_txt, mode="r") as fh:
        for n, line in enumerate(fh):
            line = line.strip()

            # skip empty lines and comments
            if not line or line.startswith("//"):
                continue

            valid = False

            if addr_match := addr_pattern.match(line):
                address = int(addr_match.group(1), 16)
                value = addr_match.group(2)
                addr_vals[address] = value
                valid = True

            elif line_match := line_pattern.match(line):
                address = int(line_match.group(2), 16)
                options = line_match.group(3)
                name = line_match.group(1)
                if attr_match := attr_pattern.findall(options):
                    attrs: dict[str, str | list[str] | bool] = {}
                    for attr_name, attr_val in cast(list[tuple[str, str]], attr_match):
                        if attr_name not in attrs:
                            attrs[attr_name] = attr_val
                        elif not isinstance(attrs[attr_name], list):
                            attrs[attr_name] = cast(list[str], [attrs[attr_name], attr_val])
                        else:
                            cast(list[str], attrs[attr_name]).append(attr_val)
                    for key, value in attrs.items():
                        if not isinstance(value, bool) and len(value) == 0:
                            attrs[key] = True
                    try:
                        data_var = DataVar(address=address, name=name, **attrs)  # pyright: ignore[reportArgumentType]
                        data_vars.append(data_var)
                        if address not in addr_vals:
                            if (
                                isinstance(data_var.numel, list) or data_var.numel > 0 or data_var.type == sceVu0FVECTOR  # type: ignore
                            ):
                                addr_vals[address] = name
                            else:
                                addr_vals[address] = f"&{name}"
                        valid = True
                    except pydantic.ValidationError as ve:
                        print(f"syntax error at line {n}")
                        raise ve

            if strict and not valid:
                raise ValueError(f"syntax error at line {n}: {line}")

    return data_vars, addr_vals


def parse_data(lang: str):
    data_vars_txt = Path("config") / lang / "data_vars.txt"

    data_vars, addr_vals = parse_data_vars(data_vars_txt)

    if not data_vars:
        return

    if len(set(data_var.name for data_var in data_vars)) != len(data_vars):
        raise RuntimeError("duplicate names in data_vars.txt")

    elf_path = Path("config") / lang / elf_names[lang]

    include_path = Path("config") / lang / "include" / "data"
    include_path.mkdir(parents=True, exist_ok=True)

    c_addr_ptr.set_addresses(addr_vals)

    extract_filter: list[str] = []
    extract_filter_str = os.environ.get("EXTRACT_DATA_FILTER", "").strip()
    if extract_filter_str:
        extract_filter: list[str] = extract_filter_str.split(",")

    with open(elf_path, mode="rb") as elf:
        vram_elf = VRamElf(elf)
        CStructure.__elf__ = vram_elf

        for data_var in tqdm.tqdm(data_vars, desc="Extracting data"):
            header_path = include_path / f"{data_var.name}.h"

            if extract_filter and data_var.name not in extract_filter:
                continue

            with header_path.open(mode="w") as fw:
                data_var._elf = vram_elf  # pyright: ignore[reportPrivateUsage]

                fw.write(data_var.data_var_dumps(addr_vals))

            try:
                subprocess.run(["clang-format", "-i", "--style=Microsoft", header_path])
            except Exception:
                # formatting is best effort for now
                pass


def main():
    class Language(str, enum.Enum):
        US = "us"
        EU = "eu"

    languages = [lang.value for lang in Language]

    class ArgsProtocol(Protocol):
        language: Language

    parser = argparse.ArgumentParser(description="fixes asm removing gp_rel macro")
    parser.add_argument("language", type=str, choices=languages, help="language to extract data for")

    args = cast(ArgsProtocol, parser.parse_args())

    parse_data(args.language)


if __name__ == "__main__":
    main()