"""ks-xlsx-parser -> Markdown, with mermaid newline restoration.

Parses an .xls/.xlsx with `ks_xlsx_parser`, concatenates each chunk's
`render_html`, converts to Markdown via `markdownify`, and restores any
multi-line mermaid blocks that were collapsed during the round-trip.

Outputs in CWD:
    <base>.ks.md
"""
from __future__ import annotations

import argparse
import logging
import os
import re
import shutil
import subprocess
import sys
import tempfile
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path
from typing import List, Tuple

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

from src.extraction_v2.xlsx_header_injection import (
    NS,
    _load_shared_strings,
    _read_cell_text,
)

logger = logging.getLogger("inject_headers_and_parse_ks")


# ---------------------------------------------------------------------------
# .xls -> .xlsx (LibreOffice headless)
# ---------------------------------------------------------------------------

def convert_xls_to_xlsx(xls_path: str, outdir: str) -> str:
    soffice = shutil.which("soffice") or shutil.which("libreoffice")
    if not soffice:
        raise RuntimeError("LibreOffice (soffice) not found; required to convert .xls")
    os.makedirs(outdir, exist_ok=True)
    proc = subprocess.run(
        [soffice, "--headless", "--convert-to", "xlsx", "--outdir", outdir, xls_path],
        capture_output=True, text=True, timeout=300,
    )
    if proc.returncode != 0:
        raise RuntimeError(f"soffice conversion failed:\n{proc.stderr}")
    converted = os.path.join(outdir, Path(xls_path).stem + ".xlsx")
    if not os.path.exists(converted):
        raise RuntimeError(f"Expected converted file not found: {converted}")
    return converted


# ---------------------------------------------------------------------------
# Mermaid block collection + restoration
# ---------------------------------------------------------------------------

def collect_mermaid_blocks(xlsx_path: str) -> List[str]:
    workdir = tempfile.mkdtemp(prefix="mermaid_collect_")
    try:
        with zipfile.ZipFile(xlsx_path) as z:
            z.extractall(workdir)

        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 []

        out: List[str] = []
        ws_dir = os.path.join(workdir, "xl/worksheets")
        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"):
                    text = _read_cell_text(c, shared)
                    if text and text.lstrip().startswith("```mermaid"):
                        out.append(text)
        return out
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


def restore_mermaid_blocks(md: str, mermaid_texts: List[str]) -> Tuple[str, int]:
    restored = 0
    for original in mermaid_texts:
        tokens = original.split()
        if not tokens:
            continue
        flex = r"\s+".join(re.escape(t) for t in tokens)
        pattern = r"\|?\s*" + flex + r"\s*\|?"
        replacement = "\n\n" + original + "\n\n"
        new_md, n = re.subn(pattern, replacement, md, count=1)
        if n == 0:
            logger.warning("could not relocate mermaid block (%d chars)", len(original))
            continue
        md = new_md
        restored += 1
    return md, restored


# ---------------------------------------------------------------------------
# Parsing
# ---------------------------------------------------------------------------

def ksparser_to_md(xlsx_path: str) -> str:
    from ks_xlsx_parser import parse_workbook
    from markdownify import markdownify

    result = parse_workbook(path=xlsx_path)
    html_body = "\n".join(chunk.render_html for chunk in result.chunks)
    return markdownify(html_body, heading_style="ATX")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("xlsx", help="Input .xls/.xlsx/.xlsm")
    ap.add_argument("-v", "--verbose", action="store_true")
    args = ap.parse_args()
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    src = args.xlsx
    if not os.path.exists(src):
        logger.error("input not found: %s", src)
        return 2

    xls_tmp_dir: str | None = None
    if src.lower().endswith(".xls"):
        xls_tmp_dir = tempfile.mkdtemp(prefix="xls2xlsx_")
        logger.info("converting legacy .xls -> .xlsx via LibreOffice")
        src = convert_xls_to_xlsx(src, xls_tmp_dir)
        logger.info("converted: %s", src)
    elif not src.lower().endswith((".xlsx", ".xlsm")):
        logger.error("input must be .xls/.xlsx/.xlsm")
        return 2

    base_name = Path(src).stem
    if base_name.endswith(".preprocessed"):
        base_name = base_name[: -len(".preprocessed")]

    out_md = os.path.join(os.getcwd(), f"{base_name}.ks.md")

    mermaid_blocks = collect_mermaid_blocks(src)
    logger.info("found %d mermaid block(s) embedded in xlsx", len(mermaid_blocks))

    md_raw = ksparser_to_md(src)
    md, n_restored = restore_mermaid_blocks(md_raw, mermaid_blocks)
    logger.info("restored mermaid newlines for %d/%d block(s)", n_restored, len(mermaid_blocks))

    with open(out_md, "w", encoding="utf-8") as fh:
        fh.write(md)
    logger.info("wrote %s (%d chars)", out_md, len(md))

    if xls_tmp_dir:
        shutil.rmtree(xls_tmp_dir, ignore_errors=True)

    print()
    print(f"Markdown: {out_md}")
    return 0


if __name__ == "__main__":
    sys.exit(main())