# -------------------------------------------------------------------------------- #
# Model Costs
# -------------------------------------------------------------------------------- #
# This file is auto-generated by scripts/generate_cost_constants.py
# DO NOT EDIT MANUALLY
# -------------------------------------------------------------------------------- #

# -------------------------------------------------------------------------------- #
# Imports
# -------------------------------------------------------------------------------- #
# Built-in imports
from typing import Dict, List, TypedDict, Optional, Union, Any
from datetime import datetime, time

# Module imports
from astral_ai.constants._models import ModelProvider, ModelId, ModelName


# -------------------------------------------------------------------------------- #
# Cost Type Definitions
# -------------------------------------------------------------------------------- #

class ModelSpecificCosts(TypedDict, total=False):
    """TypedDict for capturing model-specific cost details."""
    time_period: Optional[str]  # A datetime-parseable string or None if not applicable
    input_base_cost: float
    input_cache_hit_cost: float
    input_cache_write_cost: Optional[float]
    output_base_cost: float


class TimePeriodCosts(TypedDict):
    """TypedDict for capturing different cost periods for time-based pricing."""
    time_based: bool
    standard: ModelSpecificCosts
    discount: Optional[ModelSpecificCosts]


# -------------------------------------------------------------------------------- #
# Helper Functions for Time-Based Costs
# -------------------------------------------------------------------------------- #

def parse_time_window(time_window: Optional[str]) -> Optional[tuple[time, time]]:
    """
    Parse a time window string (e.g., 'UTC 00:30–16:30') into a tuple of start and end times.
    Returns None if the input is None or not parseable.
    """
    if time_window is None:
        return None
    
    try:
        # Example format: 'UTC 00:30–16:30'
        # Strip UTC prefix if present
        if time_window.startswith('UTC '):
            time_window = time_window[4:]
        
        # Split by the dash/hyphen
        parts = time_window.split('–')
        if len(parts) != 2:
            return None
        
        # Parse start and end times
        start_str, end_str = parts[0].strip(), parts[1].strip()
        start_time = datetime.strptime(start_str, '%H:%M').time()
        end_time = datetime.strptime(end_str, '%H:%M').time()
        
        return (start_time, end_time)
    except (ValueError, IndexError):
        return None


def is_current_time_in_window(time_window: Optional[str]) -> bool:
    """
    Check if the current time is within the specified time window.
    If time_window is None, always returns True.
    """
    if time_window is None:
        return True
    
    window = parse_time_window(time_window)
    if window is None:
        return True
    
    start_time, end_time = window
    current_time = datetime.now().time()
    
    # Handle cases where the window spans midnight
    if start_time <= end_time:
        return start_time <= current_time <= end_time
    else:
        return current_time >= start_time or current_time <= end_time


# -------------------------------------------------------------------------------- #
# Provider Cost Mappings
# -------------------------------------------------------------------------------- #

OPENAI_COSTS: Dict[ModelId, TimePeriodCosts] = {
    "gpt-4.5-preview-2025-02-27": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 75.0,
            "input_cache_hit_cost": 37.5,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 150.0,
        },
        "discount": None,
    },
    "gpt-4o-2024-05-13": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 5.0,
            "input_cache_hit_cost": 0.0,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 15.0,
        },
        "discount": None,
    },
    "gpt-4o-2024-08-06": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 2.5,
            "input_cache_hit_cost": 1.25,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 10.0,
        },
        "discount": None,
    },
    "gpt-4o-2024-11-20": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 2.5,
            "input_cache_hit_cost": 1.25,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 10.0,
        },
        "discount": None,
    },
    "gpt-4o-mini-2024-07-18": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 0.15,
            "input_cache_hit_cost": 0.075,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 0.6,
        },
        "discount": None,
    },
    "o1-2024-12-17": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 15.0,
            "input_cache_hit_cost": 7.5,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 60.0,
        },
        "discount": None,
    },
    "o1-mini-2024-09-12": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 1.1,
            "input_cache_hit_cost": 0.55,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 4.4,
        },
        "discount": None,
    },
    "o1-preview-2024-09-12": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 15.0,
            "input_cache_hit_cost": 7.5,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 60.0,
        },
        "discount": None,
    },
    "o1-pro-2025-03-19": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 150.0,
            "input_cache_hit_cost": 0.0,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 600.0,
        },
        "discount": None,
    },
    "o3-mini-2025-01-31": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 1.1,
            "input_cache_hit_cost": 0.55,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 4.4,
        },
        "discount": None,
    },
}

