from typing import Dict, List, Union, Any, Optional

from pydantic import BaseModel, Field

from ...router import ModelGroupInfo


class ModelGroupInfoProxy(ModelGroupInfo):
    is_public_model_group: bool = Field(default=False)
    health_status: Optional[str] = Field(default=None)
    health_response_time: Optional[float] = Field(default=None)
    health_checked_at: Optional[str] = Field(default=None)


class UpdateUsefulLinksRequest(BaseModel):
    # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
    # New format: { "displayName": { "url": "...", "index": 0 } }
    # Old format: { "displayName": "url" } (for backward compatibility)
    useful_links: Dict[str, Union[str, Dict[str, Any]]]


class NewModelGroupRequest(BaseModel):
    access_group: str  # The access group name (e.g., "production-models")
    model_names: Optional[List[str]] = (
        None  # Existing model groups to include - tags ALL deployments for each name
    )
    model_ids: Optional[List[str]] = (
        None  # Specific deployment IDs to tag (more precise than model_names)
    )


class NewModelGroupResponse(BaseModel):
    access_group: str
    model_names: Optional[List[str]] = None
    model_ids: Optional[List[str]] = None
    models_updated: int  # Number of models updated


class UpdateModelGroupRequest(BaseModel):
    model_names: Optional[List[str]] = (
        None  # Updated list of model groups to include - tags ALL deployments for each name
    )
    model_ids: Optional[List[str]] = (
        None  # Specific deployment IDs to tag (more precise than model_names)
    )


class DeleteModelGroupResponse(BaseModel):
    access_group: str
    models_updated: int  # Number of deployments where the access group was removed
    message: str


class AccessGroupInfo(BaseModel):
    access_group: str
    model_names: List[str]  # List of model names in this access group
    deployment_count: int  # Total number of deployments with this access group


class ListAccessGroupsResponse(BaseModel):
    access_groups: List[AccessGroupInfo]
