"""Provider-agnostic structured output from LLM responses.

Models that support tool/function calling use `with_structured_output`.
All others get prompt-based JSON instructions with Pydantic parsing.
This keeps generators decoupled from provider-specific APIs.
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass
from typing import TYPE_CHECKING, Generic, TypeVar

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts.chat import HumanMessagePromptTemplate
from langchain_core.runnables import RunnableLambda
from pydantic import BaseModel, ValidationError

from unfamiliar.llm.providers.factory import provider_supports_tool_calling

if TYPE_CHECKING:
    from langchain_core.language_models.chat_models import BaseChatModel
    from langchain_core.runnables import Runnable

logger = logging.getLogger(__name__)

T = TypeVar("T", bound=BaseModel)

# Models whose with_structured_output does NOT work reliably.
# Most providers support tool calling; we denylist the exceptions.
# Note: GLM now goes through ChatOpenAI (via init_chat_model with base_url),
# which supports tool calling, so it's no longer on this list.
_NO_TOOL_CALLING_MODELS: set[str] = set()


@dataclass
class TruncationResult(Generic[T]):
    """Result of parsing LLM response with truncation detection.

    Attributes:
        recovered_data: Parsed model if any complete data recovered, else None.
        is_truncated: True if response was incomplete/truncated.
    """

    recovered_data: T | None
    is_truncated: bool


def detect_truncation(
    text: str,
    schema: type[T],
) -> TruncationResult[T]:
    """Parse LLM response and detect if JSON was truncated.

    Attempts to parse the response, recovering partial data if truncated.

    Args:
        text: Raw LLM response text.
        schema: Pydantic model class for the expected response shape.

    Returns:
        TruncationResult with recovered data and truncation metadata.
    """
    if not text or not text.strip():
        return TruncationResult(
            recovered_data=None,
            is_truncated=False,
        )

    # Try to extract JSON
    try:
        json_str = _extract_json(text)
        # Try direct parse
        try:
            data = schema.model_validate_json(json_str)
            return TruncationResult(
                recovered_data=data,
                is_truncated=False,
            )
        except ValidationError as exc:
            logger.debug("Schema validation failed, attempting recovery: %s", exc)

    except ValueError as exc:
        logger.debug("JSON extraction failed, attempting recovery: %s", exc)

    # Attempt recovery from truncated JSON
    brace_start = text.find("{")
    if brace_start == -1:
        return TruncationResult(
            recovered_data=None,
            is_truncated=True,
        )

    # Try to recover and get metadata
    recovery_info = _recover_truncated_json_with_info(text, brace_start, schema)
    if recovery_info.recovered_data is not None:
        return recovery_info

    return TruncationResult(
        recovered_data=None,
        is_truncated=True,
    )


def _recover_truncated_json_with_info(
    text: str,
    start_pos: int,
    schema: type[T],
) -> TruncationResult:
    """Attempt to recover usable data from truncated JSON with metadata.

    Delegates structural recovery to _recover_truncated_json. and adds
    schema validation plus TruncationResult wrapping.

    Args:
        text: Full response text.
        start_pos: Position of opening brace.
        schema: Target Pydantic schema.

    Returns:
        TruncationResult with recovered data and resume info.
    """
    # Delegate structural recovery
    recovered = _recover_truncated_json(text, start_pos)

    if recovered is None:
        return TruncationResult(
            recovered_data=None,
            is_truncated=True,
        )

    # Validate against schema
    try:
        data = schema.model_validate_json(recovered)
    except ValidationError as e:
        logger.warning("Recovered JSON failed schema validation: %s", e)
        return TruncationResult(
            recovered_data=None,
            is_truncated=True,
        )

    return TruncationResult(
        recovered_data=data,
        is_truncated=True,
    )


def create_structured_chain(
    model: BaseChatModel,
    prompt: ChatPromptTemplate,
    schema: type[T],
    provider_name: str | None = None,
) -> Runnable:
    """Build a chain that returns validated structured output.

    For models with reliable function calling, uses with_structured_output.
    Otherwise, appends JSON format instructions to the prompt and parses
    the text response with Pydantic.

    Args:
        model: Any BaseChatModel instance.
        prompt: The ChatPromptTemplate (system + human messages).
        schema: Pydantic model class for the expected response shape.
        provider_name: Provider name (e.g., "glm") to check provider-level
            tool calling support. When set, provider config takes precedence
            over class-name denylist.

    Returns:
        A Runnable whose .invoke() returns a validated schema instance.
    """
    if _should_use_structured_output(model, provider_name):
        structured_model = model.with_structured_output(schema)
        return prompt | structured_model

    augmented_prompt = _augment_prompt_with_json_instructions(prompt, schema)
    parser = _JsonOutputParser(schema)
    return augmented_prompt | model | RunnableLambda(parser.invoke)


def _should_use_structured_output(
    model: BaseChatModel,
    provider_name: str | None = None,
) -> bool:
    """Decide whether to use tool-calling structured output or JSON fallback.

    Checks provider-level config first (explicit flag), then falls back
    to class-name denylist for backwards compatibility.
    """
    if provider_name is not None and not provider_supports_tool_calling(provider_name):
        return False
    return _supports_tool_calling(model)


def _supports_tool_calling(model: BaseChatModel) -> bool:
    """Check whether a model implementation reliably handles function calling."""
    model_class = type(model).__name__
    return model_class not in _NO_TOOL_CALLING_MODELS


def _augment_prompt_with_json_instructions(
    prompt: ChatPromptTemplate,
    schema: type[BaseModel],
) -> ChatPromptTemplate:
    """Return a new prompt with JSON schema instructions appended.

    Appends JSON format instructions to the human message.
    """

    example = _build_example_structure(schema)
    example_text = json.dumps(example, indent=2)
    # Escape braces so LangChain doesn't treat JSON as template variables
    example_escaped = example_text.replace("{", "{{").replace("}", "}}")

    json_instructions = (
        "\n\n---\n"
        "## CRITICAL: Response Format\n\n"
        "Respond with ONLY valid JSON. No markdown fences, no commentary — pure JSON.\n"
        "Fill in every field with real content based on your analysis.\n"
        "Do NOT echo the schema back. Produce actual values.\n\n"
        f"Example structure (replace ALL placeholder values with your analysis):\n"
        f"{example_escaped}"
    )

    # Rebuild with augmented human message
    rebuilt: list[tuple[str, str]] = []
    for msg in prompt.messages:
        if isinstance(msg, HumanMessagePromptTemplate):
            prompt_obj = msg.prompt
            if hasattr(prompt_obj, "template"):
                template_text = str(prompt_obj.template)
                rebuilt.append(("human", template_text + json_instructions))
            else:
                logger.warning("Unexpected prompt type %s", type(prompt_obj))
                continue
        elif hasattr(msg, "prompt") and hasattr(msg.prompt, "template"):
            template_text = str(msg.prompt.template)
            rebuilt.append(("system", template_text))
        else:
            logger.debug("Skipping message type %s", type(msg).__name__)

    return ChatPromptTemplate.from_messages(rebuilt)


def _build_example_structure(schema: type[BaseModel]) -> dict:
    """Build a concrete example JSON from a Pydantic model.

    Produces placeholder values that show the LLM what shape to fill in,
    rather than dumping the raw JSON Schema (which models tend to echo back).
    """
    json_schema = schema.model_json_schema()
    defs = json_schema.get("$defs", {})
    return _example_from_schema(json_schema, defs)


def _example_from_schema(
    schema: dict,
    defs: dict,
) -> dict | list | str | int | float | bool | None:
    """Recursively build an example value from a JSON Schema node."""
    # Resolve $ref
    if "$ref" in schema:
        ref_name = schema["$ref"].split("/")[-1]
        if ref_name in defs:
            return _example_from_schema(defs[ref_name], defs)
        return "..."

    schema_type = schema.get("type")

    if schema_type == "object":
        result = {}
        for prop_name, prop_schema in schema.get("properties", {}).items():
            desc = prop_schema.get("description", "")
            hint = desc[:60] + "..." if len(desc) > 60 else desc
            result[prop_name] = _example_from_property(
                prop_name, prop_schema, defs, hint
            )
        return result

    if schema_type == "array":
        items = schema.get("items", {})
        return [_example_from_schema(items, defs)]

    if schema_type == "string":
        return "<your value here>"
    if schema_type == "integer":
        return 0
    if schema_type == "number":
        return 0.0
    if schema_type == "boolean":
        return False

    # anyOf / oneOf (e.g. Optional fields)
    for key in ("anyOf", "oneOf"):
        variants = schema.get(key, [])
        non_null = [v for v in variants if v.get("type") != "null"]
        if non_null:
            return _example_from_schema(non_null[0], defs)

    return "..."


def _example_from_property(
    name: str,
    prop_schema: dict,
    defs: dict,
    hint: str,
) -> dict | list | str | int | float | bool | None:
    """Build an example for a single property, using the hint for strings."""
    # Resolve $ref
    if "$ref" in prop_schema:
        return _example_from_schema(prop_schema, defs)

    prop_type = prop_schema.get("type")

    if prop_type == "array":
        items = prop_schema.get("items", {})
        return [_example_from_schema(items, defs)]

    if prop_type == "object":
        return _example_from_schema(prop_schema, defs)

    if prop_type == "string":
        return f"<{hint}>" if hint else "<your value here>"

    if prop_type == "integer":
        return 0
    if prop_type == "number":
        return 0.0
    if prop_type == "boolean":
        return False

    # anyOf / oneOf
    for key in ("anyOf", "oneOf"):
        variants = prop_schema.get(key, [])
        non_null = [v for v in variants if v.get("type") != "null"]
        if non_null:
            return _example_from_schema(non_null[0], defs)

    return "..."


class _JsonOutputParser:
    """Parse LLM text response as JSON and validate against a Pydantic schema."""

    def __init__(self, schema: type[T]) -> None:
        self._schema = schema

    def invoke(self, message: object) -> BaseModel:
        """Extract JSON from LLM response and validate against schema."""
        # LangChain messages have a .content attribute
        text = message.content if hasattr(message, "content") else str(message)
        json_str = _extract_json(text)
        try:
            return self._schema.model_validate_json(json_str)
        except ValidationError as exc:
            raise ValueError(
                f"LLM response did not match expected schema "
                f"{self._schema.__name__}: {exc}"
            ) from exc


def _extract_json(text: str) -> str:
    """Extract a JSON object or array from LLM text output.

    Handles common LLM response patterns:
    - Pure JSON
    - JSON wrapped in markdown code fences
    - JSON preceded/followed by commentary
    - Truncated JSON (attempts recovery)
    """
    text = text.strip()

    # Strip markdown code fences first (may be unclosed if truncated)
    fence_content = _extract_from_code_fence(text)
    if fence_content:
        text = fence_content

    # Try the raw/extracted text as-is
    if (text.startswith("{") and text.endswith("}")) or (
        text.startswith("[") and text.endswith("]")
    ):
        return text

    # Attempt to find complete JSON object
    brace_start = text.find("{")
    if brace_start != -1:
        # Walk forward to find the matching closing brace
        depth = 0
        in_string = False
        escape_next = False
        for i in range(brace_start, len(text)):
            char = text[i]
            if escape_next:
                escape_next = False
                continue
            if char == "\\":
                escape_next = True
                continue
            if char == '"':
                in_string = not in_string
                continue
            if in_string:
                continue
            if char == "{":
                depth += 1
            elif char == "}":
                depth -= 1
                if depth == 0:
                    return text[brace_start : i + 1]

        # If we get here, JSON is truncated - try recovery
        recovered = _recover_truncated_json(text, brace_start)
        if recovered:
            return recovered

    raise ValueError(
        f"Could not extract JSON from LLM response. First 500 chars: {text[:500]}"
    )


def _extract_from_code_fence(text: str) -> str | None:
    """Extract content from markdown code fence, handling unclosed fences.

    Returns None if no code fence found, or the content (possibly truncated).
    """
    # Try complete fence first
    fence_match = re.search(
        r"```(?:json)?\s*\n?(.*?)\n?\s*```",
        text,
        re.DOTALL,
    )
    if fence_match:
        return fence_match.group(1).strip()

    # Check for unclosed fence (truncated response)
    unclosed_match = re.search(
        r"```(?:json)?\s*\n?(.*)",
        text,
        re.DOTALL,
    )
    if unclosed_match:
        content = unclosed_match.group(1).strip()
        logger.warning(
            "Detected unclosed markdown fence - response likely truncated. "
            "Attempting recovery from %d chars.",
            len(content),
        )
        return content

    return None


def _find_last_complete_object(
    json_text: str,
    array_start: int,
) -> tuple[int | None, int]:
    """Find the last complete JSON object in an array using brace-depth tracking.

    Args:
        json_text: The JSON text to search.
        array_start: Position where the array content starts (after the '[' character).

    Returns:
        Tuple of (position after last complete object's '}', or None if none found,
        and count of complete objects found).
    """
    depth = 0
    in_string = False
    escape_next = False
    last_complete_end = None
    object_start = None
    object_count = 0

    for i in range(array_start, len(json_text)):
        char = json_text[i]

        if escape_next:
            escape_next = False
            continue
        if char == "\\":
            escape_next = True
            continue
        if char == '"':
            in_string = not in_string
            continue
        if in_string:
            continue

        if char == "{":
            if depth == 0:
                object_start = i
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0 and object_start is not None:
                # Found a complete object
                last_complete_end = i + 1
                object_count += 1
        elif char == "]" and depth == 0:
            # End of the array - stop scanning
            break

    return last_complete_end, object_count


def _recover_chains_from_text(
    json_text: str,
    search_start: int,
) -> tuple[str, bool]:
    """Attempt to recover chain objects from the JSON text after missions.

    Scans for '"chains"' key and extracts complete chain objects using
    the same brace-depth parser.

    Args:
        json_text: The JSON text to search.
        search_start: Position in json_text where to start searching for chains.

    Returns:
        Tuple of (JSON string of recovered chain objects joined by commas,
        bool indicating if chains were present but lost due to truncation).
    """
    # Find the chains key and array start
    chains_match = re.search(r'"chains"\s*:\s*\[', json_text[search_start:])
    if not chains_match:
        return "", False

    chains_array_start = search_start + chains_match.end()

    # Find all complete chain objects
    last_complete_end, chain_count = _find_last_complete_object(
        json_text, chains_array_start
    )

    if last_complete_end is None:
        # Chains section exists but no complete chains
        return "", True  # chains_lost = True

    # Extract the complete chain objects
    chains_text = json_text[chains_array_start:last_complete_end]

    # Validate we got complete objects
    if not chains_text.strip():
        return "", True

    return chains_text.strip().rstrip(","), False  # chains_lost = False


def _recover_truncated_json(text: str, start_pos: int) -> str | None:
    """Attempt to recover usable data from truncated JSON.

    Looks for the last complete object in an array and closes the structure.
    Returns recovered JSON string or None if recovery fails.
    """
    # Focus on the JSON part
    json_text = text[start_pos:]

    # Check if this looks like a missions response
    if '"missions"' not in json_text:
        return None

    logger.warning(
        "Attempting JSON recovery from truncated response (%d chars)",
        len(json_text),
    )

    # Find the missions array start
    missions_match = re.search(r'"missions"\s*:\s*\[', json_text)
    if not missions_match:
        return None

    array_start = missions_match.end()

    # Find the last complete mission object using the brace-depth parser
    last_complete_end, _ = _find_last_complete_object(json_text, array_start)

    if last_complete_end is None:
        logger.warning("JSON recovery failed: no complete mission objects found")
        return None

    # Build recovered JSON: truncate to last complete object and close missions array
    recovered = json_text[:last_complete_end] + "]"

    # Now try to recover chains
    recovered_chains, chains_lost = _recover_chains_from_text(
        json_text, last_complete_end
    )

    if recovered_chains:
        recovered += ', "chains": [' + recovered_chains + "]}"
        logger.warning(
            "JSON recovery: recovered %d missions and chains from truncated",
            recovered.count('"title_plain"'),
        )
    else:
        recovered += ', "chains": []}'
        if chains_lost:
            logger.warning(
                "JSON recovery: chains section was truncated, recovered missions only"
            )
        else:
            logger.warning(
                "JSON recovery: recovered %d missions from truncated (no chains)",
                recovered.count('"title_plain"'),
            )

    # Validate it's parseable
    try:
        json.loads(recovered)
        return recovered
    except json.JSONDecodeError as e:
        logger.warning("JSON recovery failed validation: %s", e)
        return None
