import os
from datetime import date
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
from google.cloud import storage

_FIGURE_SAVE_DIR = ""
_FIGURE_SAVE_PREFIX = ""

# ---------- DOCX helpers ----------

def _set_figure_save_context(output_dir: str, prefix: str = "") -> None:
    global _FIGURE_SAVE_DIR, _FIGURE_SAVE_PREFIX
    _FIGURE_SAVE_DIR = output_dir
    _FIGURE_SAVE_PREFIX = prefix


def _clear_figure_save_context() -> None:
    global _FIGURE_SAVE_DIR, _FIGURE_SAVE_PREFIX
    _FIGURE_SAVE_DIR = ""
    _FIGURE_SAVE_PREFIX = ""

def _init_report_counters(doc: Document) -> Document:
    """Initialize figure and table counters on the document object."""
    doc._report_counters = {"figure": 0, "table": 0}
    return doc


def _get_next_figure_num(doc: Document) -> int:
    """Get the next figure number and increment the counter."""
    if not hasattr(doc, '_report_counters'):
        doc._report_counters = {"figure": 0, "table": 0}
    doc._report_counters["figure"] += 1
    return doc._report_counters["figure"]


def _get_next_table_num(doc: Document) -> int:
    """Get the next table number and increment the counter."""
    if not hasattr(doc, '_report_counters'):
        doc._report_counters = {"figure": 0, "table": 0}
    doc._report_counters["table"] += 1
    return doc._report_counters["table"]


def _set_default_styles(doc: Document):
    # Normal
    style = doc.styles["Normal"]
    style.font.name = "Calibri"
    style.font.size = Pt(11)

    # Headings
    for hname, size in [("Heading 1", 14), ("Heading 2", 12)]:
        s = doc.styles[hname]
        s.font.name = "Calibri"
        s.font.size = Pt(size)
        s.font.bold = True


def _add_toc_field(paragraph):
    """
    Insert a Word TOC field. Word will render it after you:
      References -> Update Table
    """
    run = paragraph.add_run()

    fld_begin = OxmlElement("w:fldChar")
    fld_begin.set(qn("w:fldCharType"), "begin")

    instr = OxmlElement("w:instrText")
    instr.set(qn("xml:space"), "preserve")
    # TOC switches:
    #   \\o "1-3" include headings 1-3
    #   \\h hyperlinks
    #   \\z hide page numbers in web layout
    #   \\u use outline levels
    instr.text = r'TOC \o "1-3" \h \z \u'

    fld_sep = OxmlElement("w:fldChar")
    fld_sep.set(qn("w:fldCharType"), "separate")

    fld_end = OxmlElement("w:fldChar")
    fld_end.set(qn("w:fldCharType"), "end")

    run._r.append(fld_begin)
    run._r.append(instr)
    run._r.append(fld_sep)
    run._r.append(fld_end)
    return


def _add_caption(doc: Document, text: str):
    """Add a caption paragraph (for backward compatibility)."""
    p = doc.add_paragraph(text)
    p.style = doc.styles["Normal"]
    p.runs[0].italic = True
    p.alignment = WD_ALIGN_PARAGRAPH.CENTER
    return p


def _add_figure_caption(doc: Document, description: str):
    """Add an auto-numbered figure caption.

    Args:
        doc: Document object with _report_counters attached
        description: Caption text without the "Figure X." prefix
    """
    fig_num = _get_next_figure_num(doc)
    text = f"Figure {fig_num}. {description}"
    return _add_caption(doc, text)


def _add_table_title(doc: Document, description: str):
    """Get an auto-numbered table title string.

    Args:
        doc: Document object with _report_counters attached
        description: Title text without the "Table X." prefix

    Returns:
        str: The full table title with number prefix
    """
    table_num = _get_next_table_num(doc)
    return f"Table {table_num}. {description}"

def _save_fig_to_tmp(fig, basename="figure", dpi=300):
    # Keep function name for backward compatibility, but always persist figures.
    fig_dir = _FIGURE_SAVE_DIR
    if not fig_dir:
        fallback_gene_dir = _FIGURE_SAVE_PREFIX or "gene"
        fig_dir = os.path.join(".", "data", "rwe_reports", fallback_gene_dir)
    os.makedirs(fig_dir, exist_ok=True)
    fig_prefix = _FIGURE_SAVE_PREFIX
    fig_name = f"{fig_prefix}_{basename}.png" if fig_prefix else f"{basename}.png"
    path = os.path.join(fig_dir, fig_name)
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    return path

def _set_table_autofit(table, autofit: bool = False):
    # python-docx doesn't expose this cleanly; set tblLayout to fixed
    tbl_pr = table._tbl.tblPr
    tbl_layout = tbl_pr.find(qn('w:tblLayout'))
    if tbl_layout is None:
        tbl_layout = OxmlElement('w:tblLayout')
        tbl_pr.append(tbl_layout)
    tbl_layout.set(qn('w:type'), 'autofit' if autofit else 'fixed')

def _add_table_to_doc(doc, df, title, columns, col_widths=None, table_width_inches=None):
    """
    col_widths: list same length as columns, items are Inches/Cm, e.g. [Inches(1.2), Inches(3.0), ...]
    table_width_inches: optional total width; if provided and col_widths is None, widths are evenly split
    """
    # Add title
    p = doc.add_paragraph()
    run = p.add_run(title)
    run.bold = True

    table = doc.add_table(rows=1, cols=len(columns))
    table.style = 'Light Grid Accent 1'
    _set_table_autofit(table, autofit=False)

    # decide widths
    if col_widths is None and table_width_inches is not None:
        w = Inches(table_width_inches / len(columns))
        col_widths = [w] * len(columns)

    # Header
    header_cells = table.rows[0].cells
    for i, col in enumerate(columns):
        header_cells[i].text = col
        header_cells[i].paragraphs[0].runs[0].bold = True
        if col_widths is not None:
            header_cells[i].width = col_widths[i]

    # Rows
    for _, row in df.iterrows():
        row_cells = table.add_row().cells
        for i, col in enumerate(columns):
            value = row[col]
            if isinstance(value, float):
                if col in ('P-Value (Burden)', 'P-value', 'P value', 'p_value'):
                    row_cells[i].text = f"{value:.2e}"
                else:
                    row_cells[i].text = f"{value:.2f}"
            else:
                row_cells[i].text = str(value)

            if col_widths is not None:
                row_cells[i].width = col_widths[i]

    doc.add_paragraph()
    return doc

# AoU helpers
def _relpath(gs_uri: str, bucket: str) -> str:
    # gs://bucket/prefix/file -> prefix/file
    return gs_uri.replace(f"{bucket}/", "")

def _gcs_size(gs_uri: str, bucket: str) -> int:
    storage_client = storage.Client()
    bucket_obj = storage_client.bucket(bucket.lstrip("gs://"))
    rel = _relpath(gs_uri, bucket)
    blob = bucket_obj.blob(rel)
    if not blob.exists():
        return 0
    blob.reload()
    return blob.size or 0
