"""SQLAlchemy 2.0 declarative models.

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

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(Text)   # "&"-joined VEP terms
    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(Text)   # ","-joined co-located ClinVar terms
    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")