"""Unit tests for src.extraction_v2.xlsx_header_injection.

Covers the I/O matrix from docs/implementation-artifacts/spec-textiq-627-xlsx-font-header-injection.md.
"""
from __future__ import annotations

import os
import re
import shutil
import sys
import tempfile
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path

import pytest
from openpyxl import Workbook
from openpyxl.styles import Font
from openpyxl.utils import get_column_letter

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT))

from src.extraction_v2.xlsx_header_injection import (  # noqa: E402
    DEFAULT_BODY_SIZE,
    MAX_HEADER_LEVEL,
    SENTINEL_RE,
    _detect_body_size,
    _rank_heading_sizes,
    convert_sentinels_to_atx,
    inject_header_sentinels,
)

NS = "http://schemas.openxmlformats.org/spreadsheetml/2006/main"


def _save_workbook(wb: Workbook, dir_: str, name: str = "test.xlsx") -> str:
    path = os.path.join(dir_, name)
    wb.save(path)
    return path


def _read_inline_strings(xlsx_path: str) -> list[str]:
    """Return every inline-string cell text across all sheets (sentinels survive here)."""
    workdir = tempfile.mkdtemp(prefix="t_read_")
    try:
        with zipfile.ZipFile(xlsx_path) as z:
            z.extractall(workdir)
        ws_dir = os.path.join(workdir, "xl/worksheets")
        out: list[str] = []
        for fname in sorted(os.listdir(ws_dir)):
            if not fname.endswith(".xml"):
                continue
            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"):
                    if c.get("t") == "inlineStr":
                        t = c.find(f"{{{NS}}}is/{{{NS}}}t")
                        if t is not None and t.text:
                            out.append(t.text)
        return out
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


@pytest.fixture
def tmp_xlsx_dir(tmp_path):
    return str(tmp_path)


# ---------------------------------------------------------------------------
# Helpers under test
# ---------------------------------------------------------------------------


def test_rank_heading_sizes_three_tiers():
    assert _rank_heading_sizes({14.0, 11.0, 9.0}) == {14.0: 1, 11.0: 2, 9.0: 3}


def test_rank_heading_sizes_empty():
    assert _rank_heading_sizes(set()) == {}


def test_rank_heading_sizes_tail_merge_into_h6():
    mapping = _rank_heading_sizes({16.0, 14.0, 12.0, 11.0, 10.0, 9.0, 8.0})
    # 5 distinct levels available before merge; smallest 2 sizes both → H6
    assert mapping[16.0] == 1
    assert mapping[14.0] == 2
    assert mapping[12.0] == 3
    assert mapping[11.0] == 4
    assert mapping[10.0] == 5
    assert mapping[9.0] == MAX_HEADER_LEVEL
    assert mapping[8.0] == MAX_HEADER_LEVEL


def test_detect_body_size_prefers_non_bold_mode():
    styles = [(11.0, False), (11.0, False), (14.0, True), (11.0, False), (9.0, True)]
    assert _detect_body_size(styles) == 11.0


def test_detect_body_size_falls_back_to_all_when_all_bold():
    styles = [(12.0, True), (12.0, True), (16.0, True)]
    assert _detect_body_size(styles) == 12.0


def test_detect_body_size_default_when_empty():
    assert _detect_body_size([]) == DEFAULT_BODY_SIZE


# ---------------------------------------------------------------------------
# I/O matrix: pre-Docling injection
# ---------------------------------------------------------------------------


def test_happy_path_3tier_injection(tmp_xlsx_dir):
    """HAPPY_PATH_3TIER: bold sizes {14,11,9}, body cells at 11."""
    wb = Workbook()
    ws = wb.active
    ws.title = "Sheet1"
    ws["A1"] = "Title"
    ws["A1"].font = Font(size=14, bold=True)
    ws["A2"] = "Section"
    ws["A2"].font = Font(size=11, bold=True)
    ws["A3"] = "Step"
    ws["A3"].font = Font(size=9, bold=True)
    ws["A4"] = "body text 1"
    ws["A4"].font = Font(size=11, bold=False)
    ws["A5"] = "body text 2"
    ws["A5"].font = Font(size=11, bold=False)
    path = _save_workbook(wb, tmp_xlsx_dir)

    count, body, size_map = inject_header_sentinels(path)

    assert count == 3
    assert body == 11.0
    assert size_map == {14.0: 1, 11.0: 2, 9.0: 3}

    strings = _read_inline_strings(path)
    # H1 cell text replaced with the sheet display name
    assert any("§§H1§§Sheet1§§/H1§§" in s for s in strings)
    assert any("§§H2§§Section§§/H2§§" in s for s in strings)
    assert any("§§H3§§Step§§/H3§§" in s for s in strings)


