"""
Cost estimation module for OpenAI API responses.

This is the main public interface of the library, providing a simple
function to estimate costs for various types of OpenAI API responses.
"""

from __future__ import annotations

from typing import Any, Dict, Iterable, List, Optional, Union

from .calculation import calculate_cost, format_usd
from .response_parser import extract_model_details, extract_usage
from .pricing_data import load_pricing, get_model_pricing


class CostEstimateError(Exception):
    """
    Unified exception for all cost estimation errors.

    This exception is raised for any recoverable errors during cost estimation,
    including invalid inputs, missing attributes, and pricing data issues.
    """

    pass


def _find_last_chunk_with_usage(stream: Iterable[Any]) -> Any:
    """
    Extract the last chunk from a stream that contains usage information.

    Args:
        stream: An iterable of response chunks

    Returns:
        The last chunk containing usage information

    Raises:
        CostEstimateError: If no chunks contain usage information
    """
    last_chunk = None

    for chunk in stream:
        if hasattr(chunk, "usage"):
            last_chunk = chunk

    if last_chunk is None:
        raise CostEstimateError("Stream contained no chunks with usage information")

    return last_chunk


def _get_model_rates(model_name: str, model_date: str) -> Dict[str, float]:
    """
    Find the appropriate pricing rates for a model.

    Uses a multi-stage lookup strategy:
    1. Exact match with model name and date
    2. Match using the full versioned model name
    3. Match using base model name
    4. Fuzzy match based on model name prefix/substring
    5. Find the latest available pricing for that model

    Args:
        model_name: Base model name (e.g., "gpt-4o-mini")
        model_date: Version date or "latest"

    Returns:
        Dict containing input_price, cached_input_price, and output_price

    Raises:
        CostEstimateError: If no pricing data can be found for the model
    """
    pricing_data = load_pricing()

    # Strategy 1: Exact match
    exact_key = (model_name, model_date)
    if exact_key in pricing_data:
        return pricing_data[exact_key]

    # Strategy 2: Full versioned model name
    if model_date != "latest" and "-" in model_name:
        versioned_key = (f"{model_name}-{model_date}", "latest")
        if versioned_key in pricing_data:
            return pricing_data[versioned_key]

    # Strategy 3: Base model match
    general_key = (model_name, "latest")
    if general_key in pricing_data:
        return pricing_data[general_key]

    # Strategy 4: Fuzzy match based on prefix/substring
    for (name, date), rates in pricing_data.items():
        # Get first part of multiline model names
        name_first_part = name.split("\n")[0].strip()

        # Check if our model name is in the name or vice versa
        if model_name in name_first_part or name_first_part in model_name:
            return rates

    # Strategy 5: Latest available for this model (if date specified)
    candidates = [
        (date, rates)
        for (name, date), rates in pricing_data.items()
        if name == model_name and date <= model_date
    ]

    if candidates:
        # Pick the newest among older dates
        _, rates = max(candidates, key=lambda x: x[0])
        return rates

    raise CostEstimateError(
        f"No pricing data found for model '{model_name}' (date: {model_date})"
    )


def ctoken(response: Any) -> Dict[str, Any]:
    """
    Estimate token usage and cost for an OpenAI API response.

    Supports multiple response types:
    - Single ChatCompletion response object
    - Stream of ChatCompletionChunk objects
    - Responses API object

    Args:
        response: An OpenAI API response object or stream

    Returns:
        Dict containing detailed cost breakdown:
            - prompt_tokens: Number of input tokens
            - completion_tokens: Number of output tokens
            - total_tokens: Total token count
            - cached_tokens: Number of cached tokens
            - prompt_cost_uncached: Cost of non-cached prompt tokens (USD as string)
            - prompt_cost_cached: Cost of cached prompt tokens (USD as string)
            - completion_cost: Cost of completion tokens (USD as string)
            - total_cost: Total cost (USD as string)

    Raises:
        CostEstimateError: For any issues during estimation
    """
    try:
        # Handle different response types
        if hasattr(response, "__iter__") and not hasattr(response, "model"):
            # This is a stream of chunks - get the last one with usage data
            chunk = _find_last_chunk_with_usage(response)
        else:
            # This is a single response object
            chunk = response

        # Extract token usage and model details
        usage_data = extract_usage(chunk)
        model_info = extract_model_details(chunk.model)
        pricing_rates = _get_model_rates(
            model_info["model_name"], model_info["model_date"]
        )

        # Calculate and return cost breakdown
        result = calculate_cost(usage_data, pricing_rates)

        # Only format cost values as strings, keeping token counts as numbers
        for key, value in result.items():
            if isinstance(value, (int, float)) and key in [
                "prompt_cost_uncached",
                "prompt_cost_cached",
                "completion_cost",
                "total_cost",
            ]:
                result[key] = format_usd(value)

        return result

    except Exception as e:
        # Wrap all exceptions in our unified error type
        if isinstance(e, CostEstimateError):
            raise
        raise CostEstimateError(str(e)) from e


