import json
import os
from pathlib import Path

from backend.config import MCP_SUPPORT, _load_settings


_mcp_tool_cache: dict[str, dict] = {}


def _mcp_effective_env(command: str, srv_env: dict | None = None, *, base_env: dict | None = None) -> dict:
    merged: dict[str, str] = {
        str(key): str(value)
        for key, value in dict(base_env or os.environ).items()
    }

    for raw_key, raw_value in dict(srv_env or {}).items():
        key = str(raw_key or "").strip()
        if not key:
            continue
        merged[key] = str(raw_value or "")

    command_path = Path(str(command or "").strip())
    if command_path.is_file():
        command_dir = str(command_path.parent)
        current_path = str(merged.get("PATH") or "")
        existing_parts = [part for part in current_path.split(os.pathsep) if part]
        if command_dir and not any(part.lower() == command_dir.lower() for part in existing_parts):
            merged["PATH"] = command_dir if not current_path else command_dir + os.pathsep + current_path

    return merged


def _mcp_discover_server(server_name: str, srv_cfg: dict) -> dict:
    """Discover tools on a single MCP server. Returns cache entry dict."""
    import subprocess
    import sys
    import time as _time

    command = srv_cfg.get("command", "")
    srv_args = srv_cfg.get("args", [])
    srv_env = _mcp_effective_env(command, srv_cfg.get("env"))
    if not command:
        return {"ok": False, "tools": [], "error": "no command configured", "ts": _time.time()}
    list_script = (
        "import sys, json, asyncio\n"
        "from mcp import ClientSession, StdioServerParameters\n"
        "from mcp.client.stdio import stdio_client\n"
        "async def main():\n"
        f"    params = StdioServerParameters(command={command!r}, args={srv_args!r}, env={srv_env!r})\n"
        "    async with stdio_client(params) as (read, write):\n"
        "        async with ClientSession(read, write) as session:\n"
        "            await session.initialize()\n"
        "            tools = await session.list_tools()\n"
        "            out = []\n"
        "            for t in tools.tools:\n"
        '                out.append({"name": t.name, "description": t.description or ""})\n'
        '            print(json.dumps({"ok": True, "tools": out}, ensure_ascii=False))\n'
        "try:\n"
        "    asyncio.run(main())\n"
        "except Exception as e:\n"
        "    import json as _j\n"
        '    print(_j.dumps({"ok": False, "error": str(e)}, ensure_ascii=False))\n'
    )
    try:
        proc = subprocess.run(
            [sys.executable, "-c", list_script],
            capture_output=True,
            text=True,
            timeout=30,
        )
        output = (proc.stdout or "").strip()
        if not output:
            err = (proc.stderr or "").strip()[:500]
            return {"ok": False, "tools": [], "error": err or "no output", "ts": _time.time()}
        data = json.loads(output)
        if data.get("ok"):
            return {"ok": True, "tools": data.get("tools", []), "error": None, "ts": _time.time()}
        return {"ok": False, "tools": [], "error": data.get("error", "unknown"), "ts": _time.time()}
    except subprocess.TimeoutExpired:
        return {"ok": False, "tools": [], "error": "timeout (30s)", "ts": _time.time()}
    except Exception as e:
        return {"ok": False, "tools": [], "error": str(e), "ts": _time.time()}


def _mcp_discover_all() -> dict[str, dict]:
    """Discover tools on all configured MCP servers. Updates and returns cache."""
    if not MCP_SUPPORT:
        return {}
    settings = _load_settings()
    servers = settings.get("mcp_servers") or {}
    if not servers:
        _mcp_tool_cache.clear()
        return _mcp_tool_cache
    for old in list(_mcp_tool_cache):
        if old not in servers:
            del _mcp_tool_cache[old]
    for name, cfg in servers.items():
        if not isinstance(cfg, dict):
            continue
        result = _mcp_discover_server(name, cfg)
        _mcp_tool_cache[name] = result
        status = "ok" if result["ok"] else f"error: {result.get('error', '?')}"
        n_tools = len(result.get("tools", []))
        print(f"[mcp] {name}: {status} ({n_tools} tools)")
    return _mcp_tool_cache


def _mcp_get_cache() -> dict[str, dict]:
    """Return the current MCP discovery cache (read-only)."""
    return dict(_mcp_tool_cache)


