"""
Agent management module for creating and managing pydantic-ai agents.

This module handles the lifecycle of AI agents including:
- Creating pydantic-ai Agent instances with specified configurations
- Managing MCP server toolsets
- Managing global agent instances
"""

import asyncio
import os
from typing import Optional, Dict, Any, List, Set

from pydantic_ai import Agent, Tool as PydanticTool
from pydantic_ai.mcp import MCPServerStdio, MCPServerStreamableHTTP
from pydantic_ai.tools import DeferredToolRequests, RunContext
from pydantic_ai.toolsets import FunctionToolset

from suzent.core.agent_deps import AgentDeps
from suzent.core.model_factory import create_pydantic_ai_model
from suzent.core.providers import get_enabled_models_from_db

from suzent.config import CONFIG
from suzent.logger import get_logger
from suzent.prompts import (
    STATIC_INSTRUCTIONS,
    register_dynamic_instructions,
)
from suzent.skills import get_skill_manager
from suzent.tools.tool_search_tool import tool_search

# Import memory lifecycle functions (for backward compatibility re-exports)
from suzent.memory.lifecycle import (
    get_memory_manager,
)

# Suppress LiteLLM's verbose logging
os.environ["LITELLM_LOG"] = "ERROR"

logger = get_logger(__name__)


# --- Agent State ---
agent_instance: Optional[Agent] = None
agent_config: Optional[dict] = None
agent_lock = asyncio.Lock()

# --- Deferred (agent-activated) tool state ---
_unlocked_tools_by_chat: Dict[str, Set[str]] = {}


def get_unlocked_tools(chat_id: str) -> Set[str]:
    return _unlocked_tools_by_chat.get(chat_id, set())


def unlock_tool(chat_id: str, tool_name: str) -> None:
    _unlocked_tools_by_chat.setdefault(chat_id, set()).add(tool_name)


def remove_unlocked_tool(chat_id: str, tool_name: str) -> None:
    _unlocked_tools_by_chat.get(chat_id, set()).discard(tool_name)


def clear_unlocked_tools(chat_id: str) -> None:
    _unlocked_tools_by_chat.pop(chat_id, None)


def _build_mcp_servers(config: Dict[str, Any]) -> List:
    """
    Build MCP server toolset instances from the enabled servers in config.

    Returns a list of MCPServerStreamableHTTP / MCPServerStdio instances
    that can be passed as ``toolsets`` to a pydantic-ai Agent.
    """
    from suzent.core import mcp_store as _mcp_store

    # JSON file is the source of truth; config dict can override for special callers
    _defaults = _mcp_store.as_agent_config()
    mcp_enabled = config.get("mcp_enabled") or _defaults["mcp_enabled"]
    mcp_headers = config.get("mcp_headers") or _defaults["mcp_headers"]
    mcp_stdio_params = config.get("mcp_stdio_params") or _defaults["mcp_stdio_params"]

    raw_mcp_urls = config.get("mcp_urls") or _defaults["mcp_urls"]

    # Parse mcp_urls to handle both simple and nested formats (legacy config.yaml support)
    mcp_urls = {}

    if isinstance(raw_mcp_urls, list):
        for i, url in enumerate(raw_mcp_urls):
            mcp_urls[f"mcp-url-{i}"] = url
    elif isinstance(raw_mcp_urls, dict):
        for name, value in raw_mcp_urls.items():
            if isinstance(value, str):
                mcp_urls[name] = value
            elif isinstance(value, dict):
                mcp_urls[name] = value.get("url", "")
                if value.get("headers") and name not in mcp_headers:
                    mcp_headers[name] = value["headers"]

    servers = []

    # Build URL servers
    for name, url in mcp_urls.items():
        if mcp_enabled.get(name, False) and url:
            headers = mcp_headers.get(name)
            server = MCPServerStreamableHTTP(
                url,
                headers=headers,
                tool_prefix=name,
            )
            servers.append(server)

    # Build stdio servers
    if mcp_stdio_params:
        for name, params in mcp_stdio_params.items():
            if mcp_enabled.get(name, False):
                server = MCPServerStdio(
                    params["command"],
                    args=params.get("args", []),
                    env=params.get("env"),
                    tool_prefix=name,
                )
                servers.append(server)

    return servers


