#!/usr/bin/env python3

import io
import json
import os
import re
import sys

from dataclasses import dataclass, field
from enum import Enum, auto as EnumAuto
from pathlib import Path
from pprint import pprint
from typing import *
from datetime import datetime, timezone

from clang import cindex

TEMPLATE_STANZA = (
    f"""-- DO NOT EDIT - THIS FILE IS AUTOGENERATED VIA './etlegacy/misc/generate-lua-definitions.py'
--               MANUAL CHANGES WILL BE LOST!
--
--        Timestamp : {datetime.now(timezone.utc).isoformat()}
--        Latest    : https://github.com/etlegacy/etlegacy/blob/master/misc/luals-definitions/generated.lua
"""
    + """
---@meta legacy

---@class et
et = {}

---@class Vec2
---@field [1] number
---@field [2] number

---@class Vec3
---@field [1] number
---@field [2] number
---@field [3] number

---@class Vec4
---@field [1] number
---@field [2] number
---@field [3] number
---@field [4] number

---@class Vec5
---@field [1] number
---@field [2] number
---@field [3] number
---@field [4] number
---@field [5] number

"""
)


class DefinitionType(Enum):
    HOOK = EnumAuto()
    FUNCTION = EnumAuto()
    CONSTANT = EnumAuto()


def clean_typed_name(name: str):
    return name.rstrip("?")


@dataclass
class LuaCommentAnnotatedParam:
    name: str
    type: str
    description: str


@dataclass
class LuaComment:
    name: str = None
    description: str = ""
    params: List[str] = field(default_factory=list)
    params_annotated: dict[str, LuaCommentAnnotatedParam] = field(default_factory=dict)
    returns_annotated: List[LuaCommentAnnotatedParam] = field(default_factory=list)
    raw_annotations: List[str] = field(default_factory=list)

    def __str__(self):
        result = ""

        for line in io.StringIO(self.description):
            result += "---" + line

        for param in self.params_annotated.values():
            result += f"---@param {param.name} {param.type}"
            if param.description:
                result += f" {param.description}"
            result += "\n"

        for param in self.returns_annotated:
            result += f"---@return {param.type} {param.name}"
            if param.description:
                result += f" {param.description}"
            result += "\n"

        for line in self.raw_annotations:
            result += line

        return result


