from __future__ import annotations

import uuid
from datetime import datetime

from sqlalchemy import (
    CheckConstraint,
    DateTime,
    SmallInteger,
    String,
    Text,
    func,
)
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.orm import Mapped, mapped_column

from app.db.session import Base


class JobStatus:
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    DONE = "DONE"
    FAILED = "FAILED"
    EXPIRED = "EXPIRED"


class JobType:
    DOWNLOAD = "download"
    AI_ANALYZE = "ai_analyze"
    AI_EMBED = "ai_embed"
    NOTIFY = "notify"


VALID_STATUSES = (JobStatus.PENDING, JobStatus.RUNNING, JobStatus.DONE, JobStatus.FAILED, JobStatus.EXPIRED)
VALID_TYPES = (JobType.DOWNLOAD, JobType.AI_ANALYZE, JobType.AI_EMBED, JobType.NOTIFY)

# Valid state transitions
ALLOWED_TRANSITIONS: dict[str, set[str]] = {
    JobStatus.PENDING:  {JobStatus.RUNNING},
    JobStatus.RUNNING:  {JobStatus.DONE, JobStatus.FAILED},
    JobStatus.FAILED:   {JobStatus.PENDING},  # manual retry
    JobStatus.DONE:     {JobStatus.EXPIRED},
    JobStatus.EXPIRED:  set(),
}


class Job(Base):
    __tablename__ = "jobs"
    __table_args__ = (
        CheckConstraint(
            "status IN ('PENDING','RUNNING','DONE','FAILED','EXPIRED')",
            name="ck_job_status",
        ),
        CheckConstraint(
            "entity_type IN ('media','link')",
            name="ck_job_entity_type",
        ),
    )

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
    )
    entity_id: Mapped[uuid.UUID | None] = mapped_column(UUID(as_uuid=True), index=True)
    entity_type: Mapped[str | None] = mapped_column(String(16))
    user_id: Mapped[uuid.UUID | None] = mapped_column(UUID(as_uuid=True), index=True)
    canonical_url_hash: Mapped[str | None] = mapped_column(String(64))
    job_type: Mapped[str] = mapped_column(String(32), nullable=False)
    status: Mapped[str] = mapped_column(String(16), default=JobStatus.PENDING, index=True)
    attempts: Mapped[int] = mapped_column(SmallInteger, default=0, server_default="0")
    max_attempts: Mapped[int] = mapped_column(SmallInteger, default=3, server_default="3")
    error_message: Mapped[str | None] = mapped_column(Text)
    payload: Mapped[dict] = mapped_column(JSONB, default=dict, server_default="{}")
    status_history: Mapped[list] = mapped_column(JSONB, default=list, server_default="[]")
    scheduled_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now()
    )
    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))

    def __repr__(self) -> str:
        return f"<Job id={self.id} type={self.job_type} status={self.status}>"