#!/usr/bin/env python
"""Load a VEP --tab output into the rarelens Postgres schema and mark the job succeeded."""
import argparse
import json
import sys

import pandas as pd
from sqlalchemy import create_engine, text


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--tsv", required=True)
    p.add_argument("--job-id", required=True)
    p.add_argument("--db-url", required=True)
    a = p.parse_args()

    df = pd.read_csv(a.tsv, sep="\t", comment="#", header=None, dtype=str)
    with open(a.tsv) as fh:
        header = next(l for l in fh if l.startswith("#Uploaded_variation")).lstrip("#").rstrip().split("\t")
    df.columns = header

    chrom_pos = df["Location"].str.split(":", expand=True)
    rows = []
    for i, r in df.iterrows():
        ref, _, alt = r["Uploaded_variation"].partition("/") if "/" in r["Uploaded_variation"] else ("", "", r["Allele"])
        rows.append({
            "job_id": a.job_id,
            "chrom": chrom_pos.iloc[i, 0],
            "pos": int(chrom_pos.iloc[i, 1].split("-")[0]),
            "ref": ref or "-",
            "alt": r["Allele"],
            "gene": r.get("SYMBOL") if r.get("SYMBOL") != "-" else None,
            "consequence": r["Consequence"],
            "impact": r["IMPACT"],
            "hgvsc": None if r.get("HGVSc") == "-" else r.get("HGVSc"),
            "hgvsp": None if r.get("HGVSp") == "-" else r.get("HGVSp"),
            "gnomad_af": None if r.get("gnomADe_AF", "-") == "-" else float(r["gnomADe_AF"]),
            "clinvar_sig": None if r.get("CLIN_SIG", "-") == "-" else r["CLIN_SIG"],
            "annotations": json.dumps({k: v for k, v in r.items() if v != "-"}),
        })

    if a.db_url == "none":
        print(f"{len(rows)} variants parsed (dry run, no DB)")
        return

    engine = create_engine(a.db_url.replace("+asyncpg", "+psycopg"))
    with engine.begin() as conn:
        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)
        conn.execute(text("UPDATE jobs SET status='succeeded', finished_at=now() WHERE id=:id"),
                     {"id": a.job_id})
    print(f"loaded {len(rows)} variants for job {a.job_id}", file=sys.stderr)


if __name__ == "__main__":
    main()