"""MCP config wrapping and unwrapping for AvaKill proxy interception.

Rewrites agent MCP config files so all server traffic routes through
the AvaKill proxy, and can reverse the process to restore originals.
"""

from __future__ import annotations

import json
import logging
import shutil
from pathlib import Path

from avakill.mcp.config import MCPConfig, MCPServerEntry, is_already_wrapped

logger = logging.getLogger("avakill.mcp.wrapper")


def wrap_mcp_config(
    config: MCPConfig,
    policy: str | Path,
    *,
    daemon: bool = False,
    log_db: str | Path | None = None,
) -> MCPConfig:
    """Wrap all MCP server entries to route through AvaKill proxy.

    For each stdio server::

        command: "npx" args: ["-y", "@anthropic/mcp-fs", "/path"]

    Becomes::

        command: "avakill" args: ["mcp-proxy", "--policy", "...",
                 "--upstream-cmd", "npx", "--upstream-args", "-y @anthropic/mcp-fs /path"]

    This is idempotent — already-wrapped entries are skipped.
    """
    wrapped_servers: list[MCPServerEntry] = []

    for server in config.servers:
        if is_already_wrapped(server):
            wrapped_servers.append(server)
            continue

        if server.transport != "stdio":
            # Non-stdio servers pass through unchanged for now
            wrapped_servers.append(server)
            continue

        args: list[str] = ["mcp-proxy"]

        if daemon:
            args.extend(["--daemon", "~/.avakill/avakill.sock"])
        else:
            args.extend(["--policy", str(policy)])

        if log_db:
            args.extend(["--log-db", str(log_db)])

        # Preserve original command as upstream
        args.extend(["--upstream-cmd", server.command])
        if server.args:
            args.extend(["--upstream-args", " ".join(server.args)])

        wrapped = MCPServerEntry(
            name=server.name,
            command="avakill",
            args=args,
            env=server.env,
            transport=server.transport,
        )
        wrapped_servers.append(wrapped)

    return MCPConfig(
        agent=config.agent,
        config_path=config.config_path,
        servers=wrapped_servers,
    )


def unwrap_mcp_config(config: MCPConfig) -> MCPConfig:
    """Reverse wrap_mcp_config: restore original server commands.

    Extracts the ``--upstream-cmd`` and ``--upstream-args`` from wrapped
    entries and rebuilds the original server config.
    """
    unwrapped_servers: list[MCPServerEntry] = []

    for server in config.servers:
        if not is_already_wrapped(server):
            unwrapped_servers.append(server)
            continue

        # Parse the avakill mcp-proxy args to extract upstream info
        upstream_cmd = ""
        upstream_args: list[str] = []

        args_iter = iter(server.args)
        for arg in args_iter:
            if arg == "--upstream-cmd":
                upstream_cmd = next(args_iter, "")
            elif arg == "--upstream-args":
                raw = next(args_iter, "")
                upstream_args = raw.split() if raw else []

        if not upstream_cmd:
            # Can't unwrap — leave as-is
            unwrapped_servers.append(server)
            continue

        restored = MCPServerEntry(
            name=server.name,
            command=upstream_cmd,
            args=upstream_args,
            env=server.env,
            transport=server.transport,
        )
        unwrapped_servers.append(restored)

    return MCPConfig(
        agent=config.agent,
        config_path=config.config_path,
        servers=unwrapped_servers,
    )


def write_mcp_config(config: MCPConfig) -> None:
    """Write the modified config back to disk.

    Creates a ``.bak`` backup of the original file before overwriting.
    Preserves non-MCP keys in the original config.
    """
    config_path = config.config_path
    if not config_path.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")

    # Read original to preserve non-MCP keys
    raw = json.loads(config_path.read_text())

    # Create backup
    backup_path = config_path.with_suffix(config_path.suffix + ".bak")
    shutil.copy2(config_path, backup_path)
    logger.info("Backup created: %s", backup_path)

    # Rebuild mcpServers dict
    servers_dict: dict[str, dict] = {}
    for server in config.servers:
        entry: dict = {"command": server.command}
        if server.args:
            entry["args"] = server.args
        if server.env:
            entry["env"] = server.env
        if server.url:
            entry["url"] = server.url
        if server.transport != "stdio":
            entry["transport"] = server.transport
        servers_dict[server.name] = entry

    # Write back — detect which key the original used
    if "mcpServers" in raw:
        raw["mcpServers"] = servers_dict
    elif "servers" in raw:
        raw["servers"] = servers_dict
    else:
        raw["mcpServers"] = servers_dict

    config_path.write_text(json.dumps(raw, indent=2) + "\n")
    logger.info("Config written: %s", config_path)
