# SPDX-FileCopyrightText: © 2026 Joe T. Sylve, Ph.D. <joe.sylve@gmail.com>
#
# SPDX-License-Identifier: MIT

"""Function prototype and calling convention tools."""

from __future__ import annotations

import ida_nalt
import ida_typeinf
import idc
from mcp.server.fastmcp import FastMCP

from ida_mcp.helpers import format_address, get_func_name, resolve_function
from ida_mcp.session import session

_CC_NAMES = {
    ida_typeinf.CM_CC_CDECL: "cdecl",
    ida_typeinf.CM_CC_STDCALL: "stdcall",
    ida_typeinf.CM_CC_PASCAL: "pascal",
    ida_typeinf.CM_CC_FASTCALL: "fastcall",
    ida_typeinf.CM_CC_THISCALL: "thiscall",
}

_CC_MAP = {v: k for k, v in _CC_NAMES.items()}


def register(mcp: FastMCP):
    @mcp.tool()
    @session.require_open
    def get_function_type(address: str) -> dict:
        """Get the full type signature (prototype) of a function.

        Returns the return type, parameters, and calling convention.

        Args:
            address: Address or name of the function.
        """
        func, err = resolve_function(address)
        if err:
            return err

        tinfo = ida_typeinf.tinfo_t()
        if not ida_nalt.get_tinfo(tinfo, func.start_ea) and not ida_typeinf.guess_tinfo(
            tinfo, func.start_ea
        ):
            type_str = idc.get_type(func.start_ea) or ""
            return {
                "address": format_address(func.start_ea),
                "name": get_func_name(func.start_ea),
                "type": type_str,
                "details": None,
            }

        # Extract function details
        fi = ida_typeinf.func_type_data_t()
        if tinfo.get_func_details(fi):
            params = []
            for i in range(fi.size()):
                param = fi[i]
                params.append(
                    {
                        "name": param.name or f"arg{i}",
                        "type": str(param.type),
                    }
                )

            return {
                "address": format_address(func.start_ea),
                "name": get_func_name(func.start_ea),
                "type": str(tinfo),
                "return_type": str(fi.rettype),
                "calling_convention": _CC_NAMES.get(fi.get_cc() & 0xF0, f"cc_{fi.get_cc():#x}"),
                "parameters": params,
            }

        return {
            "address": format_address(func.start_ea),
            "name": get_func_name(func.start_ea),
            "type": str(tinfo),
            "details": None,
        }

    @mcp.tool()
    @session.require_open
    def set_function_type(address: str, type_string: str) -> dict:
        """Set the full type signature (prototype) of a function.

        Args:
            address: Address or name of the function.
            type_string: C function declaration, e.g. "int __cdecl foo(int a, char *b)".
        """
        func, err = resolve_function(address)
        if err:
            return err

        success = idc.SetType(func.start_ea, type_string)
        if not success:
            return {
                "error": "Failed to set function type",
                "error_type": "SetTypeFailed",
            }

        return {
            "address": format_address(func.start_ea),
            "name": get_func_name(func.start_ea),
            "type": type_string,
        }

    @mcp.tool()
    @session.require_open
    def set_function_calling_convention(address: str, convention: str) -> dict:
        """Change the calling convention of a function.

        Args:
            address: Address or name of the function.
            convention: Calling convention — "cdecl", "stdcall", "fastcall", "thiscall", "pascal".
        """
        func, err = resolve_function(address)
        if err:
            return err

        cc_val = _CC_MAP.get(convention.lower())
        if cc_val is None:
            return {
                "error": f"Unknown calling convention: {convention!r}",
                "error_type": "InvalidArgument",
                "valid_conventions": list(_CC_MAP.keys()),
            }

        tinfo = ida_typeinf.tinfo_t()
        if not ida_nalt.get_tinfo(tinfo, func.start_ea) and not ida_typeinf.guess_tinfo(
            tinfo, func.start_ea
        ):
            return {
                "error": "Cannot determine function type to change convention",
                "error_type": "NoType",
            }

        fi = ida_typeinf.func_type_data_t()
        if not tinfo.get_func_details(fi):
            return {"error": "Cannot get function details", "error_type": "NoType"}

        fi.set_cc((fi.get_cc() & 0x0F) | cc_val)
        new_tinfo = ida_typeinf.tinfo_t()
        if not new_tinfo.create_func(fi):
            return {"error": "Failed to create new function type", "error_type": "CreateFailed"}

        success = ida_typeinf.apply_tinfo(func.start_ea, new_tinfo, ida_typeinf.TINFO_DEFINITE)
        if not success:
            return {
                "error": f"Failed to apply calling convention {convention!r}",
                "error_type": "ApplyFailed",
            }
        return {
            "address": format_address(func.start_ea),
            "name": get_func_name(func.start_ea),
            "convention": convention,
        }
