#!/usr/bin/env python3
from __future__ import annotations

import argparse
import re
from pathlib import Path


DEFAULT_MCP_ROOT = (
    Path(__file__).resolve().parent
    / "biomni_web"
    / "backend"
    / "data"
    / "mcp_generated"
)


TOOL_DECORATOR_RE = re.compile(r"^(\s*)@mcp\.tool\(\)\s*$")
COMMENTED_TOOL_DECORATOR_RE = re.compile(r"^(\s*)#\s*@mcp\.tool\(\)\s*$")


def build_server_name(server_file: Path) -> str:
    stem = server_file.stem
    if stem.endswith("_server"):
        stem = stem[: -len("_server")]
    return f"local_{stem.replace('-', '_').replace('.', '_')}"


def patch_file(server_file: Path) -> bool:
    text = server_file.read_text(encoding="utf-8")
    original = text

    if "_shim_server.py" in server_file.name:
        return False

    if "mcp = FastMCP(" in text and 'mcp.run(transport="stdio")' in text:
        return False

    lines = text.splitlines()

    # Enable commented decorators first so the file becomes a real MCP server.
    lines = [COMMENTED_TOOL_DECORATOR_RE.sub(r"\1@mcp.tool()", line) for line in lines]

    first_decorator_idx = None
    for idx, line in enumerate(lines):
        if TOOL_DECORATOR_RE.match(line):
            first_decorator_idx = idx
            break

    if first_decorator_idx is None:
        return False

    if "from mcp.server.fastmcp import FastMCP" not in "\n".join(lines):
        insertion = [
            "from mcp.server.fastmcp import FastMCP",
            "",
            f"SERVER_NAME = {build_server_name(server_file)!r}",
            "mcp = FastMCP(SERVER_NAME)",
            "",
        ]
        lines[first_decorator_idx:first_decorator_idx] = insertion

    patched = "\n".join(lines).rstrip() + "\n"

    if 'mcp.run(transport="stdio")' not in patched:
        patched += '\nif __name__ == "__main__":\n    mcp.run(transport="stdio")\n'

    if patched != original:
        server_file.write_text(patched, encoding="utf-8")
        return True

    return False


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Patch raw MCP server files to be directly runnable with FastMCP."
    )
    parser.add_argument(
        "--mcp-root",
        default=str(DEFAULT_MCP_ROOT),
        help="Root directory containing mcp_* subdirectories.",
    )
    args = parser.parse_args()

    mcp_root = Path(args.mcp_root).resolve()
    if not mcp_root.is_dir():
        raise SystemExit(f"MCP root does not exist or is not a directory: {mcp_root}")

    changed = 0
    scanned = 0
    for server_file in sorted(mcp_root.glob("mcp_*/app/*_server.py")):
        scanned += 1
        if patch_file(server_file):
            changed += 1

    print(f"Scanned {scanned} server files.")
    print(f"Patched {changed} raw MCP server files.")


if __name__ == "__main__":
    main()