from __future__ import annotations

import json
import threading
from typing import Any

from .config import iter_enabled_servers_from_config
from .session_manager import SESSION_MANAGER, SessionManager

DISCOVERY_LOCK = threading.Lock()


class _SchemaShim:
    def __init__(self, schema: dict[str, Any]):
        self._schema = schema or {"type": "object", "properties": {}}

    def model_json_schema(self) -> dict[str, Any]:
        return dict(self._schema)


class McpToolWrapper:
    def __init__(self, name: str, description: str, input_schema: dict[str, Any], server_id: str, tool_name: str):
        self.name = name
        self.description = description or ""
        self.args_schema = _SchemaShim(input_schema or {"type": "object", "properties": {}})
        self.metadata = {"realtime": False}
        self._server_id = server_id
        self._tool_name = tool_name

    def invoke(self, args: dict[str, Any] | str | None) -> str:
        if isinstance(args, str):
            try:
                parsed = json.loads(args)
            except Exception:
                parsed = {}
        elif isinstance(args, dict):
            parsed = args
        else:
            parsed = {}
        ok, result = SESSION_MANAGER.call_tool(self._server_id, self._tool_name, parsed)
        return result if ok else (result if result.startswith("Tool error:") else f"Tool error: {result}")


def build_openai_schema(name: str, description: str, input_schema: dict[str, Any]) -> dict[str, Any]:
    schema = dict(input_schema or {"type": "object", "properties": {}})
    schema.pop("title", None)
    schema.setdefault("type", "object")
    schema.setdefault("properties", {})
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": description or "",
            "parameters": schema,
        },
    }


def discover_tools() -> tuple[list[McpToolWrapper], list[dict[str, Any]], list[dict[str, str]]]:
    return _discover_tools_with_manager(SESSION_MANAGER)


def _discover_tools_with_manager(
    manager: SessionManager,
    *,
    servers: list[tuple[str, dict[str, Any]]] | None = None,
) -> tuple[list[McpToolWrapper], list[dict[str, Any]], list[dict[str, str]]]:
    wrappers: list[McpToolWrapper] = []
    schemas: list[dict[str, Any]] = []
    errors: list[dict[str, str]]
    with DISCOVERY_LOCK:
        records, errors = manager.discover_all_sync(servers=servers)

    for record in records:
        server_id = str(record.get("server_id") or "")
        tool_name = str(record.get("tool_name") or "")
        if not server_id or not tool_name:
            continue
        mcp_name = f"mcp__{server_id}__{tool_name}"
        schema = record.get("input_schema")
        desc = str(record.get("description") or "")
        if not isinstance(schema, dict):
            schema = {}
        wrappers.append(McpToolWrapper(mcp_name, desc, schema, server_id, tool_name))
        schemas.append(build_openai_schema(mcp_name, desc, schema))
    return wrappers, schemas, errors


def discover_tools_from_config(cfg: dict[str, Any]) -> tuple[list[McpToolWrapper], list[dict[str, Any]], list[dict[str, str]]]:
    # Probe discovery must not mutate global SESSION_MANAGER state used by live chat.
    probe_manager = SessionManager(register_atexit=False)
    try:
        servers = iter_enabled_servers_from_config(cfg)
        return _discover_tools_with_manager(probe_manager, servers=servers)
    finally:
        probe_manager.shutdown()