@dataclass
class LuaDefinition:
    c_name: str
    c_comment: str
    type: DefinitionType
    lua_name: str = None
    lua_comment: LuaComment = field(default_factory=LuaComment)

    def __parse_lua_comment(self):
        result = LuaComment(name=self.lua_name)

        def parse_c_comment_def_prototype(line):
            line = line.strip()
            doxyflag, line = line.split(maxsplit=1)
            assert doxyflag == "@lua_def_prototype", "'doxyflag' is unexpected"

            match = re.match(re_c_comment_def_prototype, line)
            assert match, "Failed to parse doxygen comment def prototype"

            parsed_name, parsed_args_line = match.groups()
            assert (
                parsed_name == self.lua_name
            ), "Parsed Name doesn't match wrangled name"

            parsed_params = []
            if parsed_args_line:
                parsed_params = [x.strip() for x in parsed_args_line.split(",")]

            result.params = parsed_params

        def parse_c_comment_def(line):
            line = line.strip()
            doxyflag, luaflag, line = line.split(maxsplit=2)

            assert doxyflag == "@lua_def", "'doxyflag' is unexpected"

            if luaflag == "---@param":
                name, type, *rest = line.split(maxsplit=2)
                descr = rest[0] if len(rest) >= 1 else None
                result.params_annotated[clean_typed_name(name)] = (
                    LuaCommentAnnotatedParam(
                        name=name,
                        type=type,
                        description=descr,
                    )
                )
            elif luaflag == "---@return":
                type, name, *rest = line.split(maxsplit=2)
                descr = rest[0] if len(rest) >= 1 else None
                result.returns_annotated.append(
                    LuaCommentAnnotatedParam(
                        name=name,
                        type=type,
                        description=descr,
                    )
                )
            else:
                result.raw_annotations.append(luaflag + " " + line + "\n")

        if self.c_comment == None:
            return

        start_end_marker_found = 0
        lines = io.StringIO(self.c_comment)
        eat_newlines = False
        for line in lines:
            stripped_line = line.strip()
            match stripped_line:
                case "/**" | "/*" | "*/":
                    start_end_marker_found += 1
                    continue

            assert stripped_line.startswith("*"), "Expected line to start with '*'"

            line = stripped_line.removeprefix("*").strip()

            # skip empty lines
            if eat_newlines and not line:
                continue

            line += "\n"

            if not line.startswith("@"):  # description
                eat_newlines = False
                result.description += line
                continue
            elif line.startswith("@lua_def_prototype"):
                parse_c_comment_def_prototype(line)
            elif line.startswith("@lua_def"):
                parse_c_comment_def(line)
            else:
                print_debug(f"Skipping unknown comment line: {line.rstrip()}")
            eat_newlines = True

        result.description = result.description.rstrip() + "\n"

        assert (
            start_end_marker_found == 2
        ), f"Expected to match both a start and an end marker. ['/**', '*/']\n\t{self.c_name} / expected 2 found {start_end_marker_found}"

        # TODO - opt
        for x in result.params_annotated:
            assert (
                x in result.params
            ), f"Found a ---@param annotation that tried to annotate a missing parameter\n\t{self.c_name} → {x}"

        self.lua_comment = result

    def __post_init__(self):
        # wrangle 'c_name' into 'lua_name'
        match self.type:
            case DefinitionType.FUNCTION:
                self.lua_name = "et." + self.c_name.removeprefix("_et_")
            case DefinitionType.HOOK:
                self.lua_name = "et_" + self.c_name.removeprefix("G_LuaHook_")
            case _:
                self.lua_name = self.c_name

        self.__parse_lua_comment()

    def __str__(self):
        result = ""
        params = ", ".join(self.lua_comment.params)

        match self.type:
            case DefinitionType.FUNCTION:
                assert self.c_name, "Expected to have a c_name"
                if self.c_comment:
                    result += str(self.lua_comment)
                result += f"function {self.lua_name}({params}) end"
            case DefinitionType.HOOK:
                assert self.c_name, "Expected to have a c_name"
                params_typed = ", ".join(
                    [
                        f"{x}: {self.lua_comment.params_annotated[x].type}"
                        for x in self.lua_comment.params
                    ]
                )

                if self.c_comment:
                    result += str(self.lua_comment)

                returns_annotated = ""
                if self.lua_comment.returns_annotated:
                    returns_annotated = ": " + ", ".join(
                        x.type for x in self.lua_comment.returns_annotated
                    )

                result += f"---@alias {self.lua_name} fun({params_typed}){returns_annotated}\n"
                result += f"function {self.lua_name}({params}) end"

            case DefinitionType.CONSTANT:
                result += "---@type number\n"
                result += f"et.{self.lua_name} = nil"

        return result


re_c_comment_def_prototype = r"^([^\(]+)\s*\(([^\)]*)\)$"
inputstr = "function_name ( arg1, arg2 )"

output_lines = []
definitions: List[LuaDefinition] = []
script_parent_dir = Path(__file__).resolve().parent


def print_debug(s):
    # TODO
    print(s)


