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

import argparse
import re
from pathlib import Path


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

SOURCE_RE = re.compile(r"SOURCE_SERVER\s*=\s*Path\((['\"])(.*?)\1\)")
SERVER_RE = re.compile(r"SERVER_NAME\s*=\s*(['\"])(.*?)\1")

SHIM_TEMPLATE = """#!/usr/bin/env python3
from __future__ import annotations

import ast
from pathlib import Path

from mcp.server.fastmcp import FastMCP


SOURCE_SERVER = Path({source_server_path!r})
LOCAL_SERVER = Path(__file__).with_name(SOURCE_SERVER.name)
SERVER_NAME = {server_name!r}


class _ShimMCP:
    @staticmethod
    def tool(*args, **kwargs):
        if args and callable(args[0]) and len(args) == 1 and not kwargs:
            return args[0]
        def _decorator(fn):
            return fn
        return _decorator


def _resolve_source_server():
    if LOCAL_SERVER.exists() and LOCAL_SERVER.name != Path(__file__).name:
        return LOCAL_SERVER
    return SOURCE_SERVER


def _load_functions():
    source_server = _resolve_source_server()
    code = source_server.read_text(encoding="utf-8")
    tree = ast.parse(code, filename=str(source_server))
    function_names = [n.name for n in tree.body if isinstance(n, ast.FunctionDef) and not n.name.startswith("_")]
    namespace = {{
        "__name__": "__mcp_source__",
        "mcp": _ShimMCP(),
    }}
    exec(compile(code, str(source_server), "exec"), namespace, namespace)
    loaded = []
    for name in function_names:
        fn = namespace.get(name)
        if callable(fn):
            loaded.append(fn)
    return loaded


mcp = FastMCP(SERVER_NAME)
for _fn in _load_functions():
    mcp.tool()(_fn)


if __name__ == "__main__":
    mcp.run(transport="stdio")
"""


def extract_metadata(shim_path: Path) -> tuple[str, str] | None:
    text = shim_path.read_text(encoding="utf-8")
    source_match = SOURCE_RE.search(text)
    server_match = SERVER_RE.search(text)
    if not source_match or not server_match:
        return None
    return source_match.group(2), server_match.group(2)


def rewrite_shim(shim_path: Path, dry_run: bool = False) -> bool:
    metadata = extract_metadata(shim_path)
    if metadata is None:
        return False
    source_server_path, server_name = metadata
    new_text = SHIM_TEMPLATE.format(source_server_path=source_server_path, server_name=server_name)
    current = shim_path.read_text(encoding="utf-8")
    if current == new_text:
        return False
    if not dry_run:
        shim_path.write_text(new_text, encoding="utf-8")
    return True


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Rewrite generated MCP shim servers to prefer local patched raw servers."
    )
    parser.add_argument(
        "--mcp-root",
        default=str(DEFAULT_MCP_ROOT),
        help="Root directory containing mcp_* server dirs.",
    )
    parser.add_argument("--dry-run", action="store_true")
    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}")

    scanned = 0
    changed = 0
    for shim_path in sorted(mcp_root.glob("mcp_*/app/*_shim_server.py")):
        scanned += 1
        if rewrite_shim(shim_path, dry_run=args.dry_run):
            changed += 1

    print(f"Scanned {scanned} shim files.")
    print(f"Rewrote {changed} shim files.")


if __name__ == "__main__":
    main()