from enum import Enum
from pydantic import BaseModel, Field
from typing import Optional, Dict, Any

class KeyStatus(str, Enum):
    ACTIVE = "active"
    COOLDOWN = "cooldown"
    DISABLED = "disabled"

class APIKey(BaseModel):
    id: str = Field(..., description="Unique identifier for the key")
    key: str = Field(..., description="The actual API Key string")
    provider: str = Field(..., description="Provider identifier (e.g., google, zhipu)")
    api_base: Optional[str] = Field(None, description="Base URL for the API")
    model_name: str = Field(..., description="The physical model name")
    params: Dict[str, Any] = Field(default_factory=dict, description="Default model parameters (e.g., temperature)")
    extra_config: Dict[str, Any] = Field(default_factory=dict, description="Provider specific configuration")
    enabled: bool = Field(default=True, description="Whether the key is enabled in configuration")
    status: KeyStatus = Field(default=KeyStatus.ACTIVE, description="Current status of the key")
    rpm_limit: int = Field(..., description="Requests Per Minute limit")
    tpm_limit: int = Field(..., description="Tokens Per Minute limit")
    rpd_limit: int = Field(default=0, description="Requests Per Day limit (0 = unlimited)")
    max_concurrency: int = Field(..., description="Maximum concurrent requests allowed")
    config_id: Optional[str] = Field(None, description="Config ID this key belongs to")
    owner: Optional[str] = Field(None, description="Owner of this key")
    
    # Runtime stats (Not loaded from config usually, but part of the object in memory)
    current_concurrency: int = Field(default=0, description="Current concurrent requests")
    current_rpm: int = Field(default=0, description="Current RPM usage")
    current_tpm: int = Field(default=0, description="Current TPM usage")
    current_rpd: int = Field(default=0, description="Current RPD usage")
    rate_limit_penalty: int = Field(default=0, description="429 demotion order. 0 means never rate limited.")
    endpoint_idx: Optional[int] = Field(default=None, description="Endpoint index for local llama load balancing")

    class Config:
        use_enum_values = True