#!/usr/bin/env python3
"""Inspect or minimally update text cells in an existing .xlsx workbook.

The script edits OOXML members directly. It never opens and re-saves the
workbook through a spreadsheet library, so unrelated workbook features remain
untouched.
"""

from __future__ import annotations

import argparse
import copy
import html
import json
import os
from pathlib import Path
import re
import sys
import tempfile
from typing import Any
import xml.etree.ElementTree as ET
import zipfile


MAIN_NS = "http://schemas.openxmlformats.org/spreadsheetml/2006/main"
DOC_REL_NS = "http://schemas.openxmlformats.org/officeDocument/2006/relationships"
PKG_REL_NS = "http://schemas.openxmlformats.org/package/2006/relationships"
XML_NS = "http://www.w3.org/XML/1998/namespace"
CELL_REF_RE = re.compile(r"^[A-Z]{1,3}[1-9][0-9]*$")


class WorkbookError(RuntimeError):
    """Raised when the workbook cannot be edited without unsafe assumptions."""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Inspect or minimally update shared-string cells in an .xlsx template."
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    inspect_parser = subparsers.add_parser("inspect", help="List sheets, merges, and populated cells.")
    inspect_parser.add_argument("input", type=Path)
    inspect_parser.add_argument("--sheet", help="Inspect only this sheet; defaults to all sheets.")

    apply_parser = subparsers.add_parser("apply", help="Apply a JSON cell mapping and validate the result.")
    apply_parser.add_argument("input", type=Path)
    apply_parser.add_argument("output", type=Path)
    apply_parser.add_argument("--mapping", type=Path, required=True)
    apply_parser.add_argument("--sheet", help="Override the sheet name in the mapping file.")
    apply_parser.add_argument("--force", action="store_true", help="Replace an existing output file.")
    return parser.parse_args()


def require_xlsx(path: Path, must_exist: bool = True) -> Path:
    resolved = path.expanduser().resolve()
    if resolved.suffix.lower() != ".xlsx":
        raise WorkbookError(f"Only .xlsx files are supported: {resolved}")
    if must_exist and not resolved.is_file():
        raise WorkbookError(f"Workbook not found: {resolved}")
    return resolved


def read_zip_members(path: Path) -> tuple[list[zipfile.ZipInfo], dict[str, bytes]]:
    with zipfile.ZipFile(path, "r") as archive:
        bad_member = archive.testzip()
        if bad_member:
            raise WorkbookError(f"Corrupt ZIP member: {bad_member}")
        infos = archive.infolist()
        members = {info.filename: archive.read(info.filename) for info in infos}
    return infos, members


def parse_xml(data: bytes, member: str) -> ET.Element:
    try:
        return ET.fromstring(data)
    except ET.ParseError as exc:
        raise WorkbookError(f"Invalid XML in {member}: {exc}") from exc


def workbook_sheet_map(members: dict[str, bytes]) -> tuple[list[str], dict[str, str]]:
    workbook_member = "xl/workbook.xml"
    rels_member = "xl/_rels/workbook.xml.rels"
    if workbook_member not in members or rels_member not in members:
        raise WorkbookError("Workbook sheet metadata is incomplete.")

    workbook_root = parse_xml(members[workbook_member], workbook_member)
    rels_root = parse_xml(members[rels_member], rels_member)
    relationships = {
        rel.attrib["Id"]: rel.attrib["Target"]
        for rel in rels_root.findall(f"{{{PKG_REL_NS}}}Relationship")
    }

    ordered_names: list[str] = []
    mapping: dict[str, str] = {}
    for sheet in workbook_root.findall(f".//{{{MAIN_NS}}}sheet"):
        name = sheet.attrib["name"]
        rel_id = sheet.attrib.get(f"{{{DOC_REL_NS}}}id")
        if not rel_id or rel_id not in relationships:
            raise WorkbookError(f"Cannot resolve worksheet for sheet: {name}")
        target = relationships[rel_id].replace("\\", "/")
        if target.startswith("/"):
            member = target.lstrip("/")
        elif target.startswith("xl/"):
            member = target
        else:
            member = f"xl/{target}"
        ordered_names.append(name)
        mapping[name] = str(Path(member)).replace("\\", "/")
    return ordered_names, mapping


def shared_string_entries(members: dict[str, bytes]) -> tuple[ET.Element, list[ET.Element]]:
    member = "xl/sharedStrings.xml"
    if member not in members:
        raise WorkbookError("The workbook has no sharedStrings.xml; this safe editor will not guess a conversion.")
    root = parse_xml(members[member], member)
    entries = root.findall(f"{{{MAIN_NS}}}si")
    return root, entries


