"""
Canonical token extraction from LangChain LLMResult objects.

Single source of truth for extracting token usage, model name, and response
text from LangChain's LLMResult. Replaces inline extraction duplicated across
auto_instrument/langgraph.py and callback.py.

Priority order:
1. llm_output.token_usage (OpenAI standard)
2. llm_output.usage (generic)
3. generation_info.usage_metadata (Gemini via langchain_google_genai)
4. message.usage_metadata (newer LangChain / Gemini)
5. message.response_metadata.usage (Anthropic via langchain_anthropic)
"""

import logging
from dataclasses import dataclass
from typing import Any, Optional

logger = logging.getLogger(__name__)


@dataclass
class ExtractedUsage:
    """Token usage and metadata extracted from an LLMResult."""
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    model: Optional[str] = None
    response_text: Optional[str] = None


def extract_from_langchain_response(response: Any) -> ExtractedUsage:
    """Extract token usage from a LangChain LLMResult.

    Handles OpenAI, Anthropic, Gemini, and generic LangChain formats
    with a single prioritized fallback chain.

    Args:
        response: A LangChain LLMResult (or similar) object.

    Returns:
        ExtractedUsage with whatever fields could be populated.
    """
    result = ExtractedUsage()

    if response is None:
        return result

    # --- Extract model name ---
    if hasattr(response, "llm_output") and response.llm_output:
        result.model = (
            response.llm_output.get("model_name")
            or response.llm_output.get("model")
        )

    # --- Path 1: llm_output.token_usage (OpenAI) or llm_output.usage ---
    if hasattr(response, "llm_output") and response.llm_output:
        usage = (
            response.llm_output.get("token_usage")
            or response.llm_output.get("usage")
            or {}
        )
        if usage:
            result.input_tokens = usage.get("prompt_tokens") or usage.get("input_tokens", 0)
            result.output_tokens = usage.get("completion_tokens") or usage.get("output_tokens", 0)
            result.total_tokens = usage.get("total_tokens", 0)

    # --- Path 2: generation-level extraction ---
    gen = _get_first_generation(response)
    if gen is not None:
        # Extract response text
        if hasattr(gen, "text") and gen.text:
            result.response_text = gen.text

        # generation_info.usage_metadata (Gemini)
        if not result.total_tokens:
            gen_info = getattr(gen, "generation_info", None) or {}
            usage_meta = gen_info.get("usage_metadata") or {}
            if usage_meta:
                result.input_tokens = (
                    usage_meta.get("prompt_token_count")
                    or usage_meta.get("input_tokens", 0)
                )
                result.output_tokens = (
                    usage_meta.get("candidates_token_count")
                    or usage_meta.get("output_tokens", 0)
                )
                result.total_tokens = (
                    usage_meta.get("total_token_count")
                    or usage_meta.get("total_tokens", 0)
                )

        # message.usage_metadata (newer LangChain, Gemini, Anthropic)
        if not result.total_tokens:
            msg = getattr(gen, "message", None)
            if msg and hasattr(msg, "usage_metadata") and msg.usage_metadata:
                um = msg.usage_metadata
                if isinstance(um, dict):
                    result.input_tokens = um.get("input_tokens", 0)
                    result.output_tokens = um.get("output_tokens", 0)
                    result.total_tokens = um.get("total_tokens", 0)
                else:
                    result.input_tokens = getattr(um, "input_tokens", 0)
                    result.output_tokens = getattr(um, "output_tokens", 0)
                    result.total_tokens = getattr(um, "total_tokens", 0)

            # message.response_metadata.usage (Anthropic)
            if not result.total_tokens and msg:
                resp_meta = getattr(msg, "response_metadata", None) or {}
                if isinstance(resp_meta, dict):
                    usage = resp_meta.get("usage") or {}
                    if usage:
                        result.input_tokens = usage.get("input_tokens", 0)
                        result.output_tokens = usage.get("output_tokens", 0)
                        result.total_tokens = usage.get("total_tokens", 0)
                    # Also try to get model from response_metadata
                    if not result.model:
                        result.model = resp_meta.get("model") or resp_meta.get("model_id")

    # --- Compute total if missing ---
    if not result.total_tokens and (result.input_tokens or result.output_tokens):
        result.total_tokens = result.input_tokens + result.output_tokens

    return result


def _get_first_generation(response: Any) -> Any:
    """Safely get the first generation from an LLMResult."""
    try:
        if hasattr(response, "generations") and response.generations:
            first_list = response.generations[0]
            if first_list:
                return first_list[0]
    except (IndexError, AttributeError):
        pass
    return None
