"""
Pydantic schemas for Open API dataset download endpoints.
"""
from enum import Enum
from datetime import datetime
from typing import Optional
from pydantic import BaseModel, Field


class DatasetFormat(str, Enum):
    ALPACA = "alpaca"
    SHAREGPT = "sharegpt"
    JSON = "json"
    CSV = "csv"
    COCO = "coco"
    YOLO = "yolo"
    PASCAL_VOC = "pascal_voc"


TEXT_FORMATS = {DatasetFormat.ALPACA, DatasetFormat.SHAREGPT}
IMAGE_FORMATS = {DatasetFormat.JSON, DatasetFormat.CSV, DatasetFormat.COCO,
                 DatasetFormat.YOLO, DatasetFormat.PASCAL_VOC}


class DatasetDownloadRequest(BaseModel):
    format: DatasetFormat = Field(..., description="数据集格式")
    completed_only: bool = Field(default=True, description="是否只导出已完成的任务")


class DatasetDownloadResponseData(BaseModel):
    project_id: str = Field(..., description="项目ID")
    format: str = Field(..., description="导出格式")
    total_exported: int = Field(..., description="导出任务数")
    file_url: str = Field(..., description="下载链接")
    file_name: str = Field(..., description="文件名")
    file_size: Optional[int] = Field(None, description="文件大小（字节）")
    expires_at: Optional[datetime] = Field(None, description="链接过期时间")
    status: str = Field(default="completed", description="导出状态")


class DatasetDownloadResponse(BaseModel):
    code: int = Field(default=0)
    message: str = Field(default="success")
    data: DatasetDownloadResponseData