#!/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


def load(engine: Engine, job_id: str, rows: list[dict], vep: str | None) -> 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,
                                finished_at = now()
                WHERE id = :id
            """),
            {"id": job_id, "vep": vep},
        )


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()

    rows = to_rows(read_vep_tab(a.tsv), job_id=a.job_id)
    if a.dry_run:
        print(f"{len(rows)} variants parsed (dry run, no DB)")
        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))
    print(f"loaded {len(rows)} variants for job {a.job_id}", file=sys.stderr)


if __name__ == "__main__":
    main()