"""Usage extraction from LLM responses for nanomon."""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from typing import Any

logger = logging.getLogger(__name__)


class UsageExtractor(ABC):
    """Abstract base class for extracting usage information from LLM responses."""

    @abstractmethod
    def extract_usage(self, response: Any) -> tuple[int, int] | None:
        """Extract input and output token counts from a response.

        Args:
            response: The LLM response object

        Returns:
            Tuple of (input_tokens, output_tokens) or None if cannot extract
        """
        ...

    @abstractmethod
    def extract_model_name(
        self,
        response: Any,
        lm: Any,
        kwargs: dict[str, Any],
    ) -> str | None:
        """Extract the model name from response, LM, or kwargs.

        Args:
            response: The LLM response object
            lm: The language model instance
            kwargs: Arguments passed to the LM call

        Returns:
            Model name string or None if cannot extract
        """
        ...


class DefaultUsageExtractor(UsageExtractor):
    """Default extractor supporting common response shapes.

    Supports:
        - OpenAI style: response.usage.prompt_tokens / completion_tokens
        - OpenAI dict: response["usage"]["prompt_tokens"] / "completion_tokens"
        - Anthropic style: response.usage.input_tokens / output_tokens
        - Anthropic dict: response["usage"]["input_tokens"] / "output_tokens"
    """

    def extract_usage(self, response: Any) -> tuple[int, int] | None:
        """Extract usage from response using common patterns."""
        # Try object attribute access (response.usage.X)
        usage = getattr(response, "usage", None)
        if usage is not None:
            result = self._extract_from_usage_object(usage)
            if result is not None:
                return result

        # Try dictionary access (response["usage"])
        if isinstance(response, dict) and "usage" in response:
            usage_dict = response["usage"]
            result = self._extract_from_usage_dict(usage_dict)
            if result is not None:
                return result

        logger.debug("Could not extract usage from response")
        return None

    def _extract_from_usage_object(self, usage: Any) -> tuple[int, int] | None:
        """Extract from an object with attributes."""
        input_tokens = 0
        output_tokens = 0

        # OpenAI style: prompt_tokens / completion_tokens
        if hasattr(usage, "prompt_tokens") and hasattr(usage, "completion_tokens"):
            input_tokens = getattr(usage, "prompt_tokens", 0) or 0
            output_tokens = getattr(usage, "completion_tokens", 0) or 0
            return (input_tokens, output_tokens)

        # Anthropic style: input_tokens / output_tokens
        if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"):
            input_tokens = getattr(usage, "input_tokens", 0) or 0
            output_tokens = getattr(usage, "output_tokens", 0) or 0
            return (input_tokens, output_tokens)

        return None

    def _extract_from_usage_dict(self, usage: dict[str, Any]) -> tuple[int, int] | None:
        """Extract from a dictionary."""
        # OpenAI style
        if "prompt_tokens" in usage or "completion_tokens" in usage:
            input_tokens = usage.get("prompt_tokens", 0) or 0
            output_tokens = usage.get("completion_tokens", 0) or 0
            return (input_tokens, output_tokens)

        # Anthropic style
        if "input_tokens" in usage or "output_tokens" in usage:
            input_tokens = usage.get("input_tokens", 0) or 0
            output_tokens = usage.get("output_tokens", 0) or 0
            return (input_tokens, output_tokens)

        return None

    def extract_model_name(
        self,
        response: Any,
        lm: Any,
        kwargs: dict[str, Any],
    ) -> str | None:
        """Extract model name with priority: lm.model -> kwargs -> response."""
        # Priority 1: LM's model attribute
        if hasattr(lm, "model") and lm.model:
            return str(lm.model)

        # Priority 2: kwargs model parameter
        if "model" in kwargs and kwargs["model"]:
            return str(kwargs["model"])

        # Priority 3: Response model attribute
        if hasattr(response, "model") and response.model:
            return str(response.model)

        # Priority 4: Response dict model key
        if isinstance(response, dict) and "model" in response:
            return str(response["model"])

        return None


class ExtractorChain:
    """Chain of extractors that tries each in order until one succeeds."""

    def __init__(self, extractors: list[UsageExtractor] | None = None) -> None:
        """Initialize with a list of extractors.

        Args:
            extractors: List of extractors to try in order.
                       If None, uses DefaultUsageExtractor.
        """
        self._extractors: list[UsageExtractor] = extractors or [DefaultUsageExtractor()]

    def register(self, extractor: UsageExtractor) -> None:
        """Register an additional extractor at highest priority.

        Args:
            extractor: Extractor to add to the front of the chain
        """
        self._extractors.insert(0, extractor)

    def extract_usage(self, response: Any) -> tuple[int, int]:
        """Extract usage from response using the chain of extractors.

        Returns:
            Tuple of (input_tokens, output_tokens), defaults to (0, 0)
        """
        for extractor in self._extractors:
            result = extractor.extract_usage(response)
            if result is not None:
                return result

        logger.debug("No extractor could extract usage, returning (0, 0)")
        return (0, 0)

    def extract_model_name(
        self,
        response: Any,
        lm: Any,
        kwargs: dict[str, Any],
    ) -> str:
        """Extract model name using the chain of extractors.

        Returns:
            Model name string, defaults to "unknown"
        """
        for extractor in self._extractors:
            result = extractor.extract_model_name(response, lm, kwargs)
            if result is not None:
                return result

        logger.debug("No extractor could extract model name, returning 'unknown'")
        return "unknown"
