"""
Utilities for model name normalization and usage recording.

This module provides centralized functions for normalizing model names
and recording model usage for consistent tracking and reporting.
"""

import logging
import re

import litellm
from inspect_ai.log import ModelEvent, transcript
from inspect_ai.model import GenerateConfig, ModelOutput, ModelUsage
from inspect_ai.model._model import record_and_check_model_usage

logger = logging.getLogger(__name__)

# Regex to validate model names; allows names with or without provider prefix
VALID_MODEL_NAME_REGEX = re.compile(r"^([^\s\/]+/)*[^\s\/]+$")


def normalize_model_name(model_name: str) -> str:
    """
    Normalize model name to the format <provider>/<model_name>.

    Uses litellm's get_llm_provider function to identify the provider
    for a given model name and format it as provider/model_name.

    Args:
        model_name: The model name to normalize

    Returns:
        A normalized model name in the format <provider>/<model_name>
        or the original model name if provider couldn't be determined
    """
    if not VALID_MODEL_NAME_REGEX.match(model_name):
        raise ValueError(
            f"Invalid model name cannot be normalized: '{model_name}'; must match regex '{VALID_MODEL_NAME_REGEX.pattern}'"
        )

    # If already has provider prefix, return as-is
    # Note: it will have a false-positive in cases like
    #
    # "meta-llama/Llama-4-Scout-17B-16E-Instruct"
    # which should map to: "together/meta-llama/Llama-4-Scout-17B-16E-Instruct"
    #
    # since it is hard to know if "meta-llama" could be a provider.
    if "/" in model_name:
        return model_name

    try:
        # get_llm_provider returns (model, provider, api_base, api_base_from_map)
        _, provider, _, _ = litellm.get_llm_provider(model_name)

        if provider:
            return f"{provider}/{model_name}"
        else:
            logger.warning(f"Unable to determine provider for model {model_name}")
            return model_name
    except Exception as e:
        logger.warning(f"Error normalizing model name '{model_name}': {e}")
        return model_name


def _is_valid_model_usage(usage: ModelUsage) -> tuple[bool, str | None]:
    """
    Check if the ModelUsage object has valid token counts.

    Args:
        usage: ModelUsage object to validate

    Returns:
        Tuple of:
        is_valid: True if usage is valid, False otherwise
        error_message: Error message if invalid, None if valid
    """
    if not all(
        isinstance(x, int)
        for x in [usage.input_tokens, usage.output_tokens, usage.total_tokens]
    ):
        return (
            False,
            "ModelUsage must include input_tokens, output_tokens, and total_tokens",
        )

    if usage.input_tokens < 0 or usage.output_tokens < 0 or usage.total_tokens < 0:
        return False, "Token counts must be non-negative"

    # Some providers give token counts that partially overlap, so may add to
    # more than the "total".  However, it is never acceptable to get:
    # - `total_tokens` greater than the sum of other fields, since that means
    #   there are tokens not accounted for by a more specific category.
    # - `total_tokens` less than the larger of `input_tokens` and
    #   `output_tokens` (and possibly others, but it's unclear if providers
    #   consistently include `reasoning_tokens` in the total)
    all_counts = {
        k: v
        for k, v in usage.model_dump().items()
        if isinstance(v, int) and k != "total_tokens"
    }
    if usage.total_tokens > sum(all_counts.values()):
        return (
            False,
            "total_tokens cannot be greater than the sum of specific token types",
        )

    if usage.total_tokens < max(usage.input_tokens, usage.output_tokens):
        return (
            False,
            "total_tokens cannot be less than the larger of input_tokens and output_tokens",
        )

    return True, None


def record_model_usage_with_inspect(
    model_name: str,
    usage: ModelUsage,
    normalize_name: bool = True,
    allow_invalid: bool = False,
) -> None:
    """Convenience function to record model usage with inspect_ai for tracking
    and reporting.

    See record_model_usage_event_with_inspect for more details; this just
    creates a ModelEvent with default blank values other than the name/usage
    info.

    By default, this function performs basic validation of the token counts to
    make sure they are consistent and complete (must include at least
    input/output/total, total should be no more than sum of other fields) and
    throws a ValueError if they are not.  To log a warning instead of erroring,
    use `allow_invalid=True`.
    """

    is_valid, error_message = _is_valid_model_usage(usage)
    if not is_valid:
        if allow_invalid:
            logger.warning(
                f"Invalid ModelUsage for model '{model_name}': {error_message}"
            )
        else:
            raise ValueError(
                f"Invalid ModelUsage for model '{model_name}': {error_message}"
            )

    # missing finer details like input/output/tools so we leave those blank
    event = ModelEvent(
        model=model_name,
        input=[],
        tools=[],
        tool_choice="auto",
        config=GenerateConfig(),
        output=ModelOutput(model=model_name, usage=usage),
        cache=None,
        call=None,
        pending=False,
    )

    record_model_usage_event_with_inspect(event, normalize_name=normalize_name)


def record_model_usage_event_with_inspect(
    event: ModelEvent, normalize_name: bool = True
) -> None:
    """
    Record ModelEvent with inspect_ai for usage tracking and reporting.

    Note: This should not be used for models implemented via Inspect, since
    they already log automatically; this function is useful for cases where
    e.g. an external system exposes internal model API calls.

    This function normalizes the model name by default and records token usage
    in a format that will appear in both the ModelEvents transcript and the
    usage summary.

    Args:
        event: ModelEvent to record; `event.output.usage` should contain usage info
        normalize_name: Whether to normalize the model name (default: True)
    """

    if normalize_name:
        event = event.model_copy()
        event.model = normalize_model_name(event.model)
        event.output = event.output.model_copy()
        event.output.model = normalize_model_name(event.output.model)

    # ModelEvents show up in `inspect view`
    transcript()._event(event)

    # Model usage is tracked via record_model_usage as well; this determines
    # what gets shown in the cli printout at the end of `inspect eval`
    assert (
        event.output.usage is not None
    ), "ModelEvent output must have usage information"
    record_and_check_model_usage(event.model, event.output.usage)
