"""SQLAlchemy 2.0 declarative models.

A case is a proband: a VCF plus the phenotype terms observed in that patient. One case -> many
jobs (pipeline runs); one job -> many variants; one variant -> one prediction and one decision.
"""
import enum
import uuid
from datetime import datetime

from sqlalchemy import (
    Boolean,
    DateTime,
    Enum,
    Float,
    ForeignKey,
    Integer,
    String,
    Text,
    UniqueConstraint,
    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 DecisionState(str, enum.Enum):
    shortlisted = "shortlisted"
    dismissed = "dismissed"


class Case(Base):
    __tablename__ = "cases"
    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="case")
    phenotypes: Mapped[list["CasePhenotype"]] = relationship(
        back_populates="case", cascade="all, delete-orphan", order_by="CasePhenotype.hpo_id"
    )


class CasePhenotype(Base):
    """An HPO term observed in this patient; the ranking matches these against gene annotations."""

    __tablename__ = "case_phenotypes"
    __table_args__ = (UniqueConstraint("case_id", "hpo_id", name="uq_case_phenotypes_case_term"),)
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    case_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("cases.id", ondelete="CASCADE"), index=True)
    hpo_id: Mapped[str] = mapped_column(String(20))
    label: Mapped[str] = mapped_column(String(200))
    case: Mapped[Case] = relationship(back_populates="phenotypes")


class HpoTerm(Base):
    """One HPO term and its information content (scripts/load-hpo.py); read-only reference data.

    ic is -ln(fraction of annotated genes carrying the term), so a term shared by nearly every
    gene is worth almost nothing and a near-pathognomonic one is worth a lot.
    """

    __tablename__ = "hpo_terms"
    hpo_id: Mapped[str] = mapped_column(String(20), primary_key=True)
    name: Mapped[str] = mapped_column(String(200))
    ic: Mapped[float] = mapped_column(Float)


class GenePhenotype(Base):
    """HPO's gene-to-phenotype annotations (scripts/load-hpo.py); read-only reference data.

    Propagated up the ontology at load time: a gene annotated with a term also carries that term's
    ancestors, so a case term matches a gene annotated with anything more specific.
    """

    __tablename__ = "gene_phenotypes"
    __table_args__ = (UniqueConstraint("gene_symbol", "hpo_id", name="uq_gene_phenotypes_gene_term"),)
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    gene_symbol: Mapped[str] = mapped_column(String(60), index=True)
    hpo_id: Mapped[str] = mapped_column(String(20), index=True)
    hpo_name: Mapped[str] = mapped_column(String(200))


class Job(Base):
    __tablename__ = "jobs"
    id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
    case_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("cases.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))
    # What this run's annotation actually produced. The ranking refuses to score a line of
    # evidence the run never looked up, so it has to be recorded rather than assumed.
    has_frequencies: Mapped[bool] = mapped_column(Boolean, default=False, server_default="false")
    has_effect_scores: Mapped[bool] = mapped_column(Boolean, default=False, server_default="false")
    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))
    case: Mapped[Case] = 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)
    decision: Mapped["VariantDecision | 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")


class VariantDecision(Base):
    """A reviewer's triage call. No authentication here, so decisions are shared by all visitors."""

    __tablename__ = "variant_decisions"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    variant_id: Mapped[int] = mapped_column(ForeignKey("variants.id", ondelete="CASCADE"), unique=True)
    state: Mapped[DecisionState] = mapped_column(Enum(DecisionState))
    reason: Mapped[str | None] = mapped_column(String(120))
    note: Mapped[str | None] = mapped_column(Text)
    decided_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    variant: Mapped[Variant] = relationship(back_populates="decision")