def regenerate():
    global output_lines
    global definitions

    with open(
        script_parent_dir.joinpath("luals-definitions", "generated.lua"),
        "w",
    ) as f:
        f.write(TEMPLATE_STANZA)
        f.writelines(output_lines)

        def definition_sort(defns: List[DefinitionType]):
            return sorted(defns, key=lambda defn: defn.lua_name)

        definitions_filtered = {
            "Functions": definition_sort(
                x for x in definitions if x.type == DefinitionType.FUNCTION
            ),
            "Hooks": definition_sort(
                x for x in definitions if x.type == DefinitionType.HOOK
            ),
            "Constants": definition_sort(
                x for x in definitions if x.type == DefinitionType.CONSTANT
            ),
        }

        for k, v in definitions_filtered.items():
            f.write(f"-- {{{{{{ {k}\n\n")
            for definition in v:
                f.write(str(definition))
                f.write("\n")
                f.write("\n")
            f.write(f"-- }}}}}} {k}\n")


def load_compile_commands(file_path, ccjson_path="compile_commands.json"):
    with open(ccjson_path, "r") as f:
        commands = json.load(f)

    file_path = os.path.abspath(file_path)
    for entry in commands:
        if os.path.abspath(entry["file"]) == file_path:
            return entry["arguments"][1:]  # skip the compiler binary
    raise ValueError(f"No compile command found for {file_path}")


def is_g_lua_function(cursor):
    return (
        cursor.kind == cindex.CursorKind.FUNCTION_DECL
        and cursor.is_definition()
        # TODO
        and (cursor.spelling.startswith("G_Lua") or cursor.spelling.startswith("_et_"))
    )


def get_doc_comment(cursor):
    raw_comment = cursor.raw_comment
    return raw_comment.strip() if raw_comment else None


def register_constant(ident):
    global definitions
    definitions.append(
        LuaDefinition(
            c_name=ident,
            c_comment=None,
            type=DefinitionType.CONSTANT,
        )
    )


def register_definition(cursor, type):
    global definitions

    name: str = cursor.spelling

    comment: str = get_doc_comment(cursor)

    match type:
        case DefinitionType.HOOK:
            definitions.append(
                LuaDefinition(
                    c_name=name,
                    c_comment=comment,
                    type=type,
                )
            )
        case DefinitionType.FUNCTION:
            definitions.append(
                LuaDefinition(
                    c_name=name,
                    c_comment=comment,
                    type=type,
                )
            )


def find_lua_macro_calls(tu):
    tokens = list(tu.get_tokens(extent=tu.cursor.extent))
    i = 0
    while i < len(tokens):
        t = tokens[i]
        if (
            t.kind == cindex.TokenKind.IDENTIFIER
            and t.spelling == "lua_regconstinteger"
        ):
            # Pattern: lua_regconstinteger ( vm->L , IDENT ) ;
            if (
                i + 6 < len(tokens)
                and tokens[i + 1].spelling == "("
                and tokens[i + 2].spelling == "vm"
                and tokens[i + 3].spelling == "->"
                and tokens[i + 4].spelling == "L"
                and tokens[i + 5].spelling == ","
                and tokens[i + 6].kind == cindex.TokenKind.IDENTIFIER
            ):
                ident = tokens[i + 6].spelling
                register_constant(ident)
            i += 7
        else:
            i += 1


def walk(cursor):
    global output_lines

    # print(
    #     f"Cursor: {cursor.spelling}, Kind: {cursor.kind}, Line: {cursor.location.line}"
    # )

    if is_g_lua_function(cursor):
        name: str = cursor.spelling

        if name.startswith("G_LuaHook_"):
            register_definition(cursor, DefinitionType.HOOK)
        elif name.startswith("_et_"):
            register_definition(cursor, DefinitionType.FUNCTION)
        else:
            pass  # TODO - not a function we care about

    for child in cursor.get_children():
        walk(child)


def main(args):
    assert args.path, "Require a path to be passed. See -h"
    path = args.path

    if not os.path.isfile(path):
        print("File does not exist:", path)
        sys.exit(1)

    index = cindex.Index.create()
    tu = index.parse(
        path,
        # TODO - settle on flags
        args=[
            "-xc",
            # assume some features
            "-DFEATURE_LUA",
            "-DFEATURE_MULTIVIEW",
            # for 'stddef.h'
            "-isystem",
            "./.pixi/envs/clang/lib/clang/20/include",
            "-isystem",
            "./.pixi/envs/clang/include/",
            # for build-generated version header
            "-I",
            "./etmain/ui/",
        ],
    )

    if not tu or not tu.cursor:
        print("Failed to parse. Check libclang path or include errors.")
        return

    for diag in tu.diagnostics:
        print("libclang-diagnostic: " + str(diag))

    find_lua_macro_calls(tu)

    walk(tu.cursor)

    regenerate()