def entry_text(entry: ET.Element) -> str:
    return "".join(node.text or "" for node in entry.iter(f"{{{MAIN_NS}}}t"))


def cell_elements(sheet_root: ET.Element) -> dict[str, ET.Element]:
    cells: dict[str, ET.Element] = {}
    for cell in sheet_root.findall(f".//{{{MAIN_NS}}}c"):
        reference = cell.attrib.get("r")
        if reference:
            cells[reference] = cell
    return cells


def cell_shared_index(cell: ET.Element, reference: str) -> int:
    if cell.find(f"{{{MAIN_NS}}}f") is not None:
        raise WorkbookError(f"Refusing to replace formula cell: {reference}")
    if cell.attrib.get("t") != "s":
        raise WorkbookError(f"Target cell is not an existing shared-string cell: {reference}")
    value = cell.find(f"{{{MAIN_NS}}}v")
    if value is None or value.text is None or not value.text.isdigit():
        raise WorkbookError(f"Target cell has no valid shared-string index: {reference}")
    return int(value.text)


def inspect_workbook(path: Path, selected_sheet: str | None = None) -> dict[str, Any]:
    _, members = read_zip_members(path)
    sheet_names, sheet_map = workbook_sheet_map(members)
    _, strings = shared_string_entries(members)
    decoded = [entry_text(item) for item in strings]

    if selected_sheet and selected_sheet not in sheet_map:
        raise WorkbookError(f"Sheet not found: {selected_sheet}; available={sheet_names}")
    names = [selected_sheet] if selected_sheet else sheet_names
    sheets: list[dict[str, Any]] = []
    for name in names:
        member = sheet_map[name]
        root = parse_xml(members[member], member)
        values: dict[str, Any] = {}
        for reference, cell in cell_elements(root).items():
            cell_type = cell.attrib.get("t")
            value = cell.find(f"{{{MAIN_NS}}}v")
            formula = cell.find(f"{{{MAIN_NS}}}f")
            if formula is not None:
                values[reference] = {"formula": formula.text or ""}
            elif value is not None and value.text is not None:
                if cell_type == "s" and value.text.isdigit():
                    index = int(value.text)
                    values[reference] = decoded[index] if index < len(decoded) else {"invalid_shared_index": index}
                else:
                    values[reference] = value.text
        merges = [
            item.attrib["ref"]
            for item in root.findall(f".//{{{MAIN_NS}}}mergeCell")
            if "ref" in item.attrib
        ]
        sheets.append({"name": name, "member": member, "merges": merges, "cells": values})
    return {"file": str(path), "sheet_names": sheet_names, "sheets": sheets}


def load_mapping(path: Path) -> dict[str, Any]:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise WorkbookError(f"Cannot read mapping JSON {path}: {exc}") from exc
    if not isinstance(payload, dict) or not isinstance(payload.get("cells"), dict):
        raise WorkbookError("Mapping must be an object with a 'cells' object.")
    if not payload["cells"]:
        raise WorkbookError("Mapping contains no target cells.")
    for reference, spec in payload["cells"].items():
        if not isinstance(reference, str) or not CELL_REF_RE.fullmatch(reference):
            raise WorkbookError(f"Invalid cell reference in mapping: {reference!r}")
        if not isinstance(spec, (str, dict)):
            raise WorkbookError(f"Cell {reference} must map to a string or an instruction object.")
    return payload


def escaped_plain_si(text: str) -> str:
    escaped = html.escape(text, quote=False)
    preserve = text[:1].isspace() or text[-1:].isspace() if text else False
    attr = ' xml:space="preserve"' if preserve else ""
    return f"<si><t{attr}>{escaped}</t></si>"


def clone_with_replacements(entry: ET.Element, replacements: dict[str, str], reference: str) -> tuple[str, str]:
    if not replacements or not all(isinstance(k, str) and isinstance(v, str) for k, v in replacements.items()):
        raise WorkbookError(f"Cell {reference} replace_runs must be a non-empty string-to-string object.")
    cloned = copy.deepcopy(entry)
    original = entry_text(cloned)
    changed = False
    replacement_pattern = re.compile(
        "|".join(re.escape(key) for key in sorted(replacements, key=len, reverse=True))
    )
    for text_node in cloned.iter(f"{{{MAIN_NS}}}t"):
        value = text_node.text or ""
        updated = replacement_pattern.sub(lambda match: replacements[match.group(0)], value)
        if updated != value:
            changed = True
        value = updated
        text_node.text = value
    if not changed:
        raise WorkbookError(f"Cell {reference} replace_runs matched no rich-text run.")
    expected = entry_text(cloned)
    if expected == original:
        raise WorkbookError(f"Cell {reference} rich-text replacement produced no change.")
    ET.register_namespace("", MAIN_NS)
    xml = ET.tostring(cloned, encoding="unicode", short_empty_elements=True)
    return xml, expected


