"""Main bridge facade: connect MCP servers, route tools to LLMs, run the loop."""

from __future__ import annotations

import asyncio
import logging
import json
import time
from typing import Any, Callable

from mcpbridge.config import BridgeConfig
from mcpbridge.exceptions import BridgeInUseError, ConfigValidationError, MCPProtocolError
from mcpbridge.llm_router import LLMRouter
from mcpbridge.loop import LoopResult, run_loop
from mcpbridge.prompt_builder import PromptBuilder
from mcpbridge.registry import MCPTool, DynamicToolRegistry
from mcpbridge.transport.auto import detect_transport
from mcpbridge.transport.base import BaseTransport

logger = logging.getLogger("mcpbridge")


class MCPBridge:
    """Universal MCP <-> LLM bridge.

    Instantiate with a config dict, then `await connect()` and finally `await run()`.
    """

    def __init__(self, config: dict[str, Any]) -> None:
        self._bridge_config = BridgeConfig.from_dict(config)

        self._adapter = None
        self._router = LLMRouter(self._bridge_config)

        self._registry: DynamicToolRegistry | None = None
        self._transport_map: dict[str, BaseTransport] = {}

        self._history: list[dict[str, Any]] = []
        self._persistent_context: dict[str, Any] = dict(self._bridge_config.prompt.context_vars or {})

        self._prompt_builder: PromptBuilder | None = None

        self._connected = False
        self._connect_lock = asyncio.Lock()

        self._run_lock = asyncio.Lock()

    async def connect(self) -> "MCPBridge":
        """Connect transport(s), discover MCP tools/resources/prompts, and prepare prompt context."""
        async with self._connect_lock:
            if self._connected:
                return self

            self._adapter = self._router.resolve()
            assert self._adapter is not None

            # Discover and merge MCP servers.
            namespace_strategy = self._bridge_config.mcp.namespace_strategy
            merged = DynamicToolRegistry(server_name="merged", namespace_strategy=namespace_strategy)

            servers = self._bridge_config.mcp.servers
            if servers:
                for srv in servers:
                    transport = detect_transport(
                        {
                            "command": srv.command,
                            "args": srv.args,
                            "env": srv.env,
                            "url": srv.url,
                            "transport": srv.transport,
                            "headers": {**self._bridge_config.mcp.headers, **srv.headers},
                            "timeout": self._bridge_config.mcp.timeout,
                        }
                    )
                    await transport.connect()
                    reg = DynamicToolRegistry(server_name=srv.name, namespace_strategy=namespace_strategy)
                    await reg.discover(transport)
                    merged.merge(reg)
                    self._transport_map[srv.name] = transport
            else:
                # Single server from top-level config.
                srv_name = "default"
                transport = detect_transport(
                    {
                        "command": self._bridge_config.mcp.command,
                        "args": self._bridge_config.mcp.args,
                        "env": self._bridge_config.mcp.env,
                        "url": self._bridge_config.mcp.url,
                        "transport": self._bridge_config.mcp.transport,
                        "headers": self._bridge_config.mcp.headers,
                        "timeout": self._bridge_config.mcp.timeout,
                    }
                )
                await transport.connect()
                reg = DynamicToolRegistry(server_name=srv_name, namespace_strategy=namespace_strategy)
                await reg.discover(transport)
                merged.merge(reg)
                self._transport_map[srv_name] = transport

            self._registry = merged

            # Prepare prompt builder.
            pb = PromptBuilder(self._bridge_config.prompt)
            pb.set_tools_context(list(self._registry.tools.values()))
            pb.set_system_in_messages(not bool(self._adapter.supports_system_prompt))
            self._prompt_builder = pb

            self._connected = True
            return self

    async def close(self) -> None:
        """Close all transport connections."""
        self._connected = False
        for _, transport in list(self._transport_map.items()):
            try:
                await transport.close()
            except Exception:
                pass
        self._transport_map.clear()
        self._registry = None

    async def __aenter__(self) -> "MCPBridge":
        return await self.connect()

    async def __aexit__(self, *_: Any) -> None:
        await self.close()

    def set_context(self, **kwargs: Any) -> None:
        """Set persistent context vars for `{placeholder}` interpolation."""
        self._persistent_context.update(kwargs)

    def clear_history(self) -> None:
        """Clear conversation history."""
        self._history = []

    def get_history(self) -> list[dict[str, Any]]:
        """Return current conversation history."""
        return list(self._history)

    def set_history(self, history: list[dict[str, Any]]) -> None:
        """Restore saved history."""
        if not isinstance(history, list):
            raise ConfigValidationError("history must be a list")
        for m in history:
            if not isinstance(m, dict):
                raise ConfigValidationError("Each history message must be a dict")
            if "role" not in m or not isinstance(m["role"], str):
                raise ConfigValidationError("Each history message must have string `role`")
            # content can be string or list depending on provider.
            if "content" in m and not isinstance(m["content"], (str, list)):
                raise ConfigValidationError("History message `content` must be string or list")
        self._history = history

    async def run(
        self,
        user_query: str,
        context: dict[str, Any] | None = None,
        stream: bool = False,
    ) -> LoopResult:
        """Run the full agentic loop for a single user query."""
        # `stream` is accepted for forward-compatibility, but provider-level
        # token streaming isn't wired up in this version. Until adapters and
        # the loop support incremental deltas, we treat it as a no-op.
        if stream or bool(self._bridge_config.llm.stream):
            logger.warning(
                "`stream=True` requested, but token streaming is not implemented; "
                "falling back to non-streaming mode."
            )
        if not self._connected:
            await self.connect()
        if self._registry is None or self._adapter is None:
            raise MCPProtocolError("Bridge not connected (missing registry/adapter).")

        if self._run_lock.locked():
            raise BridgeInUseError("This MCPBridge instance is already running.")

        async with self._run_lock:
            # Create a per-run prompt builder instance so runtime context doesn't
            # permanently mutate the bridge.
            pb = PromptBuilder(self._bridge_config.prompt)
            pb.set_tools_context(list(self._registry.tools.values()))
            pb.set_system_in_messages(not bool(self._adapter.supports_system_prompt))

            pb.set_context(**self._persistent_context)
            if context:
                pb.set_context(**context)

            loop_res = await run_loop(
                adapter=self._adapter,
                prompt_builder=pb,
                registry=self._registry,
                transport=self._transport_map if len(self._transport_map) > 1 else next(iter(self._transport_map.values())),
                user_query=user_query,
                history=self._history,
                loop_config=self._bridge_config.loop,
            )

            # History trimming (best-effort).
            if self._bridge_config.session.persist_history:
                self._history = self._trim_history(
                    self._history,
                    max_tokens=int(self._bridge_config.session.max_history_tokens),
                    strategy=str(self._bridge_config.session.history_trim_strategy),
                )
            else:
                self._history = []
            return loop_res

    def _trim_history(self, history: list[dict[str, Any]], max_tokens: int, strategy: str) -> list[dict[str, Any]]:
        """Best-effort history trimming that preserves tool call/result integrity."""
        if max_tokens <= 0:
            return history

        # Very lightweight token estimate: chars/4.
        def estimate(msgs: list[dict[str, Any]]) -> int:
            from mcpbridge.utils.token_counter import _rough_tokens  # type: ignore[attr-defined]

            total = 0
            for m in msgs:
                content = m.get("content")
                if isinstance(content, str):
                    total += _rough_tokens(content)
                elif isinstance(content, list):
                    total += _rough_tokens(str(content))
            return total

        if estimate(history) <= max_tokens:
            return history

        # `oldest_first`: keep system (if present) and the latest user message,
        # plus any messages after it.
        if strategy == "summarize":
            # Naive summarization: collapse removed prefix into a truncated string.
            keep_last_user_idx = max((i for i, m in enumerate(history) if m.get("role") == "user"), default=len(history) - 1)
            system_part = [history[0]] if history and history[0].get("role") == "system" else []
            prefix = history[1:keep_last_user_idx]
            summary_text = ""
            for pm in prefix[:5]:
                if isinstance(pm.get("content"), str):
                    summary_text += pm["content"][:200] + " "
            summary_text = summary_text.strip()[:500]
            kept = system_part + ([{"role": "assistant", "content": f"Summary: {summary_text}"}] if summary_text else []) + history[keep_last_user_idx:]
        else:
            keep_last_user_idx = max((i for i, m in enumerate(history) if m.get("role") == "user"), default=len(history) - 1)
            system_part = [history[0]] if history and history[0].get("role") == "system" else []
            kept = system_part + history[keep_last_user_idx:]

        # Ensure we don't exceed max tokens by trimming from the front while preserving the last tool block.
        while kept and estimate(kept) > max_tokens and len(kept) > 2:
            # Preserve system message at index 0 if present.
            if kept and kept[0].get("role") == "system":
                kept.pop(1)
            else:
                kept.pop(0)
        return kept

    def list_tools(self) -> list[MCPTool]:
        """Return all discovered MCP tools."""
        if not self._registry:
            return []
        return list(self._registry.tools.values())

    async def list_resources(self) -> list[dict[str, Any]]:
        """Return MCP resources list (cached after discovery)."""
        if not self._registry:
            return []
        return list(self._registry.resources.values())

    async def read_resource(self, uri: str) -> str:
        """Read an MCP resource by URI."""
        if not self._registry:
            raise MCPProtocolError("Bridge not connected")
        # Use the first transport as a fallback. Resource reads are typically
        # server-specific but the registry cache doesn't track ownership.
        transport = next(iter(self._transport_map.values()))
        req = {"jsonrpc": "2.0", "id": asyncio.get_running_loop().time(), "method": "resources/read", "params": {"uri": uri}}
        resp = await transport.send(req)
        if not isinstance(resp, dict) or resp.get("error"):
            raise MCPProtocolError(f"MCP resources/read failed: {resp.get('error') if isinstance(resp, dict) else resp}")
        result = (resp.get("result") or {}) if isinstance(resp, dict) else {}
        contents = result.get("contents") or result.get("content") or []
        if isinstance(contents, list):
            texts = []
            for b in contents:
                if isinstance(b, dict) and isinstance(b.get("text"), str):
                    texts.append(b["text"])
            if texts:
                return "\n".join(texts)
        # Fallback: return JSON.
        return json.dumps(result, ensure_ascii=False)

    async def list_prompts(self) -> list[dict[str, Any]]:
        """Return MCP prompt templates list (cached after discovery)."""
        if not self._registry:
            return []
        return list(self._registry.prompts.values())

    async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -> str:
        """Render a named MCP prompt template by calling `prompts/get`."""
        if not self._registry:
            raise MCPProtocolError("Bridge not connected")

        args = arguments or {}
        # If namespacing is enabled, the prompt name may include `server__name`.
        server_name = "default"
        if "__" in name:
            maybe_server, _rest = name.split("__", 1)
            if maybe_server in self._transport_map:
                server_name = maybe_server

        transport = self._transport_map.get(server_name) or next(iter(self._transport_map.values()))

        prompt_key = name
        prompt_def = self._registry.prompts.get(prompt_key)
        prompt_actual_name = prompt_def.get("name") if isinstance(prompt_def, dict) else name.split("__", 1)[-1]
        req = {
            "jsonrpc": "2.0",
            "id": asyncio.get_running_loop().time(),
            "method": "prompts/get",
            "params": {"name": prompt_actual_name, "arguments": args},
        }
        resp = await transport.send(req)
        if not isinstance(resp, dict) or resp.get("error"):
            raise MCPProtocolError(f"MCP prompts/get failed: {resp.get('error') if isinstance(resp, dict) else resp}")
        result = resp.get("result") or {}
        # The response typically contains `messages`.
        messages = result.get("messages") or []
        rendered_parts: list[str] = []
        if isinstance(messages, list):
            for m in messages:
                if not isinstance(m, dict):
                    continue
                content = m.get("content")
                if isinstance(content, str):
                    rendered_parts.append(content)
                elif isinstance(content, list):
                    for b in content:
                        if isinstance(b, dict) and isinstance(b.get("text"), str):
                            rendered_parts.append(b["text"])
        if rendered_parts:
            return "\n".join(rendered_parts).strip()
        return json.dumps(result, ensure_ascii=False)

    async def add_server(self, name: str, config: dict[str, Any]) -> None:
        """Dynamically add another MCP server at runtime."""
        if not self._connected:
            await self.connect()
        if self._registry is None:
            raise MCPProtocolError("Bridge not connected")

        transport = detect_transport(
            {
                **config,
                "timeout": self._bridge_config.mcp.timeout,
                "headers": {**self._bridge_config.mcp.headers, **(config.get("headers") or {})},
            }
        )
        await transport.connect()
        reg = DynamicToolRegistry(server_name=name, namespace_strategy=self._bridge_config.mcp.namespace_strategy)
        await reg.discover(transport)
        self._registry.merge(reg)
        self._transport_map[name] = transport
        # Update prompt tools context.
        if self._prompt_builder:
            self._prompt_builder.set_tools_context(list(self._registry.tools.values()))

    async def remove_server(self, name: str) -> None:
        """Disconnect and remove a server."""
        if name in self._transport_map:
            transport = self._transport_map.pop(name)
            try:
                await transport.close()
            except Exception:
                pass

        if self._registry:
            self._registry.tools = {k: v for k, v in self._registry.tools.items() if v.server_name != name}
            # Remove namespaced prompts with prefix strategy.
            self._registry.prompts = {k: v for k, v in self._registry.prompts.items() if not k.startswith(f"{name}__")}
            self._registry._conversion_cache = {}

        if self._prompt_builder and self._registry:
            self._prompt_builder.set_tools_context(list(self._registry.tools.values()))

