"""
name: MCP assistant
category: 
author: malcat
icon: wxART_SYMBOLS_FN

Expose the current analysis object to LLM agents by spawning a local MCP server, based on ZeroMCP. You can then connect your LLM agent (e.g. Claude Code) to it to help you analyse and edit the file.
"""
import json
import os
import sys
import time
import types
if sys.version_info < (3, 11):
    from typing_extensions import TypedDict
    from typing import Dict, Optional, List, Tuple, Literal, Any, Union, Annotated
else:
    from typing import Dict, Optional, List, Tuple, Literal, Any, Union, Annotated, TypedDict
import malcat
malcat.setup()
from mcp_common import MalcatMCPServer, register_tools_common, MalcatAnalysisResult, smart_ea_convert


MCP_SERVER_HOST = "127.0.0.1"
MCP_SERVER_PORT = 9009


class MalcatGUIMCPServer(MalcatMCPServer):

    def __init__(self):
        MalcatMCPServer.__init__(self, 5, None)

    def _mcp_tools_call(self, name, arguments, _meta):
        gui.print(f"Calling {name} with parameters: {json.dumps(arguments)} ... ", format=False, end="")
        if name != "show_user" and "analysis_id" in arguments and "ea" in arguments and arguments["analysis_id"] == analysis.id:
            try:
                # follow addresses in the gui
                ea = smart_ea_convert(analysis, arguments["ea"])
                gui.goto(ea)
            except: pass
        # quick hack to print out tools access and errors with nice colors
        res = super()._mcp_tools_call(name, arguments, _meta)
        if res["isError"]:
            gui.print("[error]FAIL[/error]", format=True)
            errortxt = res['content'][0]['text']
            if not "McpToolError:" in errortxt:
                gui.print("[error]{}[/error]".format(errortxt.replace("[", "\\[")), format=True)
        else:
            gui.print("[clean]OK[/clean]", format=True)
        return res

    

# ============================================================================
# Main
# ============================================================================

if __name__ == "__main__":
    mcp = MalcatGUIMCPServer()
    mcp.register_analysis(analysis)     # register current analysis
    
    register_tools_common(mcp)

    @mcp.tool
    def stop_server():
        """
        Stop the MCP server
        """
        mcp.stop()    

    @mcp.tool
    def get_current_analysis() -> int:
        """
        Retrieve the ID of the analysis currently open in Malcat's GUI
        """
        return analysis.id

    class UserPositionInGui(TypedDict):
        ea: Annotated[int, "effective address where the user is looking at"]
        englobing_function_ea: Annotated[Optional[int], "if in a function, start of the function"]
        englobing_function_name: Annotated[Optional[str], "if in a function, name of the function"]

    @mcp.tool
    def get_current_pos() -> UserPositionInGui:
        """
        Get the effective address of where the user is looking at
        """
        args = {
            "ea": gui.cursor
        }
        fn = analysis.fns.find(gui.cursor)
        if fn:
            args["englobing_function_ea"] = fn.address
            args["englobing_function_name"] = fn.name
        return UserPositionInGui(**args)

    @mcp.tool
    def show_user(
            ea:Annotated[Union[int, str], "effective address of what to show"],
            size:Annotated[Optional[int], "if you want to show an interval, size of interval"] = 0,
            view_type:Annotated[Optional[str], "what view to use: AUTO, HEXADECIMAL, STRUCTURE, DISASSEMBLY or SOURCE_CODE"] = "AUTO"
            ):
        """
        In Malcat's GUI, show the given address to the user
        if <size> is set, will also visually select the interval [ea:ea+size]
        """
        view_type = {
            "HEXADECIMAL": malcat.ViewId.HEXADECIMAL,
            "STRUCTURE": malcat.ViewId.STRUCTURE,
            "DISASSEMBLY": malcat.ViewId.DISASSEMBLY,
            "SOURCE_CODE": malcat.ViewId.SOURCE_CODE,
            "AUTO": malcat.ViewId.AUTO,
        }.get(view_type, malcat.ViewId.AUTO)
        ea = smart_ea_convert(analysis, ea)
        gui.goto(ea, view_type)
        if size:
            gui.select(ea, size)    

    @mcp.tool
    def rename_function(
            ea:Annotated[Union[int, str], "effective address of function"],
            newname:Annotated[str, "function new name"],
            ):
        """
        Rename the function defined at <ea> in the current analysis
        """
        ea = smart_ea_convert(analysis, ea)
        fn = analysis.fns.find(ea)
        if not fn:
            raise KeyError("No function defined there")
        fn.name = newname
        gui.goto(fn.address)
        analysis.run()

    @mcp.tool
    def rename_function_decompiled_variables(
            ea:Annotated[Union[int, str], "effective address of function"],
            renaming:Annotated[Dict[str, str], "a mapping old name -> new name for some or all variables of the function"],
            ):
        """
        Rename some or all the variables in the decompiled/source_code listing of a function in the current analysis
        """
        ea = smart_ea_convert(analysis, ea)
        fn = analysis.fns.find(ea)
        if not fn:
            raise KeyError("No function defined there")
        fn.set_decompile_aliases(renaming)
        gui.goto(fn.address)
        analysis.run()

    @mcp.tool
    def open_analysis_in_gui_and_stop(analysis_id:Annotated[int, "analysis ID"]):
        """
        Open an analysis object in the GUI and show it to the user. Will stop the current MCP session
        """
        a = mcp.get_analysis(analysis_id)
        gui.open_after(a.file[:], a.file.name)
        mcp.stop()

    
    gui.print("\n[block1]=== Malcat MCP Server ===[/block1]\n\n", format=True)

    gui.print("This script spawns a local HTTP (ZeroMCP-based) MCP server, allowing any LLM agent to interact with the [color2]current analysis[/color2] loaded in Malcat's GUI. Note that this is in [color2]beta[/color2] for now.\n\n", format=True)


    gui.print(f"If you want to register this MCP server using Claude Code, you can type (once): \n\n\t[color2]claude mcp add --transport http malcatgui http://{MCP_SERVER_HOST}:{MCP_SERVER_PORT}/mcp[/color2]\n", format=True)
    print("To stop the server: press the stop button on the left of this text or prompt: 'stop the MCP server'\n\n")

    mcp.serve(MCP_SERVER_HOST, MCP_SERVER_PORT, background=True)
    
    try:
        while mcp._running:
            gui.idle()
            time.sleep(0.5)
    finally:
        mcp.stop()