def replace_cell_index(sheet_xml: str, reference: str, new_index: int) -> str:
    cell_pattern = re.compile(
        rf'(<c\b[^>]*\br="{re.escape(reference)}"[^>]*>)(.*?)(</c>)',
        re.DOTALL,
    )
    match = cell_pattern.search(sheet_xml)
    if not match:
        if re.search(rf'<c\b[^>]*\br="{re.escape(reference)}"[^>]*/>', sheet_xml):
            raise WorkbookError(f"Target cell is self-closing and not a populated text input: {reference}")
        raise WorkbookError(f"Target cell not found in worksheet XML: {reference}")
    start, body, end = match.groups()
    if "<f" in body:
        raise WorkbookError(f"Refusing to replace formula cell: {reference}")
    if 't="s"' not in start:
        raise WorkbookError(f"Target cell is not a shared-string cell: {reference}")
    if re.search(r"<v>[^<]*</v>", body):
        body = re.sub(r"<v>[^<]*</v>", f"<v>{new_index}</v>", body, count=1)
    else:
        body = f"{body}<v>{new_index}</v>"
    return sheet_xml[: match.start()] + start + body + end + sheet_xml[match.end() :]


def xml_invariants(sheet_data: bytes) -> dict[str, Any]:
    root = parse_xml(sheet_data, "worksheet")
    return {
        "merges": [
            item.attrib.get("ref", "")
            for item in root.findall(f".//{{{MAIN_NS}}}mergeCell")
        ],
        "formulas": [item.text or "" for item in root.findall(f".//{{{MAIN_NS}}}f")],
        "drawing_ids": [
            item.attrib.get(f"{{{DOC_REL_NS}}}id", "")
            for item in root.findall(f".//{{{MAIN_NS}}}drawing")
        ],
    }


