import os
from datetime import date
from importlib.metadata import version, PackageNotFoundError
from docx import Document
from docx.shared import Inches, Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
import pandas as pd


import rwe.utils.helpers as uth
import rwe.clients.aou as aou
import rwe.clients.azn as azn
import rwe.clients.genebass as gbs


def generate_title_and_contents(gene: str, logo_path: str | None = None, report_date: str | None = None) -> Document:
    """
    Generate the title and table of contents pages for the RWE report.

    Parameters
    ----------
    gene : str
        Gene symbol for the report (e.g., "PCSK9").
    logo_path : str | None
        Path to logo image file to include on title page. If None or file not found,
        falls back to default logo in assets folder. If no logo is available, no logo is added.
    report_date : str | None
        Date string to include on title page. If None, defaults to today's date in "Month DD, YYYY" format.

    Returns
    -------
    Document
        A python-docx Document object with the title and TOC pages added.
    """
    if not report_date:
        # default to today's date (you can override)
        report_date = date.today().strftime("%B %d, %Y")

    doc = Document()
    uth._set_default_styles(doc)
    uth._init_report_counters(doc)

    # --- Title page ---
    # Determine which logo to use
    final_logo_path = None
    if logo_path and os.path.exists(logo_path):
        final_logo_path = logo_path
    else:
        # Fallback to default logo in assets folder
        assets_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'assets')
        default_logo = os.path.join(assets_dir, 'Arrowhead_Pharmaceuticals_logo.png')
        if os.path.exists(default_logo):
            final_logo_path = default_logo

    if final_logo_path:
        p = doc.add_paragraph()
        p.alignment = WD_ALIGN_PARAGRAPH.CENTER
        # adjust width as needed
        p.add_run().add_picture(final_logo_path, width=Inches(2.0))

    title = doc.add_paragraph()
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    r = title.add_run(f"Real World Evidence of {gene} gene pLoF carriers")
    r.bold = True
    r.font.size = Pt(24)

    doc.add_paragraph()  # spacer

    authors = doc.add_paragraph()
    authors.alignment = WD_ALIGN_PARAGRAPH.CENTER
    authors.add_run("Deepro Banerjee, Shicheng Guo").font.size = Pt(14)

    dt = doc.add_paragraph()
    dt.alignment = WD_ALIGN_PARAGRAPH.CENTER
    dt.add_run(report_date).font.size = Pt(14)

    dept = doc.add_paragraph()
    dept.alignment = WD_ALIGN_PARAGRAPH.CENTER
    dept.add_run("Translational Genetics and Data Science").font.size = Pt(14)

    org = doc.add_paragraph()
    org.alignment = WD_ALIGN_PARAGRAPH.CENTER
    org.add_run("Arrowhead Pharmaceuticals").font.size = Pt(14)

    doc.add_page_break()

    # --- Contents page ---
    p = doc.add_paragraph()
    r = p.add_run("Table of Contents")
    r.bold = True
    r.font.size = Pt(14)

    # TOC field goes in its own empty paragraph
    toc_p = doc.add_paragraph()
    uth._add_toc_field(toc_p)


    note = doc.add_paragraph(
        "Note: In Microsoft Word, right-click the table of contents and choose 'Update Field' "
        "to populate page numbers."
    )
    note.runs[0].italic = True

    doc.add_page_break()
    return doc


def generate_variant_information_and_demographics(doc: Document, chrm: str, gene: str, zygosity: str, allofus: bool = True) -> Document:
    doc.add_heading("Variant information and demographics", level=1)
    # AoU
    if allofus:
        doc = aou.generate_aou_variant_info_demographics_report(doc, chrm, gene, zygosity)   
    return doc

def generate_clinical_records(doc: Document, chrm:str, gene: str, zygosity: str,  output_dir: str, allofus: bool = True, genebass: bool = False, astrazeneca: bool = True) -> Document:
    doc.add_heading("Clinical records", level=1)
    # AoU
    if allofus:
        doc = aou.generate_aou_clinical_report(doc, chrm, gene, zygosity, output_dir)
    # GeneBass
    if genebass:
        doc = gbs.generate_genebass_clinical_report(doc, gene, output_dir)
    # AZN
    if astrazeneca:
        doc = azn.generate_azn_clinical_report(doc, gene, output_dir)
    return doc

def generate_labs_and_measurements(doc: Document, chrm: str, gene: str, zygosity: str, output_dir: str, allofus: bool = True, genebass: bool = False, astrazeneca: bool = True) -> Document:
    doc.add_heading("Labs and measurements", level=1)
    # AoU
    if allofus:
        doc = aou.generate_aou_labs_measurements_report(doc, chrm, gene, zygosity)
        doc.add_paragraph("")
    # GeneBass
    if genebass:
        doc = gbs.generate_genebass_labs_measurements_report(doc, gene, output_dir)
        doc.add_paragraph("")
    # AZN
    if astrazeneca:
        doc = azn.generate_azn_labs_measurements_report(doc, gene, output_dir)
        doc.add_paragraph("")
    return doc

def generate_survey_information(doc: Document, chrm: str, gene: str, zygosity: str, output_dir: str, allofus: bool = True) -> Document:
    doc.add_heading("Survey information", level=1)
    # AoU
    if allofus:
        doc = aou.generate_aou_survey_report(doc, chrm, gene, zygosity)
        doc.add_paragraph("")
    return doc

def generate_homozygous_lof_carriers(doc: Document, chrm: str, gene: str, output_dir: str, allofus: bool = True) -> Document:
    doc.add_heading("Homozygous loss of function carriers", level=1)
    # AoU
    if allofus:
        doc = aou.generate_aou_homozygous_lof_carriers_report(doc, chrm, gene)
        doc.add_paragraph("")
    return doc

def generate_indication_specific_report(doc: Document, chrm: str, gene: str, output_dir: str, indications: list, keywords: list, allofus: bool = True, genebass: bool = False, astrazeneca: bool = True) -> Document:
    doc.add_heading(f"Indication-specific report for {gene}", level=1)
    doc.add_paragraph("Indication(s) searched: " + ", ".join(indications))
    # AoU
    if allofus:
        doc = aou.generate_aou_indication_specific_report(doc, chrm, gene, output_dir, keywords)
    # GeneBass
    if genebass:
        doc = gbs.generate_genebass_indication_specific_report(doc, gene, output_dir, keywords)
    # AZN     
    if astrazeneca:
        doc = azn.generate_azn_indication_specific_report(doc, gene, output_dir, keywords)
    doc.add_paragraph("")
    return doc

def generate_plasma_proteomics_report(doc: Document, chrm: str, gene: str, output_dir: str, astrazeneca: bool = True) -> Document:
    doc.add_heading("Plasma proteomics", level=1)
    # AZN
    if astrazeneca:
        doc = azn.generate_azn_plasma_protein_biomarkers_report(doc, gene, output_dir)
    doc.add_paragraph("")
    return doc

def generate_end_of_report(doc: Document) -> Document:
    try:
        pkg_version = version("rwe")
    except PackageNotFoundError:
        pkg_version = "unknown (local source)"
    doc.add_heading("End of report", level=1)
    doc.add_paragraph()
    note = doc.add_paragraph(
        f"This report was generated by rwe package v{pkg_version} using data available at the time of analysis."
    )
    note.runs[0].italic = True
    return doc
