"""Response extraction utilities for GenAI service.

Provides helper functions to extract text and structured objects from
CompletionEnvelope responses, making it easier to work with GenAIService outputs.

Connected modules:
  - gen_ai_service.models.domain (CompletionEnvelope, CompletionResult)
  - External: pydantic
"""

from __future__ import annotations

from typing import Any, Optional, TypeVar

from pydantic import BaseModel

from tnh_scholar.gen_ai_service.models.domain import CompletionEnvelope
from tnh_scholar.logging_config import get_logger

logger = get_logger("tnh_scholar.gen_ai_service.utils.response_utils")

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


def extract_text(envelope: CompletionEnvelope) -> str:
    """
    Extract text content from a completion envelope.

    Args:
        envelope: CompletionEnvelope from GenAIService.generate()

    Returns:
        str: The generated text content

    Raises:
        ValueError: If envelope has no result or text content

    Example:
        >>> envelope = service.generate(request)
        >>> text = extract_text(envelope)
        >>> print(text)
        "The generated response..."
    """
    if envelope.result is None:
        raise ValueError("Completion envelope has no result")

    if envelope.result.text is None:
        raise ValueError("Completion result has no text content")

    return envelope.result.text


def extract_object(envelope: CompletionEnvelope, model_class: Optional[type[T]] = None) -> T:
    """
    Extract structured object from a completion envelope.

    When using response_format with structured outputs, this extracts
    the parsed Pydantic model.

    Args:
        envelope: CompletionEnvelope from GenAIService.generate()
        model_class: Optional Pydantic model class for type checking

    Returns:
        BaseModel: The parsed structured object

    Raises:
        ValueError: If envelope has no result or structured object
        TypeError: If object doesn't match expected model_class

    Example:
        >>> class Translation(BaseModel):
        ...     original: str
        ...     translated: str
        ...
        >>> envelope = service.generate(request)
        >>> translation = extract_object(envelope, Translation)
        >>> print(translation.translated)
        "La traduction..."
    """
    if envelope.result is None:
        raise ValueError("Completion envelope has no result")

    parsed_obj = envelope.result.parsed
    if parsed_obj is None:
        raise ValueError(
            "Structured outputs are not yet supported. Ensure response_format is enabled "
            "and the provider returned a parsed object."
        )

    if model_class is not None and not isinstance(parsed_obj, model_class):
        raise TypeError(
            f"Expected object of type {model_class.__name__}, "
            f"got {type(parsed_obj).__name__}"
        )

    return parsed_obj


def extract_usage(envelope: CompletionEnvelope) -> dict[str, int]:
    """
    Extract token usage information from a completion envelope.

    Args:
        envelope: CompletionEnvelope from GenAIService.generate()

    Returns:
        dict: Token usage with keys: tokens_in, tokens_out, total_tokens

    Example:
        >>> envelope = service.generate(request)
        >>> usage = extract_usage(envelope)
        >>> print(f"Used {usage['total_tokens']} tokens")
        Used 1234 tokens
    """
    if envelope.result is None or envelope.result.usage is None:
        return {"tokens_in": 0, "tokens_out": 0, "total_tokens": 0}

    usage = envelope.result.usage

    # Domain Usage uses prompt/completion naming; fall back to transport-style keys
    tokens_in = getattr(usage, "prompt_tokens", None)
    if tokens_in is None:
        tokens_in = getattr(usage, "tokens_in", None)

    tokens_out = getattr(usage, "completion_tokens", None)
    if tokens_out is None:
        tokens_out = getattr(usage, "tokens_out", None)

    total_tokens = getattr(usage, "total_tokens", None)
    if total_tokens is None:
        total_tokens = getattr(usage, "tokens_total", None)

    return {
        "tokens_in": tokens_in or 0,
        "tokens_out": tokens_out or 0,
        "total_tokens": total_tokens or 0,
    }


def extract_provenance(envelope: CompletionEnvelope) -> dict[str, Any]:
    """
    Extract provenance metadata from a completion envelope.

    Useful for auditing, debugging, and tracking completions.

    Args:
        envelope: CompletionEnvelope from GenAIService.generate()

    Returns:
        dict: Provenance information including provider, model, timestamps, etc.

    Example:
        >>> envelope = service.generate(request)
        >>> prov = extract_provenance(envelope)
        >>> print(f"Generated by {prov['provider']}/{prov['model']}")
        Generated by openai/gpt-4o
    """
    if envelope.provenance is None:
        return {}

    prov = envelope.provenance

    data = {
        "provider": prov.provider,
        "model": prov.model,
        "started_at": prov.started_at.isoformat() if prov.started_at else None,
        "finished_at": prov.finished_at.isoformat() if prov.finished_at else None,
        "attempt_count": prov.attempt_count,
    }

    # Add fingerprint if available
    if prov.fingerprint:
        data["fingerprint"] = {
            "prompt_key": prov.fingerprint.prompt_key,
            "prompt_name": prov.fingerprint.prompt_name,
            "prompt_content_hash": prov.fingerprint.prompt_content_hash,
        }

    # Add duration if timestamps available
    if prov.started_at and prov.finished_at:
        duration = (prov.finished_at - prov.started_at).total_seconds()
        data["duration_seconds"] = duration

    return data


def extract_finish_reason(envelope: CompletionEnvelope) -> str | None:
    """
    Extract the finish reason from a completion envelope.

    Finish reason indicates why the completion stopped:
    - "stop": Natural completion
    - "length": Hit max_tokens limit
    - "content_filter": Blocked by content filter
    - etc.

    Args:
        envelope: CompletionEnvelope from GenAIService.generate()

    Returns:
        Optional[str]: Finish reason or None if not available

    Example:
        >>> envelope = service.generate(request)
        >>> reason = extract_finish_reason(envelope)
        >>> if reason == "length":
        ...     print("Warning: Completion truncated at max_tokens")
    """
    if envelope.result is None:
        return None

    finish_reason = getattr(envelope.result, "finish_reason", None)
    if finish_reason is None:
        return None

    # Handle enum or string
    if hasattr(finish_reason, "value"):
        return finish_reason.value

    return str(finish_reason) if finish_reason else None


def is_successful(envelope: CompletionEnvelope) -> bool:
    """
    Check if a completion was successful.

    Args:
        envelope: CompletionEnvelope from GenAIService.generate()

    Returns:
        bool: True if completion has valid result, False otherwise

    Example:
        >>> envelope = service.generate(request)
        >>> if is_successful(envelope):
        ...     text = extract_text(envelope)
        ... else:
        ...     print("Completion failed")
    """
    if envelope.result is None:
        return False

    if envelope.result.text is not None:
        return True

    parsed_obj = getattr(envelope.result, "parsed", None)
    return parsed_obj is not None


__all__ = [
    "extract_text",
    "extract_object",
    "extract_usage",
    "extract_provenance",
    "extract_finish_reason",
    "is_successful",
]
