from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field

# ==============================================================================
# region Models
# ===============================================================================


class JobStatus(str, Enum):
    """Lifecycle states of a long-running job."""

    PENDING = "pending"
    STARTED = "started"
    RUNNING = "running"
    FINISHED = "finished"
    CANCELLED = "cancelled"
    FAILED = "failed"
    TIMED_OUT = "timed_out"


class JobEnvelope[TData: BaseModel | None](BaseModel):
    """Wrapper around an incoming job request payload."""

    job_id: str
    ts: datetime = Field(default_factory=datetime.now)
    # data: TData | None = Field(default=None)
    data: TData


class JobStatusUpdate(BaseModel):
    """Published on the status subject with every heartbeat or explicit update."""

    ts: datetime = Field(default_factory=datetime.now)
    job_id: str
    status: JobStatus
    message: str | None = None
    error: str | None = None
    progress: float | None = None
    estimated_time_remaining_seconds: float | None = None


class RpcRequestEnvelope[TReq: BaseModel | None](BaseModel):
    """Wrapper around an incoming RPC call."""

    rpc_id: str
    ts: datetime = Field(default_factory=datetime.now)
    data: TReq | None = None


class RpcResponseEnvelope[TRes: BaseModel | None](BaseModel):
    """Wrapper around an RPC response."""

    rpc_id: str
    ts: datetime = Field(default_factory=datetime.now)
    data: TRes | None = None
    error: str | None = None