DEEPSEEK_COSTS: Dict[ModelId, TimePeriodCosts] = {
    "deepseek-chat": {
        "time_based": True,
        "standard": {
            "time_period": "UTC 00:30–16:30",
            "input_base_cost": 0.27,
            "input_cache_hit_cost": 0.07,
            "input_cache_write_cost": None,
            "output_base_cost": 1.1,
        },
        "discount": {
            "time_period": "UTC 16:30–00:30",
            "input_base_cost": 0.135,
            "input_cache_hit_cost": 0.035,
            "input_cache_write_cost": None,
            "output_base_cost": 0.55,
        },
    },
    "deepseek-reasoner": {
        "time_based": True,
        "standard": {
            "time_period": "UTC 00:30–16:30",
            "input_base_cost": 0.55,
            "input_cache_hit_cost": 0.14,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 2.19,
        },
        "discount": {
            "time_period": "UTC 16:30–00:30",
            "input_base_cost": 0.135,
            "input_cache_hit_cost": 0.035,
            "input_cache_write_cost": 0.0,
            "output_base_cost": 0.55,
        },
    },
}

ANTHROPIC_COSTS: Dict[ModelId, TimePeriodCosts] = {
    "claude-3-5-haiku-20241022": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 0.8,
            "input_cache_hit_cost": 0.08,
            "input_cache_write_cost": 1.0,
            "output_base_cost": 4.0,
        },
        "discount": None,
    },
    "claude-3-5-sonnet-20240620": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 3.0,
            "input_cache_hit_cost": 0.3,
            "input_cache_write_cost": 3.75,
            "output_base_cost": 15.0,
        },
        "discount": None,
    },
    "claude-3-5-sonnet-20241022": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 3.0,
            "input_cache_hit_cost": 0.3,
            "input_cache_write_cost": 3.75,
            "output_base_cost": 15.0,
        },
        "discount": None,
    },
    "claude-3-7-sonnet-20250219": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 3.0,
            "input_cache_hit_cost": 0.3,
            "input_cache_write_cost": 3.75,
            "output_base_cost": 15.0,
        },
        "discount": None,
    },
    "claude-3-haiku-20240307": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 0.25,
            "input_cache_hit_cost": 0.03,
            "input_cache_write_cost": 0.3,
            "output_base_cost": 1.25,
        },
        "discount": None,
    },
    "claude-3-opus-20240229": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 15.0,
            "input_cache_hit_cost": 1.5,
            "input_cache_write_cost": 18.75,
            "output_base_cost": 75.0,
        },
        "discount": None,
    },
    "claude-3-sonnet-20240229": {
        "time_based": False,
        "standard": {
            "time_period": None,
            "input_base_cost": 3.0,
            "input_cache_hit_cost": 0.3,
            "input_cache_write_cost": 3.75,
            "output_base_cost": 15.0,
        },
        "discount": None,
    },
}

# -------------------------------------------------------------------------------- #
# Unified Cost Mapping
# -------------------------------------------------------------------------------- #

MODEL_COST_MAPPING: Dict[ModelProvider, Dict[ModelId, TimePeriodCosts]] = {
    "openai": OPENAI_COSTS,
    "deepseek": DEEPSEEK_COSTS,
    "anthropic": ANTHROPIC_COSTS,
}
