"""SQLAlchemy 2.0 declarative models.

One sample -> many jobs; one job -> many variants; one variant -> one prediction (latest).
"""
from datetime import datetime
import enum
import uuid

from sqlalchemy import DateTime, Enum, Float, ForeignKey, Integer, String, Text, func
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


class Base(DeclarativeBase):
    pass


class JobStatus(str, enum.Enum):
    queued = "queued"
    running = "running"
    succeeded = "succeeded"
    failed = "failed"


class Sample(Base):
    __tablename__ = "samples"
    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    name: Mapped[str] = mapped_column(String(120), unique=True)
    vcf_uri: Mapped[str] = mapped_column(Text)
    assembly: Mapped[str] = mapped_column(String(10), default="GRCh38")
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    jobs: Mapped[list["Job"]] = relationship(back_populates="sample")


class Job(Base):
    __tablename__ = "jobs"
    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    sample_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("samples.id", ondelete="CASCADE"))
    status: Mapped[JobStatus] = mapped_column(Enum(JobStatus), default=JobStatus.queued)
    workflow_ref: Mapped[str | None] = mapped_column(String(200))   # Argo workflow name / nf run id
    vep_version: Mapped[str | None] = mapped_column(String(40))
    log: Mapped[str | None] = mapped_column(Text)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    sample: Mapped[Sample] = relationship(back_populates="jobs")
    variants: Mapped[list["Variant"]] = relationship(back_populates="job")


class Variant(Base):
    __tablename__ = "variants"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    job_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("jobs.id", ondelete="CASCADE"), index=True)
    chrom: Mapped[str] = mapped_column(String(10), index=True)
    pos: Mapped[int] = mapped_column(Integer, index=True)
    ref: Mapped[str] = mapped_column(Text)
    alt: Mapped[str] = mapped_column(Text)
    gene: Mapped[str | None] = mapped_column(String(60), index=True)
    consequence: Mapped[str | None] = mapped_column(String(120))
    impact: Mapped[str | None] = mapped_column(String(20))
    hgvsc: Mapped[str | None] = mapped_column(Text)
    hgvsp: Mapped[str | None] = mapped_column(Text)
    gnomad_af: Mapped[float | None] = mapped_column(Float)
    clinvar_sig: Mapped[str | None] = mapped_column(String(120))
    annotations: Mapped[dict] = mapped_column(JSONB, default=dict)   # full VEP CSQ record
    job: Mapped[Job] = relationship(back_populates="variants")
    prediction: Mapped["Prediction | None"] = relationship(back_populates="variant", uselist=False)


class Prediction(Base):
    __tablename__ = "predictions"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    variant_id: Mapped[int] = mapped_column(ForeignKey("variants.id", ondelete="CASCADE"), unique=True)
    model_name: Mapped[str] = mapped_column(String(80))
    model_version: Mapped[str] = mapped_column(String(40))
    score: Mapped[float] = mapped_column(Float)         # P(pathogenic)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    variant: Mapped[Variant] = relationship(back_populates="prediction")