# Server Logging
# ---------------------------------------------------------------------------------
import logging
import sys
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path

import click
import pyghidra
from mcp.server import Server
from mcp.server.fastmcp import Context, FastMCP
from mcp.shared.exceptions import McpError
from mcp.types import INTERNAL_ERROR, INVALID_PARAMS, ErrorData

from pyghidra_mcp.__init__ import __version__
from pyghidra_mcp.context import ProgramInfo as ContextProgramInfo, PyGhidraContext
from pyghidra_mcp.models import (
    CrossReferenceInfos,
    DecompiledFunction,
    ExportInfos,
    FunctionSearchResults,
    ImportInfos,
    ProgramInfo,
    ProgramInfos,
    SymbolSearchResults,
)
from pyghidra_mcp.tools import GhidraTools

logging.basicConfig(
    level=logging.INFO,
    stream=sys.stderr,  # Critical for STDIO transport
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)
logger.info("Server initialized")


# Init Pyghidra
# ---------------------------------------------------------------------------------
@asynccontextmanager
async def server_lifespan(server: Server) -> AsyncIterator[PyGhidraContext]:
    """Manage server startup and shutdown lifecycle."""
    try:
        yield server._pyghidra_context
    finally:
        # pyghidra_context.close()
        pass


mcp = FastMCP("pyghidra-mcp", lifespan=server_lifespan)


def _get_program_info_or_raise(
    pyghidra_context: PyGhidraContext, binary_name: str
) -> ContextProgramInfo:
    """Get program info or raise McpError if not found."""
    program_info = pyghidra_context.programs.get(binary_name)
    if not program_info:
        available_progs = list(pyghidra_context.programs.keys())
        raise McpError(
            ErrorData(
                code=INVALID_PARAMS,
                message=f"Binary {binary_name} not found. Available binaries: {available_progs}",
            )
        )
    return program_info


# MCP Tools
# ---------------------------------------------------------------------------------
@mcp.tool()
async def decompile_function(binary_name: str, name: str, ctx: Context) -> DecompiledFunction:
    """Decompiles a function in a specified binary and returns its pseudo-C code.

    Args:
        binary_name: The name of the binary containing the function.
        name: The name of the function to decompile.
    """
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        program_info = _get_program_info_or_raise(pyghidra_context, binary_name)
        tools = GhidraTools(program_info)
        return tools.decompile_function(name)
    except Exception as e:
        raise McpError(
            ErrorData(code=INTERNAL_ERROR, message=f"Error decompiling function: {e!s}")
        ) from e


@mcp.tool()
def search_functions_by_name(
    binary_name: str, query: str, ctx: Context, offset: int = 0, limit: int = 100
) -> FunctionSearchResults:
    """Searches for functions within a binary by name.

    Args:
        binary_name: The name of the binary to search within.
        query: The substring to search for in function names (case-insensitive).
        offset: The number of results to skip.
        limit: The maximum number of results to return.
    """
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        program_info = _get_program_info_or_raise(pyghidra_context, binary_name)
        tools = GhidraTools(program_info)
        functions = tools.search_functions_by_name(query, offset, limit)
        return FunctionSearchResults(functions=functions)
    except Exception as e:
        raise McpError(
            ErrorData(code=INTERNAL_ERROR, message=f"Error searching for functions: {e!s}")
        ) from e


@mcp.tool()
def search_symbols_by_name(
    binary_name: str, query: str, ctx: Context, offset: int = 0, limit: int = 100
) -> SymbolSearchResults:
    """Searches for symbols within a binary by name.

    Args:
        binary_name: The name of the binary to search within.
        query: The substring to search for in symbol names (case-insensitive).
        offset: The number of results to skip.
        limit: The maximum number of results to return.
    """
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        program_info = _get_program_info_or_raise(pyghidra_context, binary_name)
        tools = GhidraTools(program_info)
        symbols = tools.search_symbols_by_name(query, offset, limit)
        return SymbolSearchResults(symbols=symbols)
    except Exception as e:
        raise McpError(
            ErrorData(code=INTERNAL_ERROR, message=f"Error searching for symbols: {e!s}")
        ) from e


@mcp.tool()
def list_project_binaries(ctx: Context) -> list[str]:
    """Lists the names of all binaries currently loaded in the Ghidra project."""
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        return list(pyghidra_context.programs.keys())
    except Exception as e:
        raise McpError(
            ErrorData(code=INTERNAL_ERROR, message=f"Error listing project binaries: {e!s}")
        ) from e


@mcp.tool()
def list_project_program_info(ctx: Context) -> ProgramInfos:
    """Retrieves detailed information for all programs (binaries) in the project."""
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        program_infos = []
        for _name, pi in pyghidra_context.programs.items():
            program_infos.append(
                ProgramInfo(
                    name=pi.name,
                    file_path=str(pi.file_path) if pi.file_path else None,
                    load_time=pi.load_time,
                    analysis_complete=pi.analysis_complete,
                    metadata=pi.metadata,
                )
            )
        return ProgramInfos(programs=program_infos)
    except Exception as e:
        raise McpError(
            ErrorData(
                code=INTERNAL_ERROR,
                message=f"Error listing project program info: {e!s}",
            )
        ) from e


