"""MCP sampling-based extraction engine for the 7 analytical lenses.

When the host LLM does not provide structured hints via tool parameters,
this module uses ``ctx.sample()`` to ask the LLM to extract analytical
structure from plain text.  Each sampling call is independent (KV₇) and
wrapped in error handling so failures degrade to text fallback (KV₃).

Usage:
    hints = await extract_structured_hints(ctx, text, ["Bayes", "FOL"])
"""
from __future__ import annotations

import json
import logging
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


async def extract_structured_hints(
    ctx: Any,
    text: str,
    fire_plan: List[str],
    timeout_ms: int = 10_000,
) -> Dict[str, dict]:
    """Use MCP sampling to extract structured models from text.

    Only samples for lenses in *fire_plan* that don't already have hints.
    Each sampling call is independent (KV₇).  Returns a dict mapping
    ``lens_id → hint dict``; entries may be empty on failure.

    Parameters
    ----------
    ctx : mcp.server.fastmcp.Context
        The MCP context (must support ``ctx.sample()``).
    text : str
        The user's input text to analyse.
    fire_plan : List[str]
        Lens IDs to attempt extraction for.
    timeout_ms : int
        Maximum time per sample call in milliseconds (unused in current
        implementation — relies on host-side timeout).

    Returns
    -------
    Dict[str, dict]
        Mapping of lens_id → extracted hints.  Missing lenses were either
        not in the fire plan or extraction failed gracefully.
    """
    from fw_server.sampling_prompts import SAMPLING_PROMPTS

    hints: Dict[str, dict] = {}

    for lens_id in fire_plan:
        prompt_template = SAMPLING_PROMPTS.get(lens_id)
        if prompt_template is None:
            continue

        try:
            prompt = prompt_template.format(text=text[:3000])
            response = await ctx.sample(prompt)

            # ctx.sample() returns a string (the LLM's response text)
            response_text = _extract_text(response)
            if not response_text:
                continue

            parsed = _parse_json_response(response_text)
            if parsed and not parsed.get("empty"):
                hints[lens_id] = parsed

        except Exception:
            # KV₃: sampling failure never propagates — degrade to text
            logger.debug(
                "Sampling extraction failed for %s", lens_id, exc_info=True
            )

    return hints


def _extract_text(response: Any) -> Optional[str]:
    """Extract plain text from a sampling response.

    Handles both string responses and structured message objects
    depending on the MCP SDK version.
    """
    if isinstance(response, str):
        return response.strip()
    # Handle CreateMessageResult / similar structured response
    if hasattr(response, "content"):
        content = response.content
        if isinstance(content, str):
            return content.strip()
        if hasattr(content, "text"):
            return content.text.strip()
        # content may be a list of content blocks
        if isinstance(content, list):
            texts = []
            for block in content:
                if hasattr(block, "text"):
                    texts.append(block.text)
                elif isinstance(block, str):
                    texts.append(block)
            return "\n".join(texts).strip() if texts else None
    if hasattr(response, "text"):
        return response.text.strip()
    return None


def _parse_json_response(text: str) -> Optional[dict]:
    """Parse a JSON response from the LLM, tolerating markdown fences."""
    text = text.strip()

    # Strip markdown code fences if present
    if text.startswith("```"):
        lines = text.split("\n")
        # Remove first line (```json or ```) and last line (```)
        lines = [
            line for i, line in enumerate(lines)
            if not (i == 0 and line.startswith("```"))
            and not (i == len(lines) - 1 and line.strip() == "```")
        ]
        text = "\n".join(lines).strip()

    # Find the first { and last } to extract JSON from surrounding text
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end <= start:
        return None

    json_str = text[start:end + 1]
    try:
        result = json.loads(json_str)
        if isinstance(result, dict):
            return result
    except (json.JSONDecodeError, ValueError):
        pass

    return None