async def probe_mcp_server(
    entry: Dict[str, Any], timeout: float | None = None
) -> Dict[str, Any]:
    """Attempt to connect to a single MCP server and list its tools.

    `entry` is a row from mcp_store (type/url/headers/command/args/env).
    On success returns {"ok": True, "count": int, "tools": [{"name", "description"}]};
    on failure {"ok": False, "error": str}.

    stdio servers get a much longer default timeout: the first `uv tool run` /
    `npx` invocation may download and install the package before the server
    starts speaking, which can take well over a minute on a cold cache.
    """
    import asyncio

    is_stdio = entry.get("type") == "stdio"
    if timeout is None:
        timeout = 90.0 if is_stdio else 15.0

    has_url = entry.get("type") == "url" and entry.get("url")
    has_cmd = is_stdio and entry.get("command")
    if not has_url and not has_cmd:
        return {"ok": False, "error": "Server has no url or command configured"}

    def _build():
        if has_url:
            return MCPServerStreamableHTTP(entry["url"], headers=entry.get("headers"))
        return MCPServerStdio(
            entry["command"], args=entry.get("args") or [], env=entry.get("env")
        )

    async def _connect() -> list[dict[str, str]]:
        # Re-create the server per attempt; a half-started transport cannot be reused.
        try:
            srv = _build()
        except Exception as e:  # noqa: BLE001 - bad config surfaces to caller
            raise RuntimeError(f"Invalid server config: {e}") from e
        async with srv:
            tools = await srv.list_tools()
            return [
                {
                    "name": getattr(t, "name", ""),
                    "description": (getattr(t, "description", "") or "").strip(),
                }
                for t in tools
            ]

    # stdio startup over `uv tool run` / `npx` can transiently break the pipe before
    # the handshake settles; retry those a couple times before giving up.
    attempts = 3 if is_stdio else 1
    last_error: str | None = None
    for attempt in range(attempts):
        try:
            tools = await asyncio.wait_for(_connect(), timeout=timeout)
            return {"ok": True, "count": len(tools), "tools": tools}
        except asyncio.TimeoutError:
            hint = (
                " — the command may still be installing its package on first run; "
                "try again once it is cached"
                if is_stdio
                else ""
            )
            return {"ok": False, "error": f"Timed out after {timeout:.0f}s{hint}"}
        except Exception as e:  # noqa: BLE001 - surface connection failure
            last_error = _flatten_error(e)
            if attempt < attempts - 1 and _is_transient_stdio_error(last_error):
                await asyncio.sleep(1.0 * (attempt + 1))
                continue
            break

    return {"ok": False, "error": last_error or "Connection failed"}


def _flatten_error(exc: BaseException) -> str:
    """Produce a readable message, unwrapping anyio/TaskGroup ExceptionGroups."""
    # ExceptionGroup (Python 3.11+) hides the real cause behind a generic message.
    inner = getattr(exc, "exceptions", None)
    if inner:
        return "; ".join(_flatten_error(e) for e in inner)
    name = exc.__class__.__name__
    msg = str(exc).strip()
    if name == "BrokenResourceError" and not msg:
        # The child process closed the pipe before the handshake completed —
        # usually it exited early (bad args, crash, or still warming up).
        return "Server process closed the connection before responding (it may have exited early or is still starting)"
    return msg or name


def _is_transient_stdio_error(message: str) -> bool:
    """True for stdio startup races worth retrying (broken/closed pipe, EOF)."""
    lowered = message.lower()
    return any(
        marker in lowered
        for marker in (
            "brokenresource",
            "closed the connection",
            "closedresource",
            "eof",
        )
    )