def cli(argv):
    import argparse

    parser = argparse.ArgumentParser(
        description="Regenerate lua definitions from .c source files via 'libclang'."
    )

    parser.add_argument("path", nargs="?", type=str, help=".c file to parse.")

    parser.add_argument(
        "-rt", "--run-tests", action="store_true", help="Run unit tests."
    )
    # TODO
    parser.add_argument(
        "-d", "--debug", action="store_true", help="Enable debug logging."
    )

    args = parser.parse_args(argv)
    if args.run_tests:
        from contextlib import contextmanager

        @contextmanager
        def expect_assertion_error(msg_expected):
            try:
                yield
            except AssertionError as e:
                msg_actual = str(e).splitlines()[0]
                assert (
                    msg_actual == msg_expected
                ), "Failed to match the assertion error message"
                return
            raise AssertionError(
                "Expected the body to raise an AssertionError, but it didn't"
            )

        def run_tests():
            ex_func = LuaDefinition(
                c_name="_et_G_Damage",
                c_comment="""/**
        * Damages target entity
        * on behalf of the attacker entity.
        *
        * Such is life
        *
        * @lua_def_prototype et.G_Damage(target, inflictor, attacker, damage, dflags, mod)
        * @lua_def ---@param target number The entity number to damage.
        * @lua_def ---@param inflictor number The entity number that does the damage.
        * @lua_def ---@param attacker number The entity number that causes the *inflictor* entity to cause damage to *target*.
        * @lua_def ---@param damage number The amount of damage to inflict.
        * @lua_def ---@param dflags number The type of damage to inflict. See `Damage bitflags` for possible values.
        * @lua_def ---@param mod number The means of death. See `et.MOD_*` constants for possible values.
        * @lua_def ---@return Your mother.
        */""",
                type=DefinitionType.FUNCTION,
            )
            # pprint(ex_func.lua_comment)

            assert (
                ex_func.lua_name == "et.G_Damage"
            ), "Expected to successfully wrangle lua_name from c_name"
            assert (
                ex_func.lua_comment.description
                == "Damages target entity\non behalf of the attacker entity.\n\nSuch is life\n"
            ), "Failed to match description"

            with expect_assertion_error("Parsed Name doesn't match wrangled name"):
                LuaDefinition(
                    c_name="_et_G_Damage",
                    c_comment="""/**
            * @lua_def_prototype et.GG_Damage(target, inflictor, attacker, damage, dflags, mod)
            */""",
                    type=DefinitionType.FUNCTION,
                )

            with expect_assertion_error(
                "Expected to match both a start and an end marker. ['/**', '*/']"
            ):
                LuaDefinition(
                    c_name="_et_G_Damage",
                    c_comment="""/**
            * @lua_def_prototype et.G_Damage(target, inflictor, attacker, damage, dflags, mod)""",
                    type=DefinitionType.FUNCTION,
                )

            with expect_assertion_error(
                "Found a ---@param annotation that tried to annotate a missing parameter"
            ):
                LuaDefinition(
                    c_name="_et_G_Damage",
                    c_comment="""/**
            * @lua_def_prototype et.G_Damage(target, inflictor, attacker, damage, dflags, mod)
            * @lua_def ---@param tarrrrrrrget number The entity number to damage.
                    */""",
                    type=DefinitionType.FUNCTION,
                )

        run_tests()
        print("Unit tests successfully passed.")
    else:
        main(args)


if __name__ == "__main__":
    cli(sys.argv[1:])