def test_different_body_size(tmp_xlsx_dir):
    wb = Workbook()
    ws = wb.active
    ws.title = "S"
    ws["A1"] = "Big"
    ws["A1"].font = Font(size=13, bold=True)
    ws["A2"] = "Section"
    ws["A2"].font = Font(size=10, bold=True)
    for i in range(3, 10):
        ws[f"A{i}"] = "body"
        ws[f"A{i}"].font = Font(size=10, bold=False)
    path = _save_workbook(wb, tmp_xlsx_dir)

    _, body, size_map = inject_header_sentinels(path)

    assert body == 10.0
    assert size_map == {13.0: 1, 10.0: 2}


def test_many_tiers_depth_cap(tmp_xlsx_dir):
    wb = Workbook()
    ws = wb.active
    ws.title = "S"
    for i, sz in enumerate([16, 14, 12, 11, 10, 9, 8], start=1):
        ws.cell(row=i, column=1, value=f"tier{sz}").font = Font(size=sz, bold=True)
    ws["B1"] = "body"
    ws["B1"].font = Font(size=11, bold=False)
    path = _save_workbook(wb, tmp_xlsx_dir)

    count, _, size_map = inject_header_sentinels(path)

    assert count == 7
    # smallest two (9, 8) merge into the bottom level
    assert size_map[9.0] == MAX_HEADER_LEVEL
    assert size_map[8.0] == MAX_HEADER_LEVEL
    # H6 sentinel appears for both
    strings = _read_inline_strings(path)
    h6_hits = [s for s in strings if "§§H6§§" in s]
    assert len(h6_hits) == 2


def test_all_bold_fallback(tmp_xlsx_dir):
    wb = Workbook()
    ws = wb.active
    ws.title = "S"
    ws["A1"] = "a"
    ws["A1"].font = Font(size=12, bold=True)
    ws["A2"] = "b"
    ws["A2"].font = Font(size=12, bold=True)
    ws["A3"] = "c"
    ws["A3"].font = Font(size=16, bold=True)
    path = _save_workbook(wb, tmp_xlsx_dir)

    _, body, size_map = inject_header_sentinels(path)

    # No non-bold cells exist, so body falls back to mode of all sizes (12.0)
    assert body == 12.0
    assert size_map == {16.0: 1, 12.0: 2}


def test_merged_h1_dedup_via_clear(tmp_xlsx_dir):
    """MERGED_H1: title cell merged across 4 cols × 2 rows. Sibling cells inside
    the merge range must be cleared so Docling can't render duplicate sentinels."""
    wb = Workbook()
    ws = wb.active
    ws.title = "MyDocSheet"
    # Set the anchor cell's value and font, then merge — openpyxl handles the
    # merge by reading from the top-left anchor only.
    ws["A1"] = "Doc Title"
    ws["A1"].font = Font(size=14, bold=True)
    ws.merge_cells(start_row=1, start_column=1, end_row=2, end_column=4)
    ws["A3"] = "body"
    ws["A3"].font = Font(size=11, bold=False)
    path = _save_workbook(wb, tmp_xlsx_dir)

    inject_header_sentinels(path)

    strings = _read_inline_strings(path)
    h1_hits = [s for s in strings if "§§H1§§" in s]
    # Anchor cell holds the sentinel with the sheet display name.
    assert len(h1_hits) == 1
    assert "§§H1§§MyDocSheet§§/H1§§" in h1_hits[0]


def test_no_bold_cells(tmp_xlsx_dir):
    wb = Workbook()
    ws = wb.active
    ws.title = "S"
    for i in range(1, 6):
        ws[f"A{i}"] = f"row {i}"
        ws[f"A{i}"].font = Font(size=11, bold=False)
    path = _save_workbook(wb, tmp_xlsx_dir)

    count, _, size_map = inject_header_sentinels(path)

    assert count == 0
    assert size_map == {}
    strings = _read_inline_strings(path)
    assert not any("§§H" in s for s in strings)


def test_mixed_format_run_uses_cell_level_style(tmp_xlsx_dir):
    """openpyxl's Font is a cell-level style. A non-bold cell-level font with
    rich-text runs would still classify as non-bold (no sentinel)."""
    wb = Workbook()
    ws = wb.active
    ws.title = "S"
    ws["A1"] = "mixed text"
    ws["A1"].font = Font(size=11, bold=False)
    ws["A2"] = "bold section"
    ws["A2"].font = Font(size=11, bold=True)
    path = _save_workbook(wb, tmp_xlsx_dir)

    inject_header_sentinels(path)

    strings = _read_inline_strings(path)
    sentinel_strings = [s for s in strings if "§§H" in s]
    # Only the cell-level bold cell produces a sentinel; non-bold cell does not.
    assert any("§§H1§§S§§/H1§§" in s for s in sentinel_strings)
    assert not any("mixed text" in s for s in sentinel_strings)