# Alias for backward compatibility
estimate_cost = ctoken


def estimate_openai_api_cost(
    model: Union[str, Any],
    messages: Optional[List[Dict[str, str]]] = None,
    prompt: Optional[str] = None,
    max_tokens: int = 0,
) -> Union[float, Dict[str, Any]]:
    """
    Estimate the cost of an OpenAI API call before making it.

    This function estimates the cost of a completion or chat completion
    request based on the model, messages/prompt, and maximum tokens.
    Can also handle a completed ChatCompletion object to calculate actual cost.

    Args:
        model: The model identifier (e.g., "gpt-4", "gpt-3.5-turbo") or a ChatCompletion object
        messages: List of message dictionaries for chat completions
        prompt: Text prompt for completions
        max_tokens: Maximum number of tokens to generate in the output

    Returns:
        When estimating a future call: Estimated cost in USD as a float
        When processing a completed call: Dict with token counts and costs

    Raises:
        ValueError: If the model is not found or inputs are invalid
    """
    if not model:
        raise ValueError("Model identifier is required")

    # Extract model name if it's a ChatCompletion or similar object
    model_name = model
    if hasattr(model, "model"):
        model_name = model.model

    # Get model pricing
    model_pricing = get_model_pricing(model_name)
    if not model_pricing:
        raise ValueError(f"Model '{model_name}' not found in pricing data")

    # If model is a ChatCompletion object with usage, calculate directly from usage
    if hasattr(model, "usage"):
        prompt_tokens = getattr(model.usage, "prompt_tokens", 0)
        completion_tokens = getattr(model.usage, "completion_tokens", 0)
        total_tokens = prompt_tokens + completion_tokens
        cached_tokens = getattr(
            getattr(model.usage, "prompt_tokens_details", None), "cached_tokens", 0
        )

        # Calculate costs
        prompt_cost_uncached = model_pricing["input_cost_per_1k"] * (
            (prompt_tokens - cached_tokens) / 1000
        )
        prompt_cost_cached = (
            model_pricing["input_cost_per_1k"] * (cached_tokens / 1000) * 0.25
        )  # Assuming cached is 25% of cost
        completion_cost = model_pricing["output_cost_per_1k"] * (
            completion_tokens / 1000
        )
        total_cost = prompt_cost_uncached + prompt_cost_cached + completion_cost

        # Format the result like the ctoken function
        result = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
            "cached_tokens": cached_tokens,
            "prompt_cost_uncached": format_usd(prompt_cost_uncached),
            "prompt_cost_cached": format_usd(prompt_cost_cached),
            "completion_cost": format_usd(completion_cost),
            "total_cost": format_usd(total_cost),
        }

        return result

    # Calculate input token count for estimation
    input_tokens = 0
    if messages:
        # Estimate tokens for chat completions
        for message in messages:
            content = message.get("content", "")
            if content:
                # Very rough estimation: 1 token ≈ 4 characters for English text
                input_tokens += len(content) // 4
    elif prompt:
        # Estimate tokens for completions
        input_tokens += len(prompt) // 4
    else:
        raise ValueError("Either messages or prompt is required for estimation")

    # Add token margin for system overhead (10%)
    input_tokens = int(input_tokens * 1.1)

    # Estimate cost
    input_cost = model_pricing["input_cost_per_1k"] * (input_tokens / 1000)
    output_cost = model_pricing["output_cost_per_1k"] * (max_tokens / 1000)

    return input_cost + output_cost


def estimate_openai_api_cost_from_response(response: Dict[str, Any]) -> float:
    """
    Calculate the exact cost of an OpenAI API response.

    Args:
        response: OpenAI API response dictionary containing:
            - model: The model identifier
            - usage: Dictionary with prompt_tokens and completion_tokens

    Returns:
        Cost in USD

    Raises:
        ValueError: If the response is invalid or missing required fields
    """
    if not isinstance(response, dict):
        raise ValueError("Response must be a dictionary")

    if "model" not in response:
        raise ValueError("Response missing 'model' field")

    if "usage" not in response or not isinstance(response["usage"], dict):
        raise ValueError("Response missing 'usage' field or invalid usage format")

    model = response["model"]
    usage = response["usage"]

    # Get model pricing
    model_pricing = get_model_pricing(model)
    if not model_pricing:
        raise ValueError(f"Model '{model}' not found in pricing data")

    # Extract token counts
    prompt_tokens = usage.get("prompt_tokens", 0)
    completion_tokens = usage.get("completion_tokens", 0)

    # Calculate cost
    input_cost = model_pricing["input_cost_per_1k"] * (prompt_tokens / 1000)
    output_cost = model_pricing["output_cost_per_1k"] * (completion_tokens / 1000)

    return input_cost + output_cost
