"""Resolve PersonaSpec via larva MCP server.

Provides a fallback persona resolution path: when a persona_ref is not a local
file path, attempt to resolve it as a spec_id via the ``larva_resolve`` MCP tool.

This works identically whether larva is connected directly or proxied through tela.
The resolver checks for a ``larva_resolve`` tool in the available MCP toolsets.
"""

# @invar:allow shell_result: MCP tool call bridge returns Result for persona resolution

from __future__ import annotations

import json
import logging
from typing import Any

from returns.result import Failure, Result, Success

from anima.core.schemas import AnimaError, InvalidInputError, PersonaSpec

logger = logging.getLogger(__name__)

LARVA_RESOLVE_TOOL = "larva_resolve"


async def resolve_persona_via_larva(
    spec_id: str,
    mcp_servers: list[Any] | None = None,
) -> Result[PersonaSpec, AnimaError]:
    """Resolve a persona spec_id via larva MCP server.

    Scans available MCP toolsets for ``larva_resolve``, calls it with the
    spec_id, and parses the result into a PersonaSpec.

    Works with both direct larva connections and tela-proxied access.

    Args:
        spec_id: Persona identifier to resolve (e.g., "code-reviewer").
        mcp_servers: List of MCP toolset objects from the agent runtime.

    Returns:
        Success(PersonaSpec) if resolved, Failure if larva not available
        or resolution fails.

    Examples:
        >>> # No MCP servers → Failure
        >>> import asyncio
        >>> r = asyncio.run(resolve_persona_via_larva("test", None))
        >>> from returns.result import Failure
        >>> isinstance(r, Failure)
        True
    """
    if not mcp_servers:
        return Failure(
            InvalidInputError(
                f"Cannot resolve spec_id '{spec_id}': no MCP servers available",
                {"spec_id": spec_id},
            )
        )

    larva_toolset = await _find_larva_toolset(mcp_servers)
    if larva_toolset is None:
        return Failure(
            InvalidInputError(
                f"Cannot resolve spec_id '{spec_id}': no larva MCP server found "
                f"(looking for '{LARVA_RESOLVE_TOOL}' tool)",
                {"spec_id": spec_id},
            )
        )

    try:
        result = await _call_larva_resolve(larva_toolset, spec_id)
    except Exception as exc:
        logger.warning(
            f"larva_resolve('{spec_id}') failed: {exc}",
            exc_info=True,
        )
        return Failure(
            InvalidInputError(
                f"larva_resolve('{spec_id}') call failed: {exc}",
                {"spec_id": spec_id, "error": str(exc)},
            )
        )

    if isinstance(result, dict) and "error" in result:
        error_info = result["error"]
        message = (
            error_info.get("message", str(error_info))
            if isinstance(error_info, dict)
            else str(error_info)
        )
        return Failure(
            InvalidInputError(
                f"larva_resolve('{spec_id}') returned error: {message}",
                {"spec_id": spec_id, "error": result["error"]},
            )
        )

    return _parse_persona_from_larva_result(result, spec_id)


# @invar:allow shell_result: Internal helper returns toolset object or None, not Result
# @shell_orchestration: MCP toolset discovery for larva server detection
async def _find_larva_toolset(mcp_servers: list[Any]) -> Any | None:
    """Find the MCP toolset that provides ``larva_resolve``.

    Iterates through available toolsets, gets their tool lists, and checks
    for the presence of ``larva_resolve``.

    Args:
        mcp_servers: List of MCP toolset objects.

    Returns:
        The toolset providing larva_resolve, or None.
    """
    for toolset in mcp_servers:
        try:
            # PydanticAI toolsets expose get_tools(ctx) → dict[str, ToolsetTool]
            # We need a RunContext-like object; pass None and catch if it fails
            tools = await toolset.get_tools(None)
            if LARVA_RESOLVE_TOOL in tools:
                return toolset
        except Exception:
            # Toolset might not be initialized or accessible
            continue
    return None


# @invar:allow shell_result: Internal helper returns raw MCP result, caller wraps in Result
# @shell_orchestration: MCP tool call bridge to larva_resolve
async def _call_larva_resolve(toolset: Any, spec_id: str) -> Any:
    """Call larva_resolve on the given toolset.

    Args:
        toolset: The MCP toolset that provides larva_resolve.
        spec_id: The persona ID to resolve.

    Returns:
        The raw result from the tool call.
    """
    tools = await toolset.get_tools(None)
    tool = tools[LARVA_RESOLVE_TOOL]

    result = await toolset.call_tool(
        LARVA_RESOLVE_TOOL,
        {"id": spec_id},
        None,  # ctx
        tool,
    )
    return result


def _parse_persona_from_larva_result(
    result: Any, spec_id: str
) -> Result[PersonaSpec, AnimaError]:
    """Parse a larva_resolve result into a PersonaSpec.

    The result may be a dict (PersonaSpec fields directly), a JSON string,
    or a list with a text content item (MCP response format).

    Args:
        result: Raw result from larva_resolve call.
        spec_id: The original spec_id for error messages.

    Returns:
        Success(PersonaSpec) or Failure with parse error.
    """
    spec_data: dict[str, Any] | None = None

    if isinstance(result, dict):
        spec_data = result
    elif isinstance(result, str):
        try:
            spec_data = json.loads(result)
        except json.JSONDecodeError as exc:
            return Failure(
                InvalidInputError(
                    f"Failed to parse larva_resolve('{spec_id}') response as JSON: {exc}",
                    {"spec_id": spec_id},
                )
            )
    elif isinstance(result, list) and result:
        # MCP response format: list of content items with .text
        first = result[0]
        text = getattr(first, "text", None) or (
            first.get("text") if isinstance(first, dict) else None
        )
        if isinstance(text, str):
            try:
                spec_data = json.loads(text)
            except json.JSONDecodeError as exc:
                return Failure(
                    InvalidInputError(
                        f"Failed to parse larva_resolve('{spec_id}') content as JSON: {exc}",
                        {"spec_id": spec_id},
                    )
                )

    if spec_data is None or not isinstance(spec_data, dict):
        return Failure(
            InvalidInputError(
                f"larva_resolve('{spec_id}') returned unexpected type: {type(result).__name__}",
                {"spec_id": spec_id},
            )
        )

    # Filter to PersonaSpec fields only (larva may include extra metadata)
    valid_fields = PersonaSpec.model_fields.keys()
    filtered = {k: v for k, v in spec_data.items() if k in valid_fields}

    try:
        persona = PersonaSpec(**filtered)
        logger.info(f"Resolved persona '{spec_id}' via larva: model={persona.model}")
        return Success(persona)
    except Exception as exc:
        return Failure(
            InvalidInputError(
                f"larva_resolve('{spec_id}') returned invalid PersonaSpec: {exc}",
                {"spec_id": spec_id, "error": str(exc)},
            )
        )
