import os
from functools import lru_cache
from importlib.resources import as_file, files

import pandas as pd
import requests

from rwe.clients.hgnc import get_gene_info


_CONSTRAINT_RESOURCE = "gnomad.v4.1.constraint_metrics.tsv"
_PLOF_SUMMARY_RESOURCE = "Combined_plof_summary_Gnomad_GeneBass_Guo_2026.csv"
_CONSTRAINT_URL = "https://storage.googleapis.com/gcp-public-data--gnomad/release/4.1/constraint/gnomad.v4.1.constraint_metrics.tsv"
_HGNC_RESOURCE = "gene_with_protein_product.txt"


def _normalize_gene(gene: str) -> str:
    return gene.strip().upper()


def _default_constraint_path() -> str:
    return os.path.expanduser(f"~/.cache/rwe/gnomad/{_CONSTRAINT_RESOURCE}")


def _default_hgnc_path() -> str:
    return os.path.abspath(_HGNC_RESOURCE)


def _download_constraint_table(constraint_path: str) -> None:
    os.makedirs(os.path.dirname(constraint_path), exist_ok=True)
    resp = requests.get(_CONSTRAINT_URL, stream=True, timeout=120)
    resp.raise_for_status()
    with open(constraint_path, "wb") as f:
        for chunk in resp.iter_content(chunk_size=1024 * 1024):
            if chunk:
                f.write(chunk)
    return

@lru_cache(maxsize=4)
def _load_gnomad_constraint_table(constraint_path: str) -> pd.DataFrame:
    if not os.path.exists(constraint_path):
        _download_constraint_table(constraint_path)

    df = pd.read_csv(
        constraint_path,
        sep="\t",
        usecols=["gene", "lof.oe_ci.upper", "canonical", "mane_select", "level"],
        low_memory=False,
    )
    df["gene"] = df["gene"].astype(str).str.upper()
    return df


@lru_cache(maxsize=1)
def _load_gnomad_plof_summary() -> pd.DataFrame:
    resource = files("rwe") / "assets" / _PLOF_SUMMARY_RESOURCE
    with as_file(resource) as path:
        df = pd.read_csv(
            path,
            usecols=["Gene", "Homozygote_Count", "Heterozygote_Count", "source"],
        )
    df["Gene"] = df["Gene"].astype(str).str.upper()
    df["source"] = df["source"].astype(str).str.lower()
    return df.loc[df["source"] == "gnomad"].copy()


def _select_best_constraint_row(gene_df: pd.DataFrame) -> pd.Series:
    ranked = gene_df.assign(
        _gene_level=gene_df["level"].isna(),
        _canonical=gene_df["canonical"].astype(str).str.lower() == "true",
        _mane_select=gene_df["mane_select"].astype(str).str.lower() == "true",
    ).sort_values(
        by=["_gene_level", "_canonical", "_mane_select"],
        ascending=[False, False, False],
    )
    return ranked.iloc[0]


def get_gnomad_constraint_info(
    gene: str,
    constraint_path: str | None = None,
    hgnc_path: str | None = None,
) -> dict:
    """
    Return gnomAD constraint and carrier count information for a gene.

    The lookup is performed against:
    - a downloaded local copy of `gnomad.v4.1.constraint_metrics.tsv` for LOEUF (`lof.oe_ci.upper`)
    - `Combined_plof_summary_Gnomad_GeneBass_Guo_2026.csv` for gnomAD carrier counts
    """
    hgnc_path = hgnc_path or _default_hgnc_path()
    approved_symbol, _ensembl_id, _entrez_id = get_gene_info(gene, hgnc_path)
    normalized_gene = _normalize_gene(approved_symbol)
    constraint_path = constraint_path or _default_constraint_path()

    constraint_df = _load_gnomad_constraint_table(constraint_path)
    constraint_hits = constraint_df.loc[constraint_df["gene"] == normalized_gene]
    if constraint_hits.empty:
        loeuf = None
    else:
        loeuf = _select_best_constraint_row(constraint_hits)["lof.oe_ci.upper"]
        loeuf = None if pd.isna(loeuf) else float(loeuf)

    plof_df = _load_gnomad_plof_summary()
    plof_hits = plof_df.loc[plof_df["Gene"] == normalized_gene]
    if plof_hits.empty:
        nhomz = None
        nhetz = None
    else:
        row = plof_hits.iloc[0]
        nhomz = None if pd.isna(row["Homozygote_Count"]) else int(row["Homozygote_Count"])
        nhetz = None if pd.isna(row["Heterozygote_Count"]) else int(row["Heterozygote_Count"])

    return {
        "input_gene": gene,
        "gene": normalized_gene,
        "loeuf": loeuf,
        "nhetz": nhetz,
        "nhomz": nhomz,
    }

if __name__ == "__main__":
    gene = "ANGPTL3"
    info = get_gnomad_constraint_info(
        gene, 
        constraint_path="/home/dbanerjee/work/rwe/data/rwe_reports/gnomad.v4.1.constraint_metrics.tsv",
        hgnc_path="/home/dbanerjee/work/rwe/data/rwe_reports/gene_with_protein_product.txt"
        )
    print(info)
