from datetime import datetime
from typing import Any

from pydantic import BaseModel


class ProjectInfo(BaseModel):
    id: str
    name: str
    storage_backends: dict[str, str] = {}
    wandb_api_key: str | None = None
    wandb_entity: str | None = None
    created_at: datetime
    updated_at: datetime


class ExperimentInfo(BaseModel):
    id: str
    name: str
    parent_name: str | None = None
    description: str | None = None
    created_at: datetime
    updated_at: datetime


class ArtifactInfo(BaseModel):
    id: str
    name: str
    type: str
    version: int
    path: str
    storage_name: str | None = None
    description: str = ""
    metadata: dict[str, Any] = {}
    tags: list[str] = []
    produced_by_id: str
    experiment_name: str
    created_at: datetime


class JobInfo(BaseModel):
    id: str
    name: str
    experiment_name: str
    git_sha: str
    status: str
    metadata: dict[str, Any] = {}
    started_at: datetime
    completed_at: datetime | None = None
    failure_reason: str | None = None
    tracking_url: str | None = None
    created_by: str | None = None


class JobDetailInfo(JobInfo):
    inputs: list[ArtifactInfo] = []
    outputs: list[ArtifactInfo] = []
    metrics: dict[str, float] = {}


class UniversalArtifactSourceInfo(BaseModel):
    project_name: str
    experiment_name: str
    artifact_name: str
    artifact_version: int
    path: str
    storage_name: str | None = None


class UniversalArtifactInfo(BaseModel):
    id: str
    name: str
    type: str
    artifact_id: str
    description: str = ""
    metadata: dict[str, Any] = {}
    tags: list[str] = []
    source: UniversalArtifactSourceInfo
    created_at: datetime
    updated_at: datetime


class MetricsQueryRow(BaseModel):
    experiment: str
    job_id: str
    job_name: str
    completed_at: datetime | None = None
    metadata: dict[str, Any] = {}
    metrics: dict[str, float] = {}
