import tiktoken
import os
from importlib import resources
from dotenv import load_dotenv

# Use importlib to get a temporary path to the installed .env file
env_path = resources.files("bilbystats.defaults").joinpath("model_costs.env")

# Load environment variables from the .env file
load_dotenv(dotenv_path=env_path)

def model_costs(input_tokens, output_tokens, model_name, ndocs):
    """
    Estimate the cost of a language model API call based on token usage.

    This function calculates the number of input and output tokens using the 
    specified model's tokenizer and computes the associated cost using 
    environment variables that specify pricing per million tokens. It returns 
    a breakdown of token counts and cost estimates.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    input_tokens : str
        The raw input text to be tokenized and used for cost estimation.
    output_tokens : str
        The raw output text to be tokenized and used for cost estimation.
    model_name : str
        The name of the language model (e.g., 'gpt-4') used to fetch the appropriate tokenizer and pricing.
    ndocs : int
        The number of documents processed (currently unused in cost calculation but reserved for future use).

    ---------------------------------------------------------------------------
    OUTPUT:
    dict : dict
        A dictionary containing:
        - 'n_input_tokens': int
              Number of tokens in the input text.
        - 'm_output_tokens': int
              Number of tokens in the output text (note: key name appears to be a typo).
        - 'input_cost': float
              Estimated input cost in currency units.
        - 'output_cost': float
              Estimated output cost in currency units.
        - 'total_cost': float
              Combined cost of input and output tokens.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    tokenizer = tiktoken.encoding_for_model(model_name)

    n_input_tokens = len(tokenizer.encode(input_tokens))
    n_output_tokens = len(tokenizer.encode(output_tokens))

    input_price_per_million = os.getenv(model_name + "_input")
    output_price_per_million = os.getenv(model_name + "_output")
    input_cost = input_price_per_million * (n_input_tokens / 1_000_000)
    output_cost = output_price_per_million * (n_output_tokens / 1_000_000)

    total_cost = input_cost + output_cost

    return {
        "n_input_tokens": n_input_tokens,
        "m_output_tokens": n_output_tokens,
        "input_cost": input_cost,
        "output_cost": output_cost,
        "total_cost": total_cost
    }