"""MCPRegistry — SQLite-backed custom MCP server registry."""
from __future__ import annotations
import json
import sqlite3
from pathlib import Path

DB_PATH = Path(__file__).parent / "mcp_servers.db"


class MCPRegistry:
    def __init__(self, db_path=None):
        self.db_path = str(db_path or DB_PATH)
        self._init_db()

    def _init_db(self):
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS mcp_servers (
                    name TEXT PRIMARY KEY,
                    url TEXT NOT NULL,
                    module_affinity TEXT NOT NULL DEFAULT '["code","perf"]',
                    enabled INTEGER NOT NULL DEFAULT 1,
                    created_at TEXT NOT NULL DEFAULT (datetime('now'))
                )
            """)

    def list_servers(self) -> list[dict]:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute("SELECT * FROM mcp_servers WHERE enabled=1").fetchall()
            return [
                {
                    "name": r["name"],
                    "url": r["url"],
                    "module_affinity": json.loads(r["module_affinity"]),
                    "enabled": bool(r["enabled"]),
                }
                for r in rows
            ]

    def add_server(self, name: str, url: str, module_affinity: list[str] = None) -> dict:
        affinity = module_affinity or ["code", "perf"]
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "INSERT OR REPLACE INTO mcp_servers (name, url, module_affinity) VALUES (?, ?, ?)",
                (name, url, json.dumps(affinity)),
            )
        return {"name": name, "url": url, "module_affinity": affinity}

    def remove_server(self, name: str) -> bool:
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("DELETE FROM mcp_servers WHERE name=?", (name,))
            return cursor.rowcount > 0

    def get_server(self, name: str) -> dict | None:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute(
                "SELECT * FROM mcp_servers WHERE name=?", (name,)
            ).fetchone()
            if row:
                return {
                    "name": row["name"],
                    "url": row["url"],
                    "module_affinity": json.loads(row["module_affinity"]),
                    "enabled": bool(row["enabled"]),
                }
        return None

    def test_connection(self, name: str) -> dict:
        """Test if an MCP server is reachable."""
        server = self.get_server(name)
        if not server:
            return {"status": "error", "message": f"Server '{name}' not found"}
        from agent.mcp_client import MCPClient, MCPConnectionError

        try:
            client = MCPClient(base_url=server["url"], timeout=5)
            tools = client.discover()
            return {"status": "ok", "tools_count": len(tools)}
        except MCPConnectionError as e:
            return {"status": "error", "message": str(e)}