def create_agent(
    config: Dict[str, Any], memory_context: Optional[str] = None
) -> Agent[AgentDeps, str]:
    """
    Creates a pydantic-ai Agent based on the provided configuration.

    Args:
        config: Configuration dictionary containing:
            - model: Model identifier (e.g., "gemini/gemini-2.5-pro")
            - tools: List of tool names to enable
            - memory_enabled: Whether to inject memory context (default: False).
              Does NOT control the memory_search tool, which is equipped via `tools`.
            - mcp_urls: Optional MCP server URLs
            - instructions: Optional custom instructions

    Returns:
        Configured pydantic-ai Agent instance.
    """
    # --- Validate model ---
    enabled_models = get_enabled_models_from_db()

    if not enabled_models:
        raise ValueError(
            "No LLM models are enabled. Please configure a provider in Settings."
        )

    model_id = config.get("model")

    if not model_id or model_id not in enabled_models:
        fallback = enabled_models[0]
        if model_id:
            logger.warning(
                f"Requested model '{model_id}' is not enabled. Falling back to '{fallback}'."
            )
        model_id = fallback

    model = create_pydantic_ai_model(model_id)

    # --- Build tool list ---
    tool_names = (config.get("tools") or CONFIG.default_tools).copy()

    from suzent.tools.registry import get_tool_function, get_tool_session_guidance

    tool_functions = []
    enabled_tool_names = set(tool_names)
    # SkillTool / SocialMessageTool are equipped by their own auto-equip logic below,
    # so skip them in the normal loop. MemorySearchTool is NOT auto-equipped: it is a
    # regular sidebar-toggleable equipment tool, equipped when present in `tools`. The
    # global memory toggle (`memory_enabled`) only governs memory *context injection*,
    # not whether the search tool is available.
    _auto_equipped = {
        "SkillTool",
        "SocialMessageTool",
    }

    for tool_name in tool_names:
        if tool_name in _auto_equipped:
            continue
        fn = get_tool_function(tool_name)
        if fn:
            tool_functions.append(fn)
        else:
            logger.warning(f"Tool function not found for: {tool_name}")

    # Auto-equip SkillTool if any skills are enabled
    skill_manager = get_skill_manager()
    if skill_manager.enabled_skills:
        fn = get_tool_function("SkillTool")
        if fn and fn not in tool_functions:
            tool_functions.append(fn)
            enabled_tool_names.add("SkillTool")
            logger.info(
                f"SkillTool equipped ({len(skill_manager.enabled_skills)} skills enabled)"
            )

    # Auto-equip SocialMessageTool
    social_ctx = config.get("social_context")
    if social_ctx or "SocialMessageTool" in tool_names:
        fn = get_tool_function("SocialMessageTool")
        if fn and fn not in tool_functions:
            tool_functions.append(fn)
            enabled_tool_names.add("SocialMessageTool")

    # --- Build MCP servers ---
    mcp_servers = _build_mcp_servers(config)

    # --- Build instructions ---
    base_instructions = config.get("instructions", CONFIG.instructions)
    session_guidance_items = get_tool_session_guidance(sorted(enabled_tool_names))
    static_instructions = config.get("static_instructions", STATIC_INSTRUCTIONS)

    # --- Create pydantic-ai Agent ---
    agent = Agent(
        model,
        deps_type=AgentDeps,
        tools=tool_functions,
        toolsets=mcp_servers if mcp_servers else [],
        instructions=static_instructions,
        output_type=[str, DeferredToolRequests],
        output_retries=3,
    )

    register_dynamic_instructions(
        agent,
        base_instructions=base_instructions,
        memory_context=memory_context,
        session_guidance_items=session_guidance_items,
        enabled_model_ids=enabled_models,
        current_model_id=model_id,
    )

    # --- Deferred toolset: injects agent-activated tools each LLM step ---
    _base_tool_names = set(enabled_tool_names)

    @agent.toolset
    async def _deferred_toolset(ctx: RunContext[AgentDeps]) -> FunctionToolset:
        ts = FunctionToolset()
        for name in get_unlocked_tools(ctx.deps.chat_id):
            if name in _base_tool_names:
                continue
            fn = get_tool_function(name)
            if fn is not None:
                if isinstance(fn, PydanticTool):
                    ts.add_tool(fn)
                else:
                    ts.add_function(fn)
        return ts

    agent.tool(tool_search)

    # Store metadata for later introspection
    agent._tool_names = [tn for tn in tool_names]  # type: ignore[attr-defined]
    agent._model_id = model_id  # type: ignore[attr-defined]

    return agent


