from datetime import datetime
from decimal import Decimal
from typing import Any, Dict, List, Optional, Union
from uuid import UUID

from pydantic import BaseModel


class LangSmithEvent(BaseModel):
    name: str
    time: Union[str, datetime]


class LangSmithRun(BaseModel):
    id: Union[str, UUID]
    name: str
    start_time: Optional[Union[str, datetime]] = None
    run_type: str
    end_time: Optional[Union[str, datetime]] = None
    extra: Optional[Dict] = None
    error: Optional[str] = None
    serialized: Optional[Union[str, Dict]] = None
    events: Optional[List[Union[LangSmithEvent, Dict]]] = []
    inputs: Dict = {}
    outputs: Optional[Dict] = None
    parent_run_id: Optional[Union[str, UUID]] = None
    tags: Optional[List[str]] = None
    attachments: Optional[Dict] = None
    session_id: Optional[Union[str, UUID]] = None
    child_run_ids: Optional[List[Union[str, UUID]]] = None
    child_runs: Optional[List['LangSmithRun']] = None
    feedback_stats: Optional[Dict] = None
    app_path: Optional[str] = None # prefix can be identifier
    manifest_id: Optional[Union[str, UUID]] = None
    status: Optional[str] = None
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    total_tokens: Optional[int] = None
    first_token_time: Optional[Union[str, datetime]] = None
    total_cost: Optional[Union[str, Decimal]] = None
    prompt_cost: Optional[Union[str, Decimal]] = None
    completion_cost: Optional[Union[str, Decimal]] = None
    parent_run_ids: Optional[List[Union[str, UUID]]] = None
    trace_id: Optional[Union[str, UUID]] = None
    dotted_order: Optional[str] = None
    in_dataset: Optional[bool] = None
    depth: Optional[int] = None

class LangSmithTrace(BaseModel):
    trace_id: str
    trace_name: str
    project_name: str
    export_time: str
    total_runs: int
    runs: List[LangSmithRun]
    run_depth: Optional[int] = None

class LangFuseObservation(BaseModel):
    id: str
    traceId: str
    type: str
    name: str
    startTime: Union[str, datetime]
    endTime: Union[str, datetime]
    completionStartTime: Optional[Union[str, datetime]]
    model: Optional[str]
    version: Optional[str]
    metadata: Dict
    input: Optional[Any]
    output: Optional[Any]
    usage: Optional[Dict]
    level: str
    statusMessage: Optional[str]
    parentObservationId: Optional[str]
    promptId: Optional[str]
    usageDetails: Dict
    costDetails: Dict
    environment: str
    promptName: Optional[str]
    promptVersion: Optional[str]
    modelId: Optional[str]
    inputPrice: float
    outputPrice: float
    totalPrice: float
    calculatedInputCost: Optional[float]
    calculatedOutputCost: Optional[float]
    calculatedTotalCost: float
    latency: float
    timeToFirstToken: Optional[float]
    updatedAt: Union[str, datetime]
    unit: str
    projectId: str
    createdAt: Union[str, datetime]
    totalTokens: int
    completionTokens: int
    promptTokens: int
    depth: Optional[int] = None

class LangFuseTrace(BaseModel):
    id: str
    timestamp: Union[str, datetime]
    name: str
    sessionId: Optional[str]
    metadata: Dict
    tags: List[str]
    public: bool
    environment: str
    htmlPath: str
    latency: float
    totalCost: float
    observations: List[LangFuseObservation]
    scores: List
    updatedAt: Union[str, datetime]
    projectId: str
    createdAt: Union[str, datetime]
    bookmarked: bool
    input: Optional[Any]
    output: Optional[Any]
    release: Optional[str]
    version: Optional[str]
    userId: Optional[str]
    externalId: Optional[str]
    observation_depth: Optional[int] = None

class LangFuseSession(BaseModel):
    session_id: str
    session_name: str
    project_name: str
    export_timestamp: str
    total_traces: int
    traces: List[LangFuseTrace]