#!/usr/bin/env python3
"""Load VEP --tab output into the rarelens Postgres schema and mark the job succeeded.

Variant identity (chrom/pos/ref/alt) comes from the VCF ID, which NORMALISE sets to
CHROM_POS_REF_ALT: VEP's own Location/Allele columns trim indel alleles and shift positions.
The database URL is read from $DATABASE_URL so it never appears on a command line or in .command.sh.
Loading replaces any rows already stored for the job, so a retried task cannot duplicate variants.
"""
import argparse
import json
import os
import re
import sys
from pathlib import Path

import pandas as pd
from sqlalchemy import Engine, create_engine, text
from sqlalchemy.engine import make_url

INSERT_CHUNK = 5000
VEP_VERSION = re.compile(r"^## ENSEMBL VARIANT EFFECT PREDICTOR v(\S+)")


def engine_for(url: str) -> Engine:
    """Accept the API's asyncpg URL (or a plain postgresql:// one) and use psycopg."""
    parsed = make_url(url).set(drivername="postgresql+psycopg")
    query = dict(parsed.query)
    if "ssl" in query:  # asyncpg's spelling; libpq (psycopg) wants sslmode
        query["sslmode"] = query.pop("ssl")
    return create_engine(parsed.set(query=query))


def read_vep_tab(path: str | Path) -> pd.DataFrame:
    """Read VEP --tab output as strings, keeping "-" (VEP's missing marker) verbatim."""
    with open(path) as fh:
        for n, line in enumerate(fh):
            if line.startswith("#Uploaded_variation"):
                break
        else:
            raise ValueError(f"{path}: no #Uploaded_variation header; is this VEP --tab output?")
    df = pd.read_csv(path, sep="\t", skiprows=n, dtype=str, keep_default_na=False)
    return df.rename(columns={"#Uploaded_variation": "Uploaded_variation"})


def vep_version(path: str | Path) -> str | None:
    with open(path) as fh:
        for line in fh:
            if not line.startswith("##"):
                return None
            if m := VEP_VERSION.match(line):
                return m.group(1)
    return None


def parse_variant_id(uid: str) -> tuple[str, int, str, str]:
    # rsplit: contig names may contain "_" (chrUn_KI270742v1); positions and alleles never do.
    parts = uid.rsplit("_", 3)
    if len(parts) != 4 or not parts[1].isdigit():
        raise ValueError(
            f"unexpected Uploaded_variation {uid!r}; NORMALISE must set VCF IDs to CHROM_POS_REF_ALT"
        )
    chrom, pos, ref, alt = parts
    return chrom, int(pos), ref, alt


def _value(v: str | None) -> str | None:
    return None if v in (None, "", "-") else v


def to_rows(df: pd.DataFrame, job_id: str) -> list[dict]:
    rows = []
    for rec in df.to_dict("records"):
        chrom, pos, ref, alt = parse_variant_id(rec["Uploaded_variation"])
        af = _value(rec.get("gnomADe_AF"))
        rows.append({
            "job_id": job_id,
            "chrom": chrom,
            "pos": pos,
            "ref": ref,
            "alt": alt,
            "gene": _value(rec.get("SYMBOL")),
            "consequence": _value(rec.get("Consequence")),
            "impact": _value(rec.get("IMPACT")),
            "hgvsc": _value(rec.get("HGVSc")),
            "hgvsp": _value(rec.get("HGVSp")),
            "gnomad_af": float(af) if af is not None else None,
            "clinvar_sig": _value(rec.get("CLIN_SIG")),
            "annotations": json.dumps({k: v for k, v in rec.items() if _value(v) is not None}),
        })
    return rows


FREQUENCY_COLUMNS = ("gnomADe_AF", "gnomADg_AF", "AF")
EFFECT_COLUMNS = ("CADD_PHRED", "am_pathogenicity")


def sources(df: pd.DataFrame) -> dict[str, bool]:
    """Which lines of evidence this run looked up at all.

    Column *presence*, not a non-empty value: a variant absent from gnomAD is strong evidence of
    rarity, but only when gnomAD was consulted. VEP's database mode emits no frequency column at
    all, and the API must be able to tell the two apart instead of scoring both as maximally rare.
    """
    return {
        "has_frequencies": any(c in df.columns for c in FREQUENCY_COLUMNS),
        "has_effect_scores": any(c in df.columns for c in EFFECT_COLUMNS),
    }


def load(engine: Engine, job_id: str, rows: list[dict], vep: str | None,
         evidence: dict[str, bool]) -> None:
    with engine.begin() as conn:
        conn.execute(text("DELETE FROM variants WHERE job_id = :id"), {"id": job_id})
        for start in range(0, len(rows), INSERT_CHUNK):
            conn.execute(
                text("""
                    INSERT INTO variants (job_id, chrom, pos, ref, alt, gene, consequence, impact,
                                          hgvsc, hgvsp, gnomad_af, clinvar_sig, annotations)
                    VALUES (:job_id, :chrom, :pos, :ref, :alt, :gene, :consequence, :impact,
                            :hgvsc, :hgvsp, :gnomad_af, :clinvar_sig, CAST(:annotations AS jsonb))
                """),
                rows[start:start + INSERT_CHUNK],
            )
        conn.execute(
            text("""
                UPDATE jobs SET status = 'succeeded', vep_version = :vep, log = NULL,
                                has_frequencies = :has_frequencies,
                                has_effect_scores = :has_effect_scores,
                                finished_at = now()
                WHERE id = :id
            """),
            {"id": job_id, "vep": vep, **evidence},
        )


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--tsv", required=True)
    p.add_argument("--job-id", required=True)
    p.add_argument("--dry-run", action="store_true", help="parse only; do not touch the database")
    a = p.parse_args()

    df = read_vep_tab(a.tsv)
    rows = to_rows(df, job_id=a.job_id)
    evidence = sources(df)
    if a.dry_run:
        print(f"{len(rows)} variants parsed (dry run, no DB); {evidence}")
        return
    url = os.environ.get("DATABASE_URL")
    if not url:
        sys.exit("DATABASE_URL is not set")
    load(engine_for(url), a.job_id, rows, vep=vep_version(a.tsv), evidence=evidence)
    print(f"loaded {len(rows)} variants for job {a.job_id}", file=sys.stderr)


if __name__ == "__main__":
    main()