def build_agent_config(
    base_config: Optional[Dict[str, Any]] = None, require_social_tool: bool = False
) -> Dict[str, Any]:
    """
    Builds the effective configuration dictionary for a ChatProcessor turn,
    merging base configs with user preferences from the database.

    Args:
        base_config: Initial configuration overrides (e.g., from request).
        require_social_tool: If True, ensures SocialMessageTool is equipped.

    Returns:
        A dictionary containing the merged configuration.
    """
    from suzent.database import get_database

    config = base_config.copy() if base_config else {}

    try:
        db = get_database()
        if prefs := db.get_user_preferences():
            if not config.get("model") and prefs.model:
                config["model"] = prefs.model
            if not config.get("agent") and prefs.agent:
                config["agent"] = prefs.agent
            if "tools" not in config and prefs.tools is not None:
                config["tools"] = prefs.tools
    except Exception as e:
        logger.warning(f"Failed to load user preferences: {e}")

    # Ensure tools list exists and is populated
    tools = config.get("tools")
    if tools is None:
        tools = CONFIG.default_tools.copy()
    elif isinstance(tools, list):
        tools = tools.copy()

    if require_social_tool and "SocialMessageTool" not in tools:
        tools.append("SocialMessageTool")

    config["tools"] = tools

    return config


async def get_or_create_agent(config: Dict[str, Any], reset: bool = False) -> Agent:
    """
    Get the current agent instance or create a new one if needed.

    Args:
        config: Agent configuration dictionary.
        reset: If True, force creation of a new agent instance.

    Returns:
        pydantic-ai Agent instance ready for use.
    """
    global agent_instance, agent_config

    _TRANSIENT_KEYS = {"_runtime", "_chat_id", "_user_id"}

    def _stable_config(cfg: dict) -> dict:
        return {k: v for k, v in cfg.items() if k not in _TRANSIENT_KEYS}

    async with agent_lock:
        config_changed = _stable_config(config) != (
            _stable_config(agent_config) if agent_config else None
        )
        if config_changed and agent_config is not None:
            logger.info("Config changed - creating new agent")

        if agent_instance is None or config_changed or reset:
            # Fetch memory context if memory system is enabled
            memory_context = None
            memory_enabled = config.get("memory_enabled", False)
            mem_manager = get_memory_manager()
            if mem_manager and memory_enabled:
                chat_id = config.get("_chat_id")
                user_id = config.get("_user_id", "default-user")
                sandbox_enabled = config.get("sandbox_enabled", CONFIG.sandbox_enabled)
                try:
                    from suzent.tools.filesystem.path_resolver import PathResolver
                    from suzent.config import get_effective_volumes

                    _path_resolver = PathResolver(
                        chat_id=chat_id or "default",
                        sandbox_enabled=sandbox_enabled,
                        custom_volumes=get_effective_volumes(
                            config.get("sandbox_volumes")
                        ),
                    )
                    memory_context = await mem_manager.format_core_memory_for_context(
                        chat_id=chat_id,
                        user_id=user_id,
                        sandbox_enabled=sandbox_enabled,
                        path_resolver=_path_resolver,
                    )
                    if memory_context:
                        logger.debug(f"Fetched core memory context for user={user_id}")
                except Exception as e:
                    logger.error(f"Error fetching memory context: {e}")
                    memory_context = None

            agent_instance = create_agent(config, memory_context=memory_context)
            agent_config = config

        return agent_instance
