from collections.abc import Sequence
from datetime import datetime
from enum import StrEnum
from typing import Protocol

from pydantic import BaseModel, Field


class OutputType(StrEnum):
    IMAGE = "image"
    VIDEO = "video"
    AUDIO = "audio"
    THREE_D = "3d"


class MediaRole(StrEnum):
    IMAGE = "image"
    START_IMAGE = "start_image"
    END_IMAGE = "end_image"
    VIDEO = "video"
    AUDIO = "audio"
    REF_ELEMENT = "ref_element"


class GenerationStatus(StrEnum):
    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    MODERATED = "moderated"
    CANCELED = "canceled"
    UNKNOWN_SUBMISSION = "unknown_submission"


class MediaReference(BaseModel):
    asset_id: str
    external_id: str
    role: MediaRole
    version: str
    checksum: str | None = None


class ModelCapabilitySnapshot(BaseModel):
    provider: str = "higgsfield"
    model_id: str
    display_name: str | None = None
    output_type: OutputType
    aspect_ratios: tuple[str, ...] = ()
    durations_sec: tuple[int, ...] = ()
    parameter_schema: dict[str, object] = Field(default_factory=dict)
    media_roles: tuple[MediaRole, ...] = ()
    fetched_at: datetime
    capability_hash: str


class GenerationRequest(BaseModel):
    request_id: str
    asset_id: str
    model_id: str
    capability_hash: str
    prompt: str
    prompt_template_id: str
    prompt_version: str
    aspect_ratio: str
    duration_sec: int | None = None
    references: tuple[MediaReference, ...] = ()
    parameters: dict[str, object] = Field(default_factory=dict)
    # Application intent deduplication only. The public provider POST currently
    # does not accept an idempotency key; uncertain submissions need reconciliation.
    idempotency_key: str


class CostEstimate(BaseModel):
    credits: float
    currency: str = "higgsfield_credit"
    estimated_at: datetime


class GenerationJob(BaseModel):
    request_id: str
    external_job_id: str
    model_id: str
    capability_hash: str
    status: GenerationStatus
    submitted_at: datetime
    estimated_cost: CostEstimate | None = None
    actual_cost: float | None = None
    output_urls: tuple[str, ...] = ()
    error_code: str | None = None
    error_message: str | None = None


class ReferenceElement(BaseModel):
    canonical_asset_id: str
    canonical_version: str
    element_id: str
    element_name: str
    element_category: str
    reference_asset_ids: tuple[str, ...]
    created_at: datetime


class HiggsfieldPort(Protocol):
    async def refresh_capabilities(self) -> Sequence[ModelCapabilitySnapshot]: ...

    async def estimate_cost(self, request: GenerationRequest) -> CostEstimate: ...

    async def submit(self, request: GenerationRequest) -> GenerationJob: ...

    async def get_job(self, external_job_id: str) -> GenerationJob: ...

    async def ensure_reference_element(
        self,
        *,
        canonical_asset_id: str,
        canonical_version: str,
        reference_asset_ids: Sequence[str],
    ) -> ReferenceElement: ...