"""
For calculating cost of fireworks ai serverless inference models.
"""

from typing import Final

from litellm.constants import (
    FIREWORKS_AI_4_B,
    FIREWORKS_AI_16_B,
    FIREWORKS_AI_56_B_MOE,
    FIREWORKS_AI_176_B_MOE,
)
from litellm.types.utils import Usage
from litellm.utils import get_model_info


# Extract the number of billion parameters from the model name
# only used for together_computer LLMs
def get_base_model_for_pricing(model_name: str) -> str:
    """
    Helper function for calculating together ai pricing.

    Returns:
    - str: model pricing category if mapped else received model name
    """
    import re

    model_name = model_name.lower()

    # Check for MoE models in the form <number>x<number>b
    moe_match: Final = re.search(r"(\d+)x(\d+)b", model_name)
    if moe_match:
        total_billion: Final = int(moe_match.group(1)) * int(moe_match.group(2))
        if total_billion <= FIREWORKS_AI_56_B_MOE:
            return "fireworks-ai-moe-up-to-56b"
        elif total_billion <= FIREWORKS_AI_176_B_MOE:
            return "fireworks-ai-56b-to-176b"

    # Check for standard models in the form <number>b
    re_params_match: Final = re.search(r"(\d+)b", model_name)
    if re_params_match is not None:
        params_match: Final = str(re_params_match.group(1))
        params_billion: Final = float(params_match)

        # Determine the category based on the number of parameters
        if params_billion <= FIREWORKS_AI_4_B:
            return "fireworks-ai-up-to-4b"
        elif params_billion <= FIREWORKS_AI_16_B:
            return "fireworks-ai-4.1b-to-16b"
        elif params_billion > FIREWORKS_AI_16_B:
            return "fireworks-ai-above-16b"

    # If no matches, return the original model_name
    return "fireworks-ai-default"


def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
    """
    Calculates the cost per token for a given model, prompt tokens, and completion tokens.

    Input:
        - model: str, the model name without provider prefix
        - usage: LiteLLM Usage block, containing anthropic caching information

    Returns:
        Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
    """
    ## check if model mapped, else use default pricing
    try:
        model_info = get_model_info(model=model, custom_llm_provider="fireworks_ai")
    except Exception:
        base_model: Final = get_base_model_for_pricing(model_name=model)

        ## GET MODEL INFO
        model_info = get_model_info(model=base_model, custom_llm_provider="fireworks_ai")

    ## CALCULATE INPUT COST
    prompt_tokens_details: Final = usage.prompt_tokens_details
    cached_tokens: Final[int] = (
        prompt_tokens_details.cached_tokens
        if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
        else 0
    )
    input_cost_per_token: Final[float] = model_info["input_cost_per_token"] or 0.0
    cache_read_input_token_cost: Final = model_info.get("cache_read_input_token_cost")
    cache_read_cost_per_token: Final[float] = (
        cache_read_input_token_cost if cache_read_input_token_cost is not None else input_cost_per_token
    )
    non_cached_prompt_tokens: Final[int] = max(usage.prompt_tokens - cached_tokens, 0)

    prompt_cost: float = non_cached_prompt_tokens * input_cost_per_token + cached_tokens * cache_read_cost_per_token

    ## CALCULATE OUTPUT COST
    output_cost_per_token: Final[float] = model_info["output_cost_per_token"] or 0.0
    completion_cost: Final[float] = usage.completion_tokens * output_cost_per_token

    return prompt_cost, completion_cost