"""Assemble a case's ranked candidates: load the latest results, rank them, attach decisions.

Everything for one job is loaded at once, which is fine for a gene panel or a chromosome — the
size of case this demo handles. A whole genome would need the narrowing pushed into SQL.
"""
import uuid
from dataclasses import dataclass

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.models import Case, GenePhenotype, Job, JobStatus, Variant
from app.services import triage

EMPTY_FUNNEL = triage.Funnel(total=0, rare=0, candidates=0, phenotype_matched=0)


@dataclass(frozen=True)
class CaseView:
    case: Case
    job: Job | None
    funnel: triage.Funnel
    candidates: list[triage.Candidate]

    @property
    def labels(self) -> dict[str, str]:
        return {p.hpo_id: p.label for p in self.case.phenotypes}


async def get_case(session: AsyncSession, case_id: uuid.UUID) -> Case | None:
    return await session.scalar(
        select(Case).where(Case.id == case_id).options(selectinload(Case.phenotypes))
    )


async def latest_job(
    session: AsyncSession, case_id: uuid.UUID, *, status: JobStatus | None = None
) -> Job | None:
    stmt = select(Job).where(Job.case_id == case_id).order_by(Job.created_at.desc()).limit(1)
    if status is not None:
        stmt = stmt.where(Job.status == status)
    return await session.scalar(stmt)


async def gene_terms_for(session: AsyncSession, genes: set[str]) -> dict[str, set[str]]:
    """gene symbol -> the HPO terms annotated to it."""
    if not genes:
        return {}
    rows = await session.execute(
        select(GenePhenotype.gene_symbol, GenePhenotype.hpo_id).where(
            GenePhenotype.gene_symbol.in_(genes)
        )
    )
    index: dict[str, set[str]] = {}
    for gene, hpo_id in rows:
        index.setdefault(gene, set()).add(hpo_id)
    return index


async def build(session: AsyncSession, case: Case) -> CaseView:
    job = await latest_job(session, case.id, status=JobStatus.succeeded)
    if job is None:
        return CaseView(case=case, job=None, funnel=EMPTY_FUNNEL, candidates=[])
    variants = (
        await session.scalars(
            select(Variant)
            .where(Variant.job_id == job.id)
            .options(selectinload(Variant.prediction), selectinload(Variant.decision))
        )
    ).all()
    case_terms = [p.hpo_id for p in case.phenotypes]
    gene_terms = await gene_terms_for(session, {v.gene for v in variants if v.gene})
    return CaseView(
        case=case,
        job=job,
        funnel=triage.funnel(variants, case_terms, gene_terms),
        candidates=triage.rank(variants, case_terms, gene_terms),
    )