import os
import shutil
import zipfile
from pathlib import Path
from playwright.sync_api import Playwright, sync_playwright
from rwe.plots.clinical import manhattan
import pandas as pd

import rwe.utils.helpers as uth
from docx.shared import Inches

def extract_only_csv_rename(zip_path: str, gene: str) -> str:
    zip_path = Path(zip_path)
    gene_dir = zip_path.parent
    target_csv = gene_dir / f"{gene}_azphewas.csv"

    # Extract into a dedicated folder
    extract_dir = gene_dir / f"__extract_{zip_path.stem}"
    if extract_dir.exists():
        shutil.rmtree(extract_dir)
    extract_dir.mkdir(parents=True, exist_ok=True)

    with zipfile.ZipFile(zip_path, "r") as z:
        z.extractall(extract_dir)

    csvs = list(extract_dir.rglob("*.csv"))
    if not csvs:
        raise FileNotFoundError(f"No CSV found after extracting {zip_path}")

    csvs.sort(key=lambda p: p.stat().st_size, reverse=True)
    chosen = csvs[0]

    if target_csv.exists():
        target_csv.unlink()
    chosen.replace(target_csv)

    # Cleanup only what we created
    shutil.rmtree(extract_dir)
    if zip_path.exists():
        zip_path.unlink()

    return str(target_csv)


def run(playwright: Playwright, gene: str, save_path: str) -> None:
    browser = playwright.chromium.launch(headless=True)
    context = browser.new_context(accept_downloads=True)
    page = context.new_page()

    url = f"https://azphewas.com/geneView/6319c068-fd59-46d8-85ee-82d82482eb14/{gene}/glr/binary"
    page.goto(url, wait_until="networkidle")

    export_control = page.locator('[data-testid="export-settings-control"] button:has-text("Export")')
    export_control.click()

    inner_export_btn = page.get_by_test_id("export_button")
    inner_export_btn.wait_for(state="visible")

    plot_checkbox = page.locator('[data-testid="plot-export-type-toggler"] input[type="checkbox"]')
    plot_checkbox.set_checked(False, force=True)
    page.wait_for_timeout(300)

    os.makedirs(save_path, exist_ok=True)

    with page.expect_download(timeout=120_000) as dl_info:
        inner_export_btn.click(force=True)

    download = dl_info.value
    fname = download.suggested_filename
    out_dir = save_path
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, fname)
    download.save_as(out_path)
    print("Saved zip:", out_path)

    final_csv = extract_only_csv_rename(out_path, gene=gene)
    print("Final CSV:", final_csv)

    context.close()
    browser.close()
    return

def clean_azn_phewas(phewas_file):
    df = pd.read_csv(phewas_file, sep=",")
    df = df.loc[df["Collapsing model"]=="ptv"]
    return df

def get_azn_manhattan(gene, phewas_file):
    df = clean_azn_phewas(phewas_file)
    name_col = "Phenotype"
    p_col    = "P value"
    odds_ratio_col = "Odds ratio"
    cat_col  = "Phenotypic category"
    AZN_CHAPTER_SHORT = {
        "Chapter I Certain infectious and parasitic diseases": "Infectious",
        "Chapter II Neoplasms": "Neoplasms",
        "Chapter III Diseases of the blood and blood-forming organs and certain disorders involving the immune mechanism": "Blood/Immune",
        "Chapter IV Endocrine nutritional and metabolic diseases": "Endocrine/Metabolic",
        "Chapter V Mental and behavioural disorders": "Mental/Behavioral",
        "Chapter VI Diseases of the nervous system": "Nervous System",
        "Chapter VII Diseases of the eye and adnexa": "Eye",
        "Chapter VIII Diseases of the ear and mastoid process": "Ear/Mastoid",
        "Chapter IX Diseases of the circulatory system": "Circulatory",
        "Chapter X Diseases of the respiratory system": "Respiratory",
        "Chapter XI Diseases of the digestive system": "Digestive",
        "Chapter XII Diseases of the skin and subcutaneous tissue": "Skin",
        "Chapter XIII Diseases of the musculoskeletal system and connective tissue": "Musculoskeletal",
        "Chapter XIV Diseases of the genitourinary system": "Genitourinary",
        "Chapter XV Pregnancy childbirth and the puerperium": "Pregnancy/Childbirth",
        "Chapter XVII Congenital malformations deformations and chromosomal abnormalities": "Congenital",
        "Chapter XVIII Symptoms signs and abnormal clinical and laboratory findings not elsewhere classified": "Symptoms/Findings",
        "Chapter XXI Factors influencing health status and contact with health services": "Health Services",
    }

    df[name_col] = df[name_col].astype(str).str.rsplit("#", n=1).str[-1]
    df[cat_col] = df[cat_col].map(AZN_CHAPTER_SHORT).fillna("Other")

    fig, ax, plot_df = manhattan(
        df,
        p_col=p_col,
        category_col=cat_col,
        label_col=name_col,
        odds_ratio_col=odds_ratio_col,
        sig_p=2.8e-4,
        top_k_labels=5,
        title=f"{gene} PheWAS AstraZeneca",
    )
    return fig, plot_df