def _mcp_call(server: str, tool: str, arguments: dict | None = None) -> str:
    """Call a tool on a configured MCP server (stdio transport)."""
    if not MCP_SUPPORT:
        return "MCP SDK is not installed. Run: pip install mcp"
    settings = _load_settings()
    servers = settings.get("mcp_servers") or {}
    srv_cfg = servers.get(server)
    if not srv_cfg:
        available = ", ".join(servers.keys()) if servers else "none"
        return f"Unknown MCP server '{server}'. Available: {available}"
    command = srv_cfg.get("command", "")
    srv_args = srv_cfg.get("args", [])
    srv_env = _mcp_effective_env(command, srv_cfg.get("env"))
    if not command:
        return f"MCP server '{server}' has no command configured."

    import subprocess
    import sys

    call_script = f"""
import sys, json, asyncio
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

async def main():
    params = StdioServerParameters(
        command={command!r},
        args={srv_args!r},
        env={srv_env!r},
    )
    async with stdio_client(params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            result = await session.call_tool({tool!r}, arguments={arguments!r})
            parts = []
            for c in result.content:
                if hasattr(c, 'text'):
                    parts.append(c.text)
                else:
                    parts.append(str(c))
            print(json.dumps({{"ok": True, "result": "\\n".join(parts)}}, ensure_ascii=False))

try:
    asyncio.run(main())
except Exception as e:
    import json as _j
    print(_j.dumps({{"ok": False, "error": str(e)}}, ensure_ascii=False))
"""
    try:
        proc = subprocess.run(
            [sys.executable, "-c", call_script],
            capture_output=True,
            text=True,
            timeout=60,
        )
        output = (proc.stdout or "").strip()
        if not output:
            err = (proc.stderr or "").strip()
            return f"MCP call failed (no output). stderr: {err[:2000]}"
        data = json.loads(output)
        if data.get("ok"):
            return data.get("result", "(empty)")
        return f"MCP error: {data.get('error', 'unknown')}"
    except subprocess.TimeoutExpired:
        return f"MCP call to '{server}.{tool}' timed out after 60s."
    except Exception as e:
        return f"MCP call error: {e}"


def _mcp_list_tools(server: str) -> str:
    """List available tools on an MCP server."""
    if not MCP_SUPPORT:
        return "MCP SDK is not installed. Run: pip install mcp"
    settings = _load_settings()
    servers = settings.get("mcp_servers") or {}
    srv_cfg = servers.get(server)
    if not srv_cfg:
        available = ", ".join(servers.keys()) if servers else "none"
        return f"Unknown MCP server '{server}'. Available: {available}"
    command = srv_cfg.get("command", "")
    srv_args = srv_cfg.get("args", [])
    srv_env = _mcp_effective_env(command, srv_cfg.get("env"))
    if not command:
        return f"MCP server '{server}' has no command configured."

    import subprocess
    import sys

    list_script = f"""
import sys, json, asyncio
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

async def main():
    params = StdioServerParameters(
        command={command!r},
        args={srv_args!r},
        env={srv_env!r},
    )
    async with stdio_client(params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            tools = await session.list_tools()
            out = []
            for t in tools.tools:
                out.append({{"name": t.name, "description": t.description or ""}})
            print(json.dumps({{"ok": True, "tools": out}}, ensure_ascii=False))

try:
    asyncio.run(main())
except Exception as e:
    import json as _j
    print(_j.dumps({{"ok": False, "error": str(e)}}, ensure_ascii=False))
"""
    try:
        proc = subprocess.run(
            [sys.executable, "-c", list_script],
            capture_output=True,
            text=True,
            timeout=30,
        )
        output = (proc.stdout or "").strip()
        if not output:
            err = (proc.stderr or "").strip()
            return f"MCP list failed. stderr: {err[:2000]}"
        data = json.loads(output)
        if data.get("ok"):
            tools = data.get("tools", [])
            if not tools:
                return f"MCP server '{server}' has no tools."
            lines = [f"Tools on MCP server '{server}':"]
            for item in tools:
                lines.append(f"  - {item['name']}: {item.get('description', '')}")
            return "\n".join(lines)
        return f"MCP error: {data.get('error', 'unknown')}"
    except subprocess.TimeoutExpired:
        return f"MCP list_tools on '{server}' timed out."
    except Exception as e:
        return f"MCP list error: {e}"