from datetime import datetime

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

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


class GenerationRecord(Base, TimestampMixin, SoftDeleteMixin):
    __tablename__ = "generation_records"

    id: Mapped[str] = mapped_column(String(32), primary_key=True)
    user_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("users.id", ondelete="CASCADE"), index=True
    )
    project_id: Mapped[str] = mapped_column(
        String(32), ForeignKey("projects.id", ondelete="CASCADE"), index=True
    )
    original_prompt: Mapped[str] = mapped_column(Text)
    optimized_prompt: Mapped[str | None] = mapped_column(Text, nullable=True)
    gen_type: Mapped[str] = mapped_column(String(16), default="video")
    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)
    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)
    image_size: Mapped[str | None] = mapped_column(String(8), nullable=True)
    image_proportion: Mapped[str | None] = mapped_column(String(8), nullable=True)
    image_px: Mapped[str | None] = mapped_column(String(10), nullable=True)

    status: Mapped[str] = mapped_column(String(32), default="prompt_optimized")
    pipeline_stage: Mapped[str | None] = mapped_column(String(48), nullable=True, index=True)
    video_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
    video_cover_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
    image_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
    media_references: Mapped[str | None] = mapped_column(Text, nullable=True)
    include_media_references: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=False, server_default="false"
    )
    video_url_expires_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    seedance_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
    credits_cost: Mapped[float] = mapped_column(Float, default=0.0)
    text_credits_cost: Mapped[float] = mapped_column(Float, default=0.0)
    text_tokens_used: Mapped[int] = mapped_column(Integer, default=0)
    video_tokens_used: Mapped[int] = mapped_column(Integer, default=0)
    image_tokens_used: Mapped[int] = mapped_column(Integer, default=0)
    generated_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
    idempotency_key: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)

    # Unified Celery generation pipeline state.
    generation_attempt_no: Mapped[int] = mapped_column(Integer, nullable=False, default=1, server_default="1")
    resource_generation_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
    deadline_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)

    engine_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
    engine_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    provider_response_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    remote_result_url: Mapped[str | None] = mapped_column(Text, nullable=True)

    provider_create_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
    provider_create_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
    provider_create_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    # retry_count is retained for API/backward compatibility. New business
    # logic uses manual_retry_count and poll_error_count separately.
    retry_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
    manual_retry_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
    poll_error_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
    poll_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
    last_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    poll_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    next_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    poll_interval_seconds: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
    poll_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
    poll_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)

    download_celery_task_id: Mapped[str | None] = mapped_column(String(160), nullable=True, index=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, index=True)
    download_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
    download_next_retry_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
    download_attempt_count: Mapped[int] = mapped_column(Integer, nullable=False, 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)

    __table_args__ = (
        Index('idx_genrec_user_status_created', 'user_id', 'status', 'created_at'),
        Index('idx_genrec_project_status', 'project_id', 'status'),
        Index(
            'idx_genrec_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"),
        ),
    )