"""MCP client — tool discovery from MCP servers."""
import json
from config import MCP_DIR, log
from database import db

async def _discover_mcp_tools(mcp_srv: dict) -> list[dict]:
    """Connect to MCP server and list available tools."""
    from mcp import ClientSession
    if mcp_srv["server_type"] in ("stdio", "module"):
        from mcp.client.stdio import stdio_client, StdioServerParameters
        cmd = mcp_srv.get("command") or "python3"
        args_raw = mcp_srv.get("args")
        args = json.loads(args_raw) if isinstance(args_raw, str) else (args_raw or [])
        env_raw = mcp_srv.get("env_vars")
        env = json.loads(env_raw) if isinstance(env_raw, str) else (env_raw or None)
        params = StdioServerParameters(command=cmd, args=args, env=env)
        async with stdio_client(params) as streams:
            async with ClientSession(*streams) as session:
                await session.initialize()
                result = await session.list_tools()
                return [{"name": t.name, "description": t.description or "",
                         "input_schema": t.inputSchema if hasattr(t, 'inputSchema') else {}} for t in result.tools]
    elif mcp_srv["server_type"] == "sse":
        from mcp.client.sse import sse_client
        async with sse_client(mcp_srv["url"]) as streams:
            async with ClientSession(*streams) as session:
                await session.initialize()
                result = await session.list_tools()
                return [{"name": t.name, "description": t.description or "",
                         "input_schema": t.inputSchema if hasattr(t, 'inputSchema') else {}} for t in result.tools]
    return []

async def _execute_mcp_tool(mcp_srv: dict, tool_name: str, arguments: dict) -> str:
    """Connect to MCP server and execute a specific tool."""
    from mcp import ClientSession
    if mcp_srv["server_type"] in ("stdio", "module"):
        from mcp.client.stdio import stdio_client, StdioServerParameters
        cmd = mcp_srv.get("command") or "python3"
        args_raw = mcp_srv.get("args")
        args = json.loads(args_raw) if isinstance(args_raw, str) else (args_raw or [])
        env_raw = mcp_srv.get("env_vars")
        env = json.loads(env_raw) if isinstance(env_raw, str) else (env_raw or None)
        params = StdioServerParameters(command=cmd, args=args, env=env)
        async with stdio_client(params) as streams:
            async with ClientSession(*streams) as session:
                await session.initialize()
                result = await session.call_tool(tool_name, arguments)
                parts = []
                for c in result.content:
                    if hasattr(c, 'text'):
                        parts.append(c.text)
                    elif hasattr(c, 'data'):
                        parts.append(str(c.data))
                return "\n".join(parts) if parts else "Tool executed successfully (no output)"
    elif mcp_srv["server_type"] == "sse":
        from mcp.client.sse import sse_client
        async with sse_client(mcp_srv["url"]) as streams:
            async with ClientSession(*streams) as session:
                await session.initialize()
                result = await session.call_tool(tool_name, arguments)
                parts = []
                for c in result.content:
                    if hasattr(c, 'text'):
                        parts.append(c.text)
                    elif hasattr(c, 'data'):
                        parts.append(str(c.data))
                return "\n".join(parts) if parts else "Tool executed successfully (no output)"
    raise ValueError(f"Unsupported MCP server type: {mcp_srv['server_type']}")

def _get_active_mcp_tools(user_id: str) -> list[dict]:
    """Get all tools from active MCP servers for a user, with server reference."""
    with db() as c:
        rows = c.execute(
            "SELECT * FROM mcp_servers WHERE is_active=1 AND (user_id=? OR is_global=1)",
            (user_id,)
        ).fetchall()
    result = []
    for r in rows:
        srv = dict(r)
        tools_raw = srv.get("tools")
        if not tools_raw:
            continue
        try:
            tools = json.loads(tools_raw) if isinstance(tools_raw, str) else tools_raw
        except Exception as e:
            log.warning(f"Failed to parse tools JSON for MCP server '{srv.get('name', '?')}': {e}")
            continue
        for t in tools:
            if isinstance(t, dict) and t.get("name"):
                result.append({"server": srv, "tool": t})
    return result

# ── ADB Vector DB Config ─────────────────────────────────────────────────────