"""MCP Server Generator — scaffolds new MCP servers from templates.

Available as a tool for the Architect to create project-specific MCP servers.
Generated servers use FastMCP and are auto-registered in config/mcp.yaml.
"""

from __future__ import annotations

import textwrap

import yaml
from langchain_core.tools import tool

from tmb.paths import TMB_ROOT, runtime_dir, user_cfg_dir

_TEMPLATES = {
    "rest_api": textwrap.dedent('''\
        """MCP Server: {name} — REST API wrapper.

        Auto-generated by TMB. Edit as needed.
        Start: uv run mcp_servers/{name}/server.py
        """

        import httpx
        from mcp.server.fastmcp import FastMCP

        mcp = FastMCP("{name}")

        BASE_URL = "{base_url}"
        HEADERS = {{"Authorization": "Bearer ${{API_TOKEN}}"}}


        @mcp.tool()
        def list_items(endpoint: str = "/items", limit: int = 20) -> str:
            """Fetch items from the API.

            Args:
                endpoint: API endpoint path (default: /items)
                limit: Maximum items to return
            """
            resp = httpx.get(f"{{BASE_URL}}{{endpoint}}", headers=HEADERS, params={{"limit": limit}})
            resp.raise_for_status()
            return resp.text


        @mcp.tool()
        def get_item(endpoint: str, item_id: str) -> str:
            """Fetch a single item by ID.

            Args:
                endpoint: API endpoint path (e.g. /items)
                item_id: The item identifier
            """
            resp = httpx.get(f"{{BASE_URL}}{{endpoint}}/{{item_id}}", headers=HEADERS)
            resp.raise_for_status()
            return resp.text


        @mcp.tool()
        def create_item(endpoint: str, data: str) -> str:
            """Create a new item via POST.

            Args:
                endpoint: API endpoint path (e.g. /items)
                data: JSON string of the item data
            """
            import json
            resp = httpx.post(
                f"{{BASE_URL}}{{endpoint}}",
                headers={{**HEADERS, "Content-Type": "application/json"}},
                content=data,
            )
            resp.raise_for_status()
            return resp.text


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

    "database": textwrap.dedent('''\
        """MCP Server: {name} — Database connector.

        Auto-generated by TMB. Edit as needed.
        Start: uv run mcp_servers/{name}/server.py
        """

        import json
        import sqlite3
        from mcp.server.fastmcp import FastMCP

        mcp = FastMCP("{name}")

        DB_PATH = "{db_path}"


        def _get_conn():
            conn = sqlite3.connect(DB_PATH)
            conn.row_factory = sqlite3.Row
            return conn


        @mcp.tool()
        def query(sql: str, params: str = "[]") -> str:
            """Run a read-only SQL query.

            Args:
                sql: SQL SELECT statement
                params: JSON array of query parameters
            """
            if not sql.strip().upper().startswith("SELECT"):
                return "[error] Only SELECT queries are allowed."
            conn = _get_conn()
            rows = conn.execute(sql, json.loads(params)).fetchall()
            conn.close()
            return json.dumps([dict(r) for r in rows], indent=2)


        @mcp.tool()
        def list_tables() -> str:
            """List all tables in the database."""
            conn = _get_conn()
            rows = conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
            ).fetchall()
            conn.close()
            return json.dumps([r["name"] for r in rows])


        @mcp.tool()
        def describe_table(table_name: str) -> str:
            """Show columns and types for a table.

            Args:
                table_name: Name of the table to describe
            """
            conn = _get_conn()
            rows = conn.execute(f"PRAGMA table_info({{table_name}})").fetchall()
            conn.close()
            return json.dumps(
                [{{"name": r["name"], "type": r["type"], "notnull": bool(r["notnull"])}} for r in rows],
                indent=2,
            )


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

    "file_based": textwrap.dedent('''\
        """MCP Server: {name} — File-based resource server.

        Auto-generated by TMB. Edit as needed.
        Start: uv run mcp_servers/{name}/server.py
        """

        from pathlib import Path
        from mcp.server.fastmcp import FastMCP

        mcp = FastMCP("{name}")

        ROOT_DIR = Path("{root_dir}")


        @mcp.tool()
        def list_files(pattern: str = "*") -> str:
            """List files matching a glob pattern.

            Args:
                pattern: Glob pattern (default: * for all files in root)
            """
            files = sorted(ROOT_DIR.glob(pattern))
            return "\\n".join(str(f.relative_to(ROOT_DIR)) for f in files[:100])


        @mcp.tool()
        def read_file(path: str) -> str:
            """Read a file's contents.

            Args:
                path: Relative path from the root directory
            """
            target = (ROOT_DIR / path).resolve()
            if not str(target).startswith(str(ROOT_DIR.resolve())):
                return "[error] Path escapes root directory."
            if not target.exists():
                return f"[error] File not found: {{path}}"
            return target.read_text()


        @mcp.resource("{name}://files")
        def resource_files() -> str:
            """List all files in the root directory."""
            return list_files("**/*")


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


def generate_mcp_server(
    name: str,
    template: str,
    description: str = "",
    **kwargs,
) -> str:
    """Generate an MCP server from a template and register it in config/mcp.yaml.

    Args:
        name: Server identifier (kebab-case)
        template: Template name — 'rest_api', 'database', or 'file_based'
        description: What this server does
        **kwargs: Template-specific variables (e.g. base_url, db_path, root_dir)

    Returns:
        Path to the generated server file
    """
    if template not in _TEMPLATES:
        raise ValueError(f"Unknown template: {template}. Available: {list(_TEMPLATES.keys())}")

    server_dir = runtime_dir() / "mcp_servers" / name
    server_dir.mkdir(parents=True, exist_ok=True)

    content = _TEMPLATES[template].format(name=name, **kwargs)
    server_file = server_dir / "server.py"
    server_file.write_text(content)

    _register_in_config(name, str(server_file))

    return str(server_file)


def _register_in_config(name: str, server_path: str):
    """Add a generated server to config/mcp.yaml."""
    config_path = user_cfg_dir() / "mcp.yaml"
    if config_path.exists():
        raw = yaml.safe_load(config_path.read_text()) or {}
    else:
        raw = {}

    servers = raw.get("servers") or {}
    if name not in servers:
        servers[name] = {
            "command": "uv",
            "args": ["run", server_path],
            "agents": ["executor"],
        }
        raw["servers"] = servers
        config_path.write_text(yaml.dump(raw, default_flow_style=False, sort_keys=False))


def create_mcp_generate_tool():
    """Return a LangChain tool for MCP server generation."""

    @tool
    def mcp_generate(name: str, template: str, description: str = "",
                     base_url: str = "", db_path: str = "", root_dir: str = "") -> str:
        """Generate a new MCP server from a template.

        Creates a FastMCP server file in mcp_servers/<name>/server.py and
        registers it in config/mcp.yaml.

        Args:
            name: Server identifier (kebab-case, e.g. 'my-api')
            template: Template to use — 'rest_api', 'database', or 'file_based'
            description: What this server does
            base_url: (rest_api) Base URL for the API
            db_path: (database) Path to the SQLite database
            root_dir: (file_based) Root directory to serve files from
        """
        kwargs = {}
        if base_url:
            kwargs["base_url"] = base_url
        if db_path:
            kwargs["db_path"] = db_path
        if root_dir:
            kwargs["root_dir"] = root_dir

        try:
            path = generate_mcp_server(name, template, description, **kwargs)
            return f"Generated MCP server at {path}. Registered in config/mcp.yaml."
        except Exception as e:
            return f"[error] Failed to generate MCP server: {e}"

    return mcp_generate