def apply_mapping(input_path: Path, output_path: Path, mapping_path: Path, sheet_override: str | None, force: bool) -> dict[str, Any]:
    if input_path == output_path:
        raise WorkbookError("Input and output must differ; preserve the source template.")
    if output_path.exists() and not force:
        raise WorkbookError(f"Output already exists; pass --force only after confirming the target: {output_path}")

    payload = load_mapping(mapping_path)
    infos, members = read_zip_members(input_path)
    sheet_names, sheet_map = workbook_sheet_map(members)
    sheet_name = sheet_override or payload.get("sheet") or sheet_names[0]
    if sheet_name not in sheet_map:
        raise WorkbookError(f"Sheet not found: {sheet_name}; available={sheet_names}")
    sheet_member = sheet_map[sheet_name]

    shared_root, shared_entries = shared_string_entries(members)
    shared_texts = [entry_text(item) for item in shared_entries]
    first_index_by_text: dict[str, int] = {}
    for index, value in enumerate(shared_texts):
        first_index_by_text.setdefault(value, index)

    sheet_root = parse_xml(members[sheet_member], sheet_member)
    cells = cell_elements(sheet_root)
    sheet_xml = members[sheet_member].decode("utf-8")
    appended: list[str] = []
    expected: dict[str, str] = {}

    for reference, spec in payload["cells"].items():
        if reference not in cells:
            raise WorkbookError(f"Target cell does not exist in template: {reference}")
        old_index = cell_shared_index(cells[reference], reference)
        if old_index >= len(shared_entries):
            raise WorkbookError(f"Cell {reference} points outside sharedStrings.xml: {old_index}")

        if isinstance(spec, str):
            expected_text = spec
            if expected_text in first_index_by_text:
                new_index = first_index_by_text[expected_text]
            else:
                new_index = len(shared_entries) + len(appended)
                appended.append(escaped_plain_si(expected_text))
                first_index_by_text[expected_text] = new_index
        else:
            unexpected = set(spec) - {"text", "replace_runs"}
            if unexpected:
                raise WorkbookError(f"Unsupported instruction keys for {reference}: {sorted(unexpected)}")
            if "replace_runs" in spec:
                if "text" in spec:
                    raise WorkbookError(f"Cell {reference} cannot use both text and replace_runs.")
                rich_xml, expected_text = clone_with_replacements(
                    shared_entries[old_index], spec["replace_runs"], reference
                )
                new_index = len(shared_entries) + len(appended)
                appended.append(rich_xml)
            elif isinstance(spec.get("text"), str):
                expected_text = spec["text"]
                if expected_text in first_index_by_text:
                    new_index = first_index_by_text[expected_text]
                else:
                    new_index = len(shared_entries) + len(appended)
                    appended.append(escaped_plain_si(expected_text))
                    first_index_by_text[expected_text] = new_index
            else:
                raise WorkbookError(f"Cell {reference} instruction must contain text or replace_runs.")

        expected[reference] = expected_text
        sheet_xml = replace_cell_index(sheet_xml, reference, new_index)

    shared_xml = members["xl/sharedStrings.xml"].decode("utf-8")
    if appended:
        insertion = "\n  " + "\n  ".join(appended)
        if "</sst>" not in shared_xml:
            raise WorkbookError("sharedStrings.xml has no closing sst element.")
        shared_xml = shared_xml.replace("</sst>", f"{insertion}\n</sst>", 1)
        new_unique_count = len(shared_entries) + len(appended)
        if re.search(r'\buniqueCount="[0-9]+"', shared_xml):
            shared_xml = re.sub(
                r'\buniqueCount="[0-9]+"',
                f'uniqueCount="{new_unique_count}"',
                shared_xml,
                count=1,
            )
        else:
            shared_xml = shared_xml.replace("<sst ", f'<sst uniqueCount="{new_unique_count}" ', 1)

    replacements = {
        sheet_member: sheet_xml.encode("utf-8"),
        "xl/sharedStrings.xml": shared_xml.encode("utf-8"),
    }
    before_invariants = xml_invariants(members[sheet_member])

    output_path.parent.mkdir(parents=True, exist_ok=True)
    temp_handle = tempfile.NamedTemporaryFile(
        prefix=f".{output_path.stem}.", suffix=".xlsx", dir=output_path.parent, delete=False
    )
    temp_path = Path(temp_handle.name)
    temp_handle.close()
    try:
        with zipfile.ZipFile(temp_path, "w") as output_zip:
            for info in infos:
                output_zip.writestr(info, replacements.get(info.filename, members[info.filename]))
        _, output_members = read_zip_members(temp_path)
        if list(members) != list(output_members):
            raise WorkbookError("Output ZIP members differ from the source workbook.")
        allowed_changes = {sheet_member, "xl/sharedStrings.xml"}
        changed_unexpectedly = [
            name for name, data in members.items()
            if name not in allowed_changes and output_members[name] != data
        ]
        if changed_unexpectedly:
            raise WorkbookError(f"Untargeted workbook members changed: {changed_unexpectedly}")
        after_invariants = xml_invariants(output_members[sheet_member])
        if before_invariants != after_invariants:
            raise WorkbookError("Worksheet formulas, merge ranges, or drawing relationships changed.")

        inspected = inspect_workbook(temp_path, sheet_name)
        output_cells = inspected["sheets"][0]["cells"]
        mismatches = {
            reference: {"expected": value, "actual": output_cells.get(reference)}
            for reference, value in expected.items()
            if output_cells.get(reference) != value
        }
        if mismatches:
            raise WorkbookError(f"Output cell verification failed: {mismatches}")
        os.replace(temp_path, output_path)
    finally:
        if temp_path.exists():
            temp_path.unlink()

    return {
        "status": "success",
        "input": str(input_path),
        "output": str(output_path),
        "sheet": sheet_name,
        "changed_cells": sorted(expected),
        "appended_shared_strings": len(appended),
        "validations": {
            "zip_integrity": True,
            "sheet_names_preserved": inspected["sheet_names"] == sheet_names,
            "zip_members_preserved": True,
            "untargeted_members_byte_identical": True,
            "formulas_merges_drawings_preserved": True,
            "target_values_verified": True,
        },
    }


def main() -> int:
    args = parse_args()
    try:
        if args.command == "inspect":
            input_path = require_xlsx(args.input)
            result = inspect_workbook(input_path, args.sheet)
        else:
            input_path = require_xlsx(args.input)
            output_path = require_xlsx(args.output, must_exist=False)
            mapping_path = args.mapping.expanduser().resolve()
            result = apply_mapping(
                input_path,
                output_path,
                mapping_path,
                args.sheet,
                args.force,
            )
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0
    except WorkbookError as exc:
        print(json.dumps({"status": "error", "message": str(exc)}, ensure_ascii=False), file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())