from datetime import datetime
from typing import Any, Union

from pydantic import BaseModel


class FileInfo(BaseModel):
    name: str
    path: str
    size: int | None = None


class QueryRequest(BaseModel):
    session_id: str
    query: str
    files: list[Union[str, FileInfo]] = []
    user_input: dict[str, Any] = {}


class QueryResponse(BaseModel):
    session_id: str
    status: str
    mode: str
    response: dict[str, Any] = {}
    messages: list[dict[str, str]] = []


class SessionCreate(BaseModel):
    user_id: str | None = None


class SessionResponse(BaseModel):
    session_id: str
    created_at: datetime
    status: str = "active"


class CaseCreate(BaseModel):
    case_id: str
    user_id: str | None = None
    gene_family: str | None = None
    species: str | None = None


class CaseResponse(BaseModel):
    case_id: str
    user_id: str | None
    gene_family: str | None
    species: str | None
    status: str
    created_at: datetime
    results: dict[str, Any] = {}


class FeedbackRequest(BaseModel):
    case_id: str
    rating: int
    comments: str | None = None
    issues: list[str] = []


class WebSocketMessage(BaseModel):
    type: str
    session_id: str
    data: dict[str, Any] = {}


class ProgressUpdate(BaseModel):
    session_id: str
    current_agent: str | None
    status: str
    progress: float
    message: str