@mcp.tool()
def list_exports(binary_name: str, ctx: Context) -> ExportInfos:
    """Lists all exported functions and symbols from a specified binary.

    Args:
        binary_name: The name of the binary to list exports from.
    """
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        program_info = _get_program_info_or_raise(pyghidra_context, binary_name)
        tools = GhidraTools(program_info)
        exports = tools.list_exports()
        return ExportInfos(exports=exports)
    except Exception as e:
        raise McpError(
            ErrorData(code=INTERNAL_ERROR, message=f"Error listing exports: {e!s}")
        ) from e


@mcp.tool()
def list_imports(binary_name: str, ctx: Context) -> ImportInfos:
    """Lists all imported functions and symbols for a specified binary.

    Args:
        binary_name: The name of the binary to list imports from.
    """
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        program_info = _get_program_info_or_raise(pyghidra_context, binary_name)
        tools = GhidraTools(program_info)
        imports = tools.list_imports()
        return ImportInfos(imports=imports)
    except Exception as e:
        raise McpError(
            ErrorData(code=INTERNAL_ERROR, message=f"Error listing imports: {e!s}")
        ) from e


@mcp.tool()
def list_cross_references(
    binary_name: str, name_or_address: str, ctx: Context
) -> CrossReferenceInfos:
    """Finds and lists all cross-references (x-refs) to a given function or address within a binary. This is crucial for understanding how code and data are used and related.

    Args:
        binary_name: The name of the binary to search for cross-references in.
        name_or_address: The name of the function or a specific address (e.g., '0x1004010') to find cross-references to.
    """
    try:
        pyghidra_context: PyGhidraContext = ctx.request_context.lifespan_context
        program_info = _get_program_info_or_raise(pyghidra_context, binary_name)
        tools = GhidraTools(program_info)
        cross_references = tools.list_cross_references(name_or_address)
        return CrossReferenceInfos(cross_references=cross_references)
    except Exception as e:
        raise McpError(
            ErrorData(code=INTERNAL_ERROR, message=f"Error listing cross-references: {e!s}")
        ) from e


def init_pyghidra_context(
    mcp: FastMCP, input_paths: list[Path], project_name: str, project_directory: str
) -> FastMCP:
    if not input_paths:
        raise ValueError("Missing Input Paths!")

    bin_paths = [Path(p) for p in input_paths]

    logger.info(f"Analyzing {', '.join(map(str, bin_paths))}")
    logger.info(f"Project: {project_name}")
    logger.info(f"Project: Location {project_directory}")

    # init pyghidra
    pyghidra.start(False)  # setting Verbose output

    # init PyGhidraContext / import + analyze binaries
    pyghidra_context = PyGhidraContext(project_name, project_directory)
    logger.info(f"Importing binaries: {project_directory}")
    pyghidra_context.import_binaries(bin_paths)
    logger.info(f"Analyize project: {pyghidra_context.project}")
    pyghidra_context.analyze_project()

    mcp._pyghidra_context = pyghidra_context

    return mcp


# MCP Server Entry Point
# ---------------------------------------------------------------------------------


@click.command(context_settings={"help_option_names": ["-h", "--help"]})
@click.version_option(
    __version__,
    "-v",
    "--version",
    help="Show version and exit.",
)
@click.option(
    "-t",
    "--transport",
    type=click.Choice(["stdio", "streamable-http", "sse"]),
    default="stdio",
    envvar="MCP_TRANSPORT",
    help="Transport protocol to use: stdio, streamable-http, or sse (legacy)",
)
@click.option(
    "--project-name",
    default="pyghidra_mcp",
    help="Name of the Ghidra project.",
)
@click.option(
    "--project-directory",
    default="pyghidra_mcp_projects",
    type=click.Path(),
    help="Directory to store the Ghidra project.",
)
@click.argument("input_paths", type=click.Path(exists=True), nargs=-1, required=True)
def main(
    transport: str, input_paths: list[Path], project_name: str, project_directory: str
) -> None:
    """PyGhidra Command-Line MCP server

    - input_paths: Path to one or more binaries to import, analyze, and expose with pyghidra-mcp
    - transport: Supports stdio, streamable-http, and sse transports.
    For stdio, it will read from stdin and write to stdout.
    For streamable-http and sse, it will start an HTTP server on port 8000.

    """

    init_pyghidra_context(mcp, input_paths, project_name, project_directory)

    try:
        if transport == "stdio":
            mcp.run(transport="stdio")
        elif transport == "streamable-http":
            mcp.run(transport="streamable-http")
        elif transport == "sse":
            mcp.run(transport="sse")
        else:
            raise ValueError(f"Invalid transport: {transport}")
    finally:
        mcp._pyghidra_context.close()


if __name__ == "__main__":
    main()
