from datetime import datetime

from sqlalchemy import Boolean, CheckConstraint, DateTime, Float, ForeignKey, Index, Integer, String, Text, text
from sqlalchemy.orm import Mapped, mapped_column

from app.models.base import Base, TimestampMixin, SoftDeleteMixin


class ApiGenerationTask(Base, TimestampMixin, SoftDeleteMixin):
    """对外开放 API 的生成任务表。

    该表设计满足 ProviderGenerationRecordLike 协议，
    使现有的 Volcano Ark SDK 封装函数可以直接复用。
    """

    __tablename__ = "api_generation_tasks"
    __table_args__ = (
        # 幂等键唯一索引
        Index(
            "uq_api_generation_tasks_key_idempotency",
            "api_key_id",
            "external_idempotency_key",
            unique=True,
            postgresql_where=text("deleted_at IS NULL AND external_idempotency_key IS NOT NULL"),
        ),
        # 视频轮询调度索引
        Index(
            "idx_api_generation_tasks_next_poll_at",
            "next_poll_at",
            postgresql_where=text(
                "deleted_at IS NULL "
                "AND status = 'generating' "
                "AND gen_type = 'video' "
                "AND next_poll_at IS NOT NULL"
            ),
        ),
        Index("idx_api_generation_tasks_api_key_created", "api_key_id", "created_at"),
        Index("idx_api_generation_tasks_provider_task_id", "provider_task_id"),
        Index("idx_api_generation_tasks_status", "status"),
        CheckConstraint("generation_count BETWEEN 1 AND 5", name="ck_api_generation_tasks_generation_count"),
    )

    id: Mapped[str] = mapped_column(String(32), primary_key=True)
    api_key_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("api_keys.id", ondelete="CASCADE"), nullable=False, index=True
    )
    external_idempotency_key: Mapped[str | None] = mapped_column(String(64), nullable=True)

    # === ProviderGenerationRecordLike 协议字段 ===
    original_prompt: Mapped[str] = mapped_column(Text, nullable=False)
    optimized_prompt: Mapped[str | None] = mapped_column(Text, nullable=True)
    gen_type: Mapped[str] = mapped_column(String(16), default="video", nullable=False)
    duration: Mapped[int | None] = mapped_column(Integer, nullable=True)
    aspect_ratio: Mapped[str | None] = mapped_column(String(8), nullable=True)
    resolution: Mapped[str | None] = mapped_column(String(8), nullable=True)
    provider_generation_resolution: Mapped[str | None] = mapped_column(String(16), nullable=True)
    image_size: Mapped[str | None] = mapped_column(String(16), nullable=True)
    image_proportion: Mapped[str | None] = mapped_column(String(8), nullable=True)
    image_px: Mapped[str | None] = mapped_column(String(16), nullable=True)
    generation_count: Mapped[int] = mapped_column(Integer, default=1, server_default="1", nullable=False)
    engine_id: Mapped[str | None] = mapped_column(String(32), nullable=True)
    model_name: Mapped[str] = mapped_column(String(128), nullable=False, comment="模型名称")
    media_references: Mapped[str | None] = mapped_column(Text, nullable=True, comment="用户原始上传的媒体URL")
    local_media_json: Mapped[str | None] = mapped_column(Text, nullable=True, comment="下载到本地的媒体文件路径JSON")
    engine_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)

    # === 请求参数快照 ===
    request_params_json: Mapped[str | None] = mapped_column(Text, nullable=True, comment="完整原始请求参数")

    # === 流水线状态（镜像 ChatGenerationTask） ===
    status: Mapped[str] = mapped_column(String(32), default="pending", nullable=False)
    pipeline_stage: Mapped[str | None] = mapped_column(String(32), nullable=True)
    generation_attempt_no: Mapped[int] = mapped_column(Integer, default=1, server_default="1", nullable=False)
    resource_generation_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    deadline_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    # === 供应商交互 ===
    provider_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
    remote_result_url: Mapped[str | None] = mapped_column(Text, nullable=True)
    provider_response_json: Mapped[str | None] = mapped_column(Text, nullable=True)

    # === 结果 ===
    image_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
    video_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
    video_cover_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    generated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    # === 超分 ===
    video_upscale_enabled_snapshot: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=False, server_default="false"
    )
    video_upscale_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)

    # === 配额消耗 ===
    credits_cost: Mapped[float] = mapped_column(Float, default=0.0, server_default="0.0")
    video_tokens_used: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
    image_tokens_used: Mapped[int] = mapped_column(Integer, default=0, server_default="0")

    # === 轮询控制 ===
    next_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    poll_interval_seconds: Mapped[int] = mapped_column(Integer, default=30, server_default="30")
    poll_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
    last_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    # === Celery 执行租约（镜像 ChatGenerationTask） ===
    provider_create_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
    provider_create_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    provider_create_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    poll_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    poll_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
    poll_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    poll_error_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
    download_celery_task_id: Mapped[str | None] = mapped_column(String(160), nullable=True)
    download_enqueued_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    download_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    download_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
    download_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    download_next_retry_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    download_attempt_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
    download_last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
    download_storage_date_dir: Mapped[str | None] = mapped_column(String(16), nullable=True)

    # === 存储 ===
    local_path: Mapped[str | None] = mapped_column(Text, nullable=True)