def test_bad_ooxml_returns_zero_without_mutation(tmp_xlsx_dir):
    bad = os.path.join(tmp_xlsx_dir, "bad.xlsx")
    with open(bad, "w") as f:
        f.write("not a zip")

    count, body, size_map = inject_header_sentinels(bad)

    assert count == 0
    assert body == DEFAULT_BODY_SIZE
    assert size_map == {}
    # File untouched
    with open(bad) as f:
        assert f.read() == "not a zip"


# ---------------------------------------------------------------------------
# I/O matrix: post-Docling sentinel → ATX conversion
# ---------------------------------------------------------------------------


def test_convert_sentinels_basic_levels():
    md = "§§H1§§Title§§/H1§§\nsome body\n§§H3§§Step§§/H3§§"
    out, counts = convert_sentinels_to_atx(md)
    assert "# Title" in out
    assert "### Step" in out
    assert counts == {"H1": 1, "H3": 1}


def test_convert_sentinels_strips_table_pipes_docling_wrap():
    """DOCLING_TABLE_WRAP: regex matches optional surrounding `|`."""
    md = "| §§H2§§Section A§§/H2§§ |\nrow continues"
    out, counts = convert_sentinels_to_atx(md)
    assert "## Section A" in out
    assert counts == {"H2": 1}
    # No leftover pipe-wrapped sentinel
    assert "§§" not in out


def test_convert_sentinels_dedup_within_sheet_h1():
    """Merge-unroll can produce the same H1 line multiple times (same sheet name)."""
    md = "§§H1§§Sheet1§§/H1§§\n\n§§H1§§Sheet1§§/H1§§\n\nbody"
    out, counts = convert_sentinels_to_atx(md)
    assert counts == {"H1": 1}
    assert out.count("# Sheet1") == 1


def test_convert_sentinels_orphan_divider_removed():
    """ORPHAN_DIVIDER: a `|---|` row immediately following a stripped header is dropped."""
    md = "intro line\n\n| §§H2§§Section§§/H2§§ |\n|---|\n\nbody"
    out, _ = convert_sentinels_to_atx(md)
    assert "## Section" in out
    assert "|---|" not in out


def test_convert_sentinels_collapses_whitespace_inside_label():
    md = "§§H2§§Multi   space   label§§/H2§§"
    out, _ = convert_sentinels_to_atx(md)
    assert "## Multi space label" in out


def test_sentinel_regex_handles_optional_pipes_either_side():
    text = "| §§H4§§X§§/H4§§ |"
    m = SENTINEL_RE.search(text)
    assert m is not None
    assert m.group("lvl") == "H4"


def test_convert_preserves_mermaid_fence_adjacent_to_header():
    """SENTINEL_NEAR_MERMAID: ensure mermaid fences are not consumed when adjacent."""
    md = (
        "intro\n\n"
        "| §§H2§§Flow§§/H2§§ |\n\n"
        "```mermaid\n"
        "graph TD; A-->B;\n"
        "```\n\n"
        "post"
    )
    out, _ = convert_sentinels_to_atx(md)
    assert "## Flow" in out
    assert "```mermaid" in out
    assert "graph TD; A-->B;" in out


def test_round_trip_two_sheets_get_distinct_h1(tmp_xlsx_dir):
    """End-to-end pre-Docling on a 2-sheet workbook: each sheet's H1 anchor
    cell holds that sheet's display name."""
    wb = Workbook()
    ws1 = wb.active
    ws1.title = "Alpha"
    ws1["A1"] = "Title1"
    ws1["A1"].font = Font(size=14, bold=True)
    ws1["A2"] = "Sec"
    ws1["A2"].font = Font(size=11, bold=True)
    ws2 = wb.create_sheet("Beta")
    ws2["A1"] = "Title2"
    ws2["A1"].font = Font(size=14, bold=True)
    ws2["A2"] = "Sec"
    ws2["A2"].font = Font(size=11, bold=True)
    path = _save_workbook(wb, tmp_xlsx_dir)

    inject_header_sentinels(path)

    strings = _read_inline_strings(path)
    assert any("§§H1§§Alpha§§/H1§§" in s for s in strings)
    assert any("§§H1§§Beta§§/H1§§" in s for s in strings)