#!/usr/bin/env python3

import argparse

import oead
from colorama import Fore
import cxxfilt
from pathlib import Path
from typing import Dict, Iterable

import common.util.checker
import common.util.elf
from common.util import utils
import ai_common


def identify(functions: Dict[str, utils.FunctionInfo], checker: common.util.checker.FunctionChecker,
             new_matches: Dict[int, str], class_names: Iterable[str], get_pairs) -> None:
    for name in class_names:
        orig_name = name
        name = name[0].upper() + name[1:]

        pairs = get_pairs(orig_name, name)
        for orig_fn_name, fn_name in pairs:
            orig_fn_info = functions.get(orig_fn_name, None)
            if orig_fn_info is None:
                continue
            if orig_fn_info.status != utils.FunctionStatus.NotDecompiled:
                continue

            orig_fn = common.util.elf.get_fn_from_base_elf(orig_fn_info.addr, orig_fn_info.size)
            try:
                decomp_fn = common.util.elf.get_fn_from_my_elf(fn_name)
            except KeyError:
                continue

            if checker.check(orig_fn, decomp_fn):
                new_matches[orig_fn_info.addr] = fn_name
                utils.print_note(f"new match: {Fore.BLUE}{cxxfilt.demangle(fn_name)}{Fore.RESET}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Identifies matching AI class functions.")
    parser.add_argument("aidef")
    parser.add_argument("--type", choices=["Action", "AI", "Behavior", "Query"], required=True)
    args = parser.parse_args()

    type_: str = args.type

    new_matches: Dict[int, str] = dict()
    checker = common.util.checker.FunctionChecker()
    functions: Dict[str, utils.FunctionInfo] = {fn.name: fn for fn in utils.get_functions()}

    aidef = oead.byml.from_text(Path(args.aidef).read_text(encoding="utf-8"))

    def get_query_pairs(orig_name, name):
        prefix = f"AI_Query_{orig_name}::"
        return [
            (f"{prefix}ctor", f"_ZN5uking5query{len(name)}{name}C1ERKN4ksys3act2ai5Query7InitArgE"),
            (f"{prefix}dtor", f"_ZN5uking5query{len(name)}{name}D1Ev"),
            (f"{prefix}dtorDelete", f"_ZN5uking5query{len(name)}{name}D0Ev"),
            (f"{prefix}m10", f"_ZN5uking5query{len(name)}{name}10loadParamsERKN4evfl8QueryArgE"),
            (f"{prefix}loadParams", f"_ZN5uking5query{len(name)}{name}10loadParamsEv"),
            (f"{prefix}rtti1",
             f"_ZNK5uking5query{len(name)}{name}27checkDerivedRuntimeTypeInfoEPKN4sead15RuntimeTypeInfo9InterfaceE"),
            (f"{prefix}rtti2", f"_ZNK5uking5query{len(name)}{name}18getRuntimeTypeInfoEv"),
            (f"AI_F_Query_{orig_name}",
             f"_ZN4ksys3act2ai12QueryFactory4makeIN5uking5query{len(name)}{name}EEEPNS1_5QueryERKNS7_7InitArgEPN4sead4HeapE"),
        ]

    def get_action_pairs(orig_name, name):
        pairs = []

        def add_pair(x):
            pairs.append((x, x))

        pairs.append(
            (f"AI_Action_{orig_name}::ctor",
             f"_ZN5uking6action{len(name)}{name}C1ERKN4ksys3act2ai10ActionBase7InitArgE"))
        pairs.append(
            (f"AI_Action{orig_name}::ctor",
             f"_ZN5uking6action{len(name)}{name}C1ERKN4ksys3act2ai10ActionBase7InitArgE"))
        pairs.append((f"AI_F_Action_{orig_name}",
                      f"_ZN4ksys3act2ai13ActionFactory4makeIN5uking6action{len(name)}{name}EEEPNS1_6ActionERKNS1_10ActionBase7InitArgEPN4sead4HeapE"))
        add_pair(f"_ZN5uking6action{len(name)}{name}D1Ev")
        add_pair(f"_ZN5uking6action{len(name)}{name}D0Ev")
        add_pair(f"_ZN5uking6action{len(name)}{name}11loadParams_Ev")
        add_pair(f"_ZN5uking6action{len(name)}{name}5init_EPN4sead4HeapE")
        add_pair(f"_ZN5uking6action{len(name)}{name}6enter_EPN4ksys3act2ai15InlineParamPackE")
        add_pair(f"_ZN5uking6action{len(name)}{name}6leave_Ev")
        add_pair(f"_ZN5uking6action{len(name)}{name}5calc_Ev")
        add_pair(
            f"_ZNK5uking6action{len(name)}{name}27checkDerivedRuntimeTypeInfoEPKN4sead15RuntimeTypeInfo9InterfaceE")
        add_pair(f"_ZNK5uking6action{len(name)}{name}18getRuntimeTypeInfoEv")
        return pairs

    def get_ai_pairs(orig_name, name):
        pairs = []

        def add_pair(x):
            pairs.append((x, x))

        pairs.append(
            (f"AI_AI_{orig_name}::ctor", f"_ZN5uking2ai{len(name)}{name}C1ERKN4ksys3act2ai10ActionBase7InitArgE"))
        pairs.append(
            (f"AI_AI{orig_name}::ctor", f"_ZN5uking2ai{len(name)}{name}C1ERKN4ksys3act2ai10ActionBase7InitArgE"))
        pairs.append((f"AI_F_AI_{orig_name}",
                      f"_ZN4ksys3act2ai9AiFactory4makeIN5uking2ai{len(name)}{name}EEEPNS1_2AiERKNS1_10ActionBase7InitArgEPN4sead4HeapE"))
        add_pair(f"_ZN5uking2ai{len(name)}{name}D1Ev")
        add_pair(f"_ZN5uking2ai{len(name)}{name}D0Ev")
        add_pair(f"_ZN5uking2ai{len(name)}{name}11loadParams_Ev")
        add_pair(f"_ZN5uking2ai{len(name)}{name}5init_EPN4sead4HeapE")
        add_pair(f"_ZN5uking2ai{len(name)}{name}6enter_EPN4ksys3act2ai15InlineParamPackE")
        add_pair(f"_ZN5uking2ai{len(name)}{name}6leave_Ev")
        add_pair(f"_ZN5uking2ai{len(name)}{name}5calc_Ev")
        add_pair(f"_ZNK5uking2ai{len(name)}{name}27checkDerivedRuntimeTypeInfoEPKN4sead15RuntimeTypeInfo9InterfaceE")
        add_pair(f"_ZNK5uking2ai{len(name)}{name}18getRuntimeTypeInfoEv")
        return pairs

    if type_ == "Action":
        action_vtable_names = ai_common.get_action_vtable_names()
        identify(functions, checker, new_matches, action_vtable_names.values(), get_action_pairs)
    if type_ == "AI":
        ai_vtable_names = ai_common.get_ai_vtable_names()
        identify(functions, checker, new_matches, ai_vtable_names.values(), get_ai_pairs)
    elif type_ == "Query":
        identify(functions, checker, new_matches, aidef["Querys"].keys(), get_query_pairs)

    utils.add_decompiled_functions(new_matches)


if __name__ == "__main__":
    main()