"""Base wrapper functionality."""

from typing import Any, Dict, Optional


def extract_mandatum_params(kwargs: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[str, Any]]:
    """
    Extract Mandatum-specific parameters from kwargs.

    Returns:
        tuple: (provider_kwargs, mandatum_params)
    """
    mandatum_params = {}
    provider_kwargs = {}

    mandatum_keys = {
        "md_tags",
        "md_environment",
        "md_metadata",
        "md_prompt_id",
        "md_version_tag",
        "md_parent_request_id",
    }

    for key, value in kwargs.items():
        if key in mandatum_keys:
            # Remove md_ prefix
            clean_key = key[3:] if key.startswith("md_") else key
            mandatum_params[clean_key] = value
        else:
            provider_kwargs[key] = value

    return provider_kwargs, mandatum_params


def calculate_cost(provider: str, model: str, input_tokens: int, output_tokens: int) -> float:
    """
    Calculate cost based on provider pricing.

    This is a simplified version. Real pricing should be kept in sync with backend.
    """
    pricing = {
        "openai": {
            "gpt-4": {"input": 0.03 / 1000, "output": 0.06 / 1000},
            "gpt-4-turbo": {"input": 0.01 / 1000, "output": 0.03 / 1000},
            "gpt-3.5-turbo": {"input": 0.0005 / 1000, "output": 0.0015 / 1000},
        },
        "anthropic": {
            "claude-3-opus": {"input": 0.015 / 1000, "output": 0.075 / 1000},
            "claude-3-sonnet": {"input": 0.003 / 1000, "output": 0.015 / 1000},
            "claude-3-haiku": {"input": 0.00025 / 1000, "output": 0.00125 / 1000},
        },
    }

    provider_pricing = pricing.get(provider, {})
    model_pricing = None

    # Find matching model (handle model variants like gpt-4-0613)
    for model_key, prices in provider_pricing.items():
        if model.startswith(model_key):
            model_pricing = prices
            break

    if not model_pricing:
        return 0.0

    input_cost = input_tokens * model_pricing["input"]
    output_cost = output_tokens * model_pricing["output"]

    return round(input_cost + output_cost, 6)
