"""Font-style-driven ATX header injection for xlsx documents.

Pre-Docling: walks the .xlsx OOXML (xl/styles.xml, xl/sharedStrings.xml,
xl/worksheets/sheetN.xml), auto-detects the workbook body size, ranks
distinct bold-cell font sizes in descending order, assigns them to
H2..H6 dynamically, and rewrites each header-styled cell with sentinels
§§HN§§...§§/HN§§. The sheet display name is always written as H1.

Post-Docling: convert_sentinels_to_atx(md) materializes the sentinels
into ATX headers (#..######), strips surrounding `| ... |` table
wrappers and orphan divider rows, and dedupes within-sheet H1
duplicates from merged-cell unrolling.

Public API:
    inject_header_sentinels(xlsx_path) -> (count, detected_body, size_to_level)
    convert_sentinels_to_atx(md)       -> (md, {'H1': n, ..., 'H6': n})

Sentinels survive stock Docling escaping (validated on JP technical
xlsx specs). No Docling fork required.
"""
from __future__ import annotations

import logging
import os
import re
import shutil
import tempfile
import xml.etree.ElementTree as ET
import zipfile
from collections import Counter
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

NS = "http://schemas.openxmlformats.org/spreadsheetml/2006/main"
NS_R = "http://schemas.openxmlformats.org/officeDocument/2006/relationships"
NS_PKG = "http://schemas.openxmlformats.org/package/2006/relationships"

ET.register_namespace("", NS)

H1_SENTINEL, H1_END = "§§H1§§", "§§/H1§§"
H2_SENTINEL, H2_END = "§§H2§§", "§§/H2§§"
H3_SENTINEL, H3_END = "§§H3§§", "§§/H3§§"
H4_SENTINEL, H4_END = "§§H4§§", "§§/H4§§"
H5_SENTINEL, H5_END = "§§H5§§", "§§/H5§§"
H6_SENTINEL, H6_END = "§§H6§§", "§§/H6§§"

_LEVEL_SENTINELS = {
    1: (H1_SENTINEL, H1_END),
    2: (H2_SENTINEL, H2_END),
    3: (H3_SENTINEL, H3_END),
    4: (H4_SENTINEL, H4_END),
    5: (H5_SENTINEL, H5_END),
    6: (H6_SENTINEL, H6_END),
}

SENTINEL_RE = re.compile(
    r"(?P<pre>\|\s*)?"
    r"§§(?P<lvl>H[1-6])§§"
    r"(?P<text>.+?)"
    r"§§/(?P=lvl)§§"
    r"(?P<post>\s*\|)?",
    re.DOTALL,
)

DEFAULT_BODY_SIZE = 11.0
MAX_HEADER_LEVEL = 6  # markdown ATX cap


def _load_fonts(styles_root) -> Dict[int, Tuple[float, bool]]:
    out: Dict[int, Tuple[float, bool]] = {}
    fonts = styles_root.find(f"{{{NS}}}fonts")
    if fonts is None:
        return out
    for i, f in enumerate(fonts.findall(f"{{{NS}}}font")):
        sz_el = f.find(f"{{{NS}}}sz")
        b_el = f.find(f"{{{NS}}}b")
        out[i] = (
            float(sz_el.get("val")) if sz_el is not None else DEFAULT_BODY_SIZE,
            b_el is not None,
        )
    return out


def _load_xf_to_font(styles_root) -> List[int]:
    out: List[int] = []
    cellXfs = styles_root.find(f"{{{NS}}}cellXfs")
    if cellXfs is None:
        return out
    for xf in cellXfs.findall(f"{{{NS}}}xf"):
        out.append(int(xf.get("fontId", "0")))
    return out


def _load_shared_strings(ss_root) -> List[str]:
    out: List[str] = []
    for si in ss_root.findall(f"{{{NS}}}si"):
        t = si.find(f"{{{NS}}}t")
        if t is not None:
            out.append(t.text or "")
            continue
        parts: List[str] = []
        for r in si.findall(f"{{{NS}}}r"):
            rt = r.find(f"{{{NS}}}t")
            if rt is not None:
                parts.append(rt.text or "")
        out.append("".join(parts))
    return out