def generate_azn_clinical_report(doc, gene, phewas_dir, phewas_filename=""):
    if not phewas_filename:
        phewas_filename = f"{gene}_azphewas.csv"
    phewas_file = os.path.join(phewas_dir, phewas_filename)

    if not os.path.exists(phewas_file):
        with sync_playwright() as playwright:
            run(playwright, gene=gene, save_path=phewas_dir)

    doc.add_heading("Phenome-wide Association Study (PheWAS) in UK Biobank cohort collected from AstraZeneca", level=2)
    if os.path.exists(phewas_file):
        fig, plot_df = get_azn_manhattan(gene, phewas_file)
        fig_path = uth._save_fig_to_tmp(fig, basename="azn_phewas", dpi=300)
        doc.add_paragraph()  # spacing
        doc.add_picture(fig_path, width=Inches(6.5))
        uth._add_figure_caption(doc, f"{gene} PheWAS results from AstraZeneca.")
        doc.add_paragraph()  # spacing

        # Table: Top 10 significant hits (sorted by p-value)
        top_significant = plot_df.nsmallest(10, 'P value')
        columns = ['Phenotype', 'Phenotypic category', 'Odds ratio', 'P value']
        # Rename columns for display
        display_df = top_significant[columns].copy()
        display_df.columns = ['Phenotype', 'Category', 'OR', 'P-value']
        uth._add_table_to_doc(doc, display_df, uth._add_table_title(doc, f"Top 10 significant associations for {gene} in AstraZeneca"),
                        ['Phenotype', 'Category', 'OR', 'P-value'], list(map(Inches, [2.75, 1.75, 0.75, 0.75]))
                        )
        # Table: Top 10 negative beta hits (odds_ratio < 0, sorted by p-value)
        negative_beta = plot_df[plot_df['Odds ratio'] < 1].nsmallest(10, 'P value')
        if len(negative_beta) > 0:
            display_df_neg = negative_beta[columns].copy()
            display_df_neg.columns = ['Phenotype', 'Category', 'OR', 'P-value']
            uth._add_table_to_doc(doc, display_df_neg, uth._add_table_title(doc, f"Top 10 protective associations for {gene} in AstraZeneca"),
                            ['Phenotype', 'Category', 'OR', 'P-value'], list(map(Inches, [2.75, 1.75, 0.75, 0.75]))
                            )
        doc.add_page_break()  # spacing

    else:
        doc.add_paragraph(f"No AstraZeneca PheWAS results found for {gene}.")
        doc.add_page_break()  # spacing
    return doc


if __name__ == "__main__":
    from docx import Document
    doc = Document()
    gene = "INHBE"
    phewas_dir = "/home/dbanerjee/deepro/rwe/data/test/"
    doc = generate_azn_clinical_report(doc, gene, phewas_dir)
    doc.save("/home/dbanerjee/deepro/rwe/data/test/INHBE_azn_report.docx")

