from pydantic import BaseModel
from typing import Dict, Optional


class ExperimentSummary(BaseModel):
    experiment_id: str
    name: str
    lifecycle_stage: str
    run_count: int


class RunSummary(BaseModel):
    run_id: str
    run_name: Optional[str] = None
    experiment_id: str
    status: str
    start_time: Optional[int] = None
    end_time: Optional[int] = None


class RunDetail(BaseModel):
    run_id: str
    run_name: Optional[str] = None
    experiment_id: str
    status: str
    start_time: Optional[int] = None
    end_time: Optional[int] = None
    params: Dict[str, str] = {}
    metrics: Dict[str, float] = {}
    tags: Dict[str, str] = {}


class MLflowLink(BaseModel):
    url: str


class ServeRequest(BaseModel):
    tracking_uri: Optional[str] = None
    model_uri: str  # e.g. "models:/model_name/1"
    port: int = 5001


class ServeStatus(BaseModel):
    id: str
    model_uri: str
    port: int
    pid: int
    status: str  # "running" or "stopped"