def _read_cell_text(c, shared: List[str]) -> str:
    t = c.get("t")
    if t == "s":
        v = c.find(f"{{{NS}}}v")
        if v is not None:
            try:
                return shared[int(v.text)]
            except (ValueError, IndexError):
                return ""
    elif t == "inlineStr":
        t_el = c.find(f"{{{NS}}}is/{{{NS}}}t")
        if t_el is not None:
            return t_el.text or ""
    else:
        v = c.find(f"{{{NS}}}v")
        if v is not None:
            return v.text or ""
    return ""


def _write_inline_string(c, text: str) -> None:
    for tag in (f"{{{NS}}}v", f"{{{NS}}}is", f"{{{NS}}}f"):
        for child in list(c.findall(tag)):
            c.remove(child)
    c.set("t", "inlineStr")
    is_el = ET.SubElement(c, f"{{{NS}}}is")
    t_el = ET.SubElement(is_el, f"{{{NS}}}t")
    t_el.text = text
    t_el.set("{http://www.w3.org/XML/1998/namespace}space", "preserve")


_CELL_REF_RE = re.compile(r"^([A-Z]+)(\d+)$")


def _split_ref(ref: str) -> Tuple[int, int]:
    m = _CELL_REF_RE.match(ref)
    if not m:
        raise ValueError(f"bad cell ref: {ref}")
    letters, row = m.group(1), int(m.group(2))
    col = 0
    for ch in letters:
        col = col * 26 + (ord(ch) - ord("A") + 1)
    return col, row


def _idx_to_col(idx: int) -> str:
    out = ""
    while idx > 0:
        idx, rem = divmod(idx - 1, 26)
        out = chr(ord("A") + rem) + out
    return out


def _expand_merge_range(rng: str) -> List[str]:
    if ":" not in rng:
        return [rng]
    a, b = rng.split(":")
    ac, ar = _split_ref(a)
    bc, br = _split_ref(b)
    refs: List[str] = []
    for col in range(min(ac, bc), max(ac, bc) + 1):
        for r in range(min(ar, br), max(ar, br) + 1):
            refs.append(f"{_idx_to_col(col)}{r}")
    return refs


def _sheet_name_map(workbook_root, rels_root) -> Dict[str, str]:
    rid_to_target: Dict[str, str] = {}
    for rel in rels_root.findall(f"{{{NS_PKG}}}Relationship"):
        rid_to_target[rel.get("Id")] = rel.get("Target")
    out: Dict[str, str] = {}
    sheets = workbook_root.find(f"{{{NS}}}sheets")
    if sheets is None:
        return out
    for s in sheets.findall(f"{{{NS}}}sheet"):
        rid = s.get(f"{{{NS_R}}}id")
        target = rid_to_target.get(rid, "")
        fname = os.path.basename(target)
        out[fname] = s.get("name", fname)
    return out


def _detect_body_size(cell_styles: List[Tuple[float, bool]]) -> float:
    """Return the workbook's body font size.

    Strategy (in order):
        1. mode of sizes from non-bold cells
        2. mode of sizes across all cells
        3. DEFAULT_BODY_SIZE (11.0) only if nothing resolvable
    """
    non_bold = [sz for sz, b in cell_styles if not b]
    if non_bold:
        return Counter(non_bold).most_common(1)[0][0]
    all_sizes = [sz for sz, _ in cell_styles]
    if all_sizes:
        return Counter(all_sizes).most_common(1)[0][0]
    return DEFAULT_BODY_SIZE


def _rank_heading_sizes(bold_sizes: set[float]) -> Dict[float, int]:
    """Assign each distinct bold font size to a markdown level 1..MAX_HEADER_LEVEL.

    Largest size -> H1 (cell text will be overridden with the sheet display
    name during injection), next -> H2, ... If more distinct sizes exist
    than available levels, the smallest tail merges into the bottom level.
    """
    if not bold_sizes:
        return {}
    ordered = sorted(bold_sizes, reverse=True)
    mapping: Dict[float, int] = {}
    if len(ordered) <= MAX_HEADER_LEVEL:
        for i, sz in enumerate(ordered):
            mapping[sz] = i + 1
    else:
        for i, sz in enumerate(ordered[: MAX_HEADER_LEVEL - 1]):
            mapping[sz] = i + 1
        tail_level = MAX_HEADER_LEVEL
        merged = len(ordered) - (MAX_HEADER_LEVEL - 1)
        for sz in ordered[MAX_HEADER_LEVEL - 1 :]:
            mapping[sz] = tail_level
        logger.info("merged %d smallest bold sizes into H%d", merged, tail_level)
    return mapping


def _wrap_sentinel(level: int, text: str) -> str:
    start, end = _LEVEL_SENTINELS[level]
    return f"{start}{text}{end}"


def inject_header_sentinels(xlsx_path: str) -> Tuple[int, float, Dict[float, int]]:
    """Rewrite header-styled cells in `xlsx_path` with §§HN§§ sentinels in place.

    Returns:
        (rewritten_cell_count, detected_body_size, size_to_level_map)

    On structural failure (missing xl/styles.xml or empty cellXfs), logs a
    warning and returns (0, DEFAULT_BODY_SIZE, {}) without mutating the file.
    """
    workdir = tempfile.mkdtemp(prefix="hdr_inject_")
    try:
        try:
            with zipfile.ZipFile(xlsx_path) as z:
                z.extractall(workdir)
        except zipfile.BadZipFile:
            logger.warning("not a valid xlsx (BadZipFile): %s", xlsx_path)
            return 0, DEFAULT_BODY_SIZE, {}

        styles_path = os.path.join(workdir, "xl/styles.xml")
        if not os.path.exists(styles_path):
            logger.warning("xl/styles.xml missing in %s; skipping header injection", xlsx_path)
            return 0, DEFAULT_BODY_SIZE, {}

        styles = ET.parse(styles_path).getroot()
        fonts = _load_fonts(styles)
        xf_to_font = _load_xf_to_font(styles)
        if not xf_to_font:
            logger.warning("xl/styles.xml has no cellXfs in %s; skipping header injection", xlsx_path)
            return 0, DEFAULT_BODY_SIZE, {}

        ss_path = os.path.join(workdir, "xl/sharedStrings.xml")
        shared = _load_shared_strings(ET.parse(ss_path).getroot()) if os.path.exists(ss_path) else []

        workbook_path = os.path.join(workdir, "xl/workbook.xml")
        rels_path = os.path.join(workdir, "xl/_rels/workbook.xml.rels")
        if not (os.path.exists(workbook_path) and os.path.exists(rels_path)):
            logger.warning("workbook.xml or rels missing in %s; skipping", xlsx_path)
            return 0, DEFAULT_BODY_SIZE, {}
        workbook = ET.parse(workbook_path).getroot()
        rels = ET.parse(rels_path).getroot()
        sheet_names = _sheet_name_map(workbook, rels)

        ws_dir = os.path.join(workdir, "xl/worksheets")
        if not os.path.isdir(ws_dir):
            logger.warning("xl/worksheets missing in %s; skipping", xlsx_path)
            return 0, DEFAULT_BODY_SIZE, {}

        sheet_files = [fn for fn in sorted(os.listdir(ws_dir)) if fn.endswith(".xml")]

        cell_styles: List[Tuple[float, bool]] = []
        bold_sizes: set[float] = set()
        for fname in sheet_files:
            root = ET.parse(os.path.join(ws_dir, fname)).getroot()
            sd = root.find(f"{{{NS}}}sheetData")
            if sd is None:
                continue
            for row in sd.findall(f"{{{NS}}}row"):
                for c in row.findall(f"{{{NS}}}c"):
                    s_idx = c.get("s")
                    if s_idx is None:
                        continue
                    try:
                        fid = xf_to_font[int(s_idx)]
                    except (ValueError, IndexError):
                        continue
                    sz, bold = fonts.get(fid, (DEFAULT_BODY_SIZE, False))
                    if not _read_cell_text(c, shared).strip():
                        continue
                    cell_styles.append((sz, bold))
                    if bold:
                        bold_sizes.add(sz)

        body_size = _detect_body_size(cell_styles)
        size_to_level = _rank_heading_sizes(bold_sizes)

        rewritten = 0
        for fname in sheet_files:
            sheet_display = sheet_names.get(fname, fname)
            path = os.path.join(ws_dir, fname)
            tree = ET.parse(path)
            root = tree.getroot()
            sd = root.find(f"{{{NS}}}sheetData")
            if sd is None:
                continue

            cell_by_ref: Dict[str, ET.Element] = {}
            for row in sd.findall(f"{{{NS}}}row"):
                for c in row.findall(f"{{{NS}}}c"):
                    if c.get("r"):
                        cell_by_ref[c.get("r")] = c

            merges: List[str] = []
            mc_el = root.find(f"{{{NS}}}mergeCells")
            if mc_el is not None:
                for m in mc_el.findall(f"{{{NS}}}mergeCell"):
                    rng = m.get("ref")
                    if rng:
                        merges.append(rng)

            sentinel_refs: List[str] = []
            for row in sd.findall(f"{{{NS}}}row"):
                for c in row.findall(f"{{{NS}}}c"):
                    s_idx = c.get("s")
                    if s_idx is None:
                        continue
                    try:
                        fid = xf_to_font[int(s_idx)]
                    except (ValueError, IndexError):
                        continue
                    sz, bold = fonts.get(fid, (DEFAULT_BODY_SIZE, False))
                    if not bold:
                        continue
                    raw = _read_cell_text(c, shared).strip()
                    if not raw:
                        continue
                    level = size_to_level.get(sz)
                    if level is None:
                        continue
                    text = sheet_display if level == 1 else raw
                    _write_inline_string(c, _wrap_sentinel(level, text))
                    rewritten += 1
                    sentinel_refs.append(c.get("r"))

            sentinel_set = set(sentinel_refs)
            for rng in merges:
                refs = _expand_merge_range(rng)
                anchor = next((r for r in refs if r in sentinel_set), None)
                if anchor is None:
                    continue
                for r in refs:
                    if r == anchor:
                        continue
                    other = cell_by_ref.get(r)
                    if other is None:
                        continue
                    for tag in (f"{{{NS}}}v", f"{{{NS}}}is", f"{{{NS}}}f"):
                        for ch in list(other.findall(tag)):
                            other.remove(ch)
                    if "t" in other.attrib:
                        del other.attrib["t"]

            tree.write(path, xml_declaration=True, encoding="UTF-8")

        out_fd, out_tmp = tempfile.mkstemp(suffix=".xlsx", prefix="hdr_inject_out_")
        os.close(out_fd)
        try:
            with zipfile.ZipFile(out_tmp, "w", zipfile.ZIP_DEFLATED) as z:
                for r, _, files in os.walk(workdir):
                    for f in files:
                        full = os.path.join(r, f)
                        arc = os.path.relpath(full, workdir)
                        z.write(full, arc)
            shutil.move(out_tmp, xlsx_path)
        finally:
            if os.path.exists(out_tmp):
                os.remove(out_tmp)

        return rewritten, body_size, size_to_level
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


_HEADER_PREFIX = {n: "#" * n for n in range(1, MAX_HEADER_LEVEL + 1)}


def convert_sentinels_to_atx(md: str) -> Tuple[str, Dict[str, int]]:
    """Materialize §§HN§§...§§/HN§§ sentinels into ATX headers.

    Strips surrounding `| ... |` table-cell pipes, collapses internal
    whitespace runs, dedupes within-sheet H1 duplicates from merged-cell
    unrolling, removes orphan table dividers immediately following a
    materialized header, and collapses runs of blank lines.

    Returns (rewritten_md, counts) where counts is keyed `H1`..`H6` for
    levels actually emitted (zero-count levels are omitted).
    """

    def _repl(m: re.Match) -> str:
        lvl = m.group("lvl")
        n = int(lvl[1])
        text = re.sub(r"\s+", " ", m.group("text").strip())
        return f"\n\n{_HEADER_PREFIX[n]} {text}\n\n"

    new_md = SENTINEL_RE.sub(_repl, md)
    new_md = re.sub(r"\n{3,}", "\n\n", new_md)

    lines = new_md.split("\n")
    deduped: List[str] = []
    last_h1: Optional[str] = None
    last_adjacent_header: Optional[str] = None
    for line in lines:
        m = re.match(r"^(#{1,6}) ", line)
        if m:
            level = len(m.group(1))
            if level == 1:
                if line == last_h1:
                    continue
                last_h1 = line
                last_adjacent_header = line
            else:
                if line == last_adjacent_header:
                    continue
                last_adjacent_header = line
        elif line.strip():
            last_adjacent_header = None
        deduped.append(line)
    new_md = "\n".join(deduped)

    new_md = re.sub(r"\n\|[-:\s|]+\|\s*\n(?=\s*\n#)", "\n", new_md)
    new_md = re.sub(r"(\n#{1,6} [^\n]+\n\n)\|[-:\s|]+\|\s*\n", r"\1", new_md)
    new_md = re.sub(r"\n{3,}", "\n\n", new_md)

    counts: Dict[str, int] = {}
    for line in new_md.split("\n"):
        m = re.match(r"^(#{1,6}) ", line)
        if m:
            key = f"H{len(m.group(1))}"
            counts[key] = counts.get(key, 0) + 1

    return new_md, counts