#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Invariant normalization engine for ornamentation Excel files.

Goal
----
Read ornamentation .xlsx workbooks and generate ONLY:
- a normalized .xlsx workbook
- a normalized .csv file
- a normalized .json file

Output schema is STRICTLY:
    sura | verse | text

No TXT output is produced.

Locked behavior
---------------
1. Read first worksheet only.
2. Expect first three header cells to represent sura / verse / text.
3. Read columns 1, 2, 3 only.
4. Normalize column 3 only.
5. Preserve columns 1, 2, 3 in output. No extra columns.
6. Hard-fail on structural anomalies or post-normalization unknown characters.

Normalization performed on text column
--------------------------------------
- Unicode NFKC normalization
- Strip BOM / invisible formatting chars
- Convert all odd whitespace to normal spaces
- Strip leading/trailing spaces
- Collapse repeated spaces
- Remove trailing Arabic-Indic / ASCII verse numerals at end of field
- Character mapping:
    أ إ آ ٱ -> ا
    ى ی ے ۓ ێ ۍ ئ -> ي
    ؤ -> و
    ة -> ه
    ٮ -> ب
    ] - -> delete
    ء -> delete
    tatweel -> delete
- Remove all combining marks (Mn, Mc, Me) including tashkil / recitation marks
- Final hard gate:
    only 28 canonical letters + single spaces allowed in text
    known forbidden offenders: ٮ ] -

Usage
-----
python normalization_engine.py --in-dir RAW_XLSX --out-dir NORMALIZED --fail-fast
python normalization_engine.py --in-file "Hafs_Ornamentation.xlsx" --out-dir NORMALIZED
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import re
import unicodedata
from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Tuple

from openpyxl import Workbook, load_workbook


CANONICAL_28 = set("ابتثجحخدذرزسشصضطظعغفقكلمنهوي")
ALLOWED_TEXT_CHARS = CANONICAL_28 | {" "}
FORBIDDEN_CHARS = {"ٮ", "]", "-"}

HEADER_TEXT_VALUES = {"text", "arabic text"}

# Character mappings locked by contract
CHAR_MAP = {
    "أ": "ا",
    "إ": "ا",
    "آ": "ا",
    "ٱ": "ا",
    "ى": "ي",
    "ی": "ي",
    "ے": "ي",
    "ۓ": "ي",
    "ێ": "ي",
    "ۍ": "ي",
    "ئ": "ي",
    "ؤ": "و",
    "ة": "ه",
    "ٮ": "ب",
    "]": "",
    "-": "",
    "ء": "",
    "ـ": "",
    "ۥ": "",
    "ۦ": "",
    "۞": "",
    "۩": "",
}

INVISIBLE_DROP = {
    "\ufeff", "\u200b", "\u200c", "\u200d", "\u200e", "\u200f",
    "\u202a", "\u202b", "\u202c", "\u202d", "\u202e", "\u2066",
    "\u2067", "\u2068", "\u2069",
}

RE_MULTI_SPACES = re.compile(r" {2,}")
RE_TRAILING_VERSE_NUMERALS = re.compile(r"[ \u00A0]*[٠-٩0-9]+[ \u00A0]*$")
RE_ORNAMENTATION_TOKEN = re.compile(r"(?i)(^|[\s_\-]+)ornamentation(?=$|[\s_\-]+)")
RE_MULTI_UNDERSCORES = re.compile(r"_+")
RE_VARIANT_SUFFIX = re.compile(r"(\([^)]+\))$")


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def ulabel(ch: str) -> str:
    return f"U+{ord(ch):04X} {unicodedata.name(ch, '<no name>')} [{unicodedata.category(ch)}]"


def output_stem_from_path(path: Path) -> str:
    stem = path.stem
    # Extract trailing (variant) suffix e.g. (Medina I) before cleaning
    m = RE_VARIANT_SUFFIX.search(stem)
    variant = m.group(1) if m else ""
    if variant:
        stem = stem[:m.start()]
    stem = RE_ORNAMENTATION_TOKEN.sub(" ", stem)
    stem = stem.replace("-", " ").replace("_", " ")
    stem = " ".join(stem.split())
    stem = stem.replace(" ", "_")
    stem = RE_MULTI_UNDERSCORES.sub("_", stem).strip("_")
    stem = stem or "output"
    if variant:
        stem = stem + variant
    return stem


def header_is_valid(a, b, c) -> bool:
    a = str(a).strip().lower() if a is not None else ""
    b = str(b).strip().lower() if b is not None else ""
    c = str(c).strip().lower() if c is not None else ""
    return a == "sura" and b == "verse" and c in HEADER_TEXT_VALUES


@dataclass
class Audit:
    file: str
    rows_read: int = 0
    rows_written: int = 0
    mapped_counts: Counter = field(default_factory=Counter)
    removed_invisible_counts: Counter = field(default_factory=Counter)
    removed_combining_counts: Counter = field(default_factory=Counter)
    removed_trailing_numerals: int = 0
    whitespace_normalizations: int = 0
    final_fail_counts: Counter = field(default_factory=Counter)
    final_fail_examples: Dict[str, List[dict]] = field(default_factory=dict)


def add_fail_example(audit: Audit, ch: str, sura, verse, text):
    key = ulabel(ch)
    audit.final_fail_examples.setdefault(key, [])
    if len(audit.final_fail_examples[key]) < 5:
        audit.final_fail_examples[key].append({
            "sura": sura,
            "verse": verse,
            "text_excerpt": text[:120]
        })


def normalize_text(text: str, audit: Audit) -> str:
    # Stage 1: Unicode hardening
    text = unicodedata.normalize("NFKC", text)

    # Stage 2: drop invisible formatting
    out = []
    for ch in text:
        if ch in INVISIBLE_DROP:
            audit.removed_invisible_counts[ulabel(ch)] += 1
            continue
        out.append(ch)
    text = "".join(out)

    # Stage 3: normalize whitespace chars to plain space
    chars = []
    replaced_any = False
    for ch in text:
        if ch == "\n" or ch == "\r":
            # normalize line breaks inside cells into spaces; cells should be one logical row
            chars.append(" ")
            replaced_any = True
        elif ch.isspace():
            if ch != " ":
                replaced_any = True
            chars.append(" ")
        else:
            chars.append(ch)
    text = "".join(chars)

    # Stage 4: remove trailing verse numerals only at end of field
    text2, n = RE_TRAILING_VERSE_NUMERALS.subn("", text)
    if n:
        audit.removed_trailing_numerals += n
        text = text2

    # Stage 5: char mapping
    mapped_out = []
    for ch in text:
        if ch in CHAR_MAP:
            mapped = CHAR_MAP[ch]
            audit.mapped_counts[f"{ulabel(ch)} -> {repr(mapped)}"] += 1
            mapped_out.append(mapped)
        else:
            mapped_out.append(ch)
    text = "".join(mapped_out)

    # Stage 6: remove all combining marks (tashkil / recitation / ornamentation)
    out = []
    for ch in text:
        if unicodedata.category(ch) in {"Mn", "Mc", "Me"}:
            audit.removed_combining_counts[ulabel(ch)] += 1
            continue
        out.append(ch)
    text = "".join(out)

    # Stage 7: trim and collapse spaces
    stripped = text.strip(" ")
    if stripped != text:
        replaced_any = True
    text = stripped

    text_new, n = RE_MULTI_SPACES.subn(" ", text)
    if n:
        replaced_any = True
    text = text_new

    if replaced_any:
        audit.whitespace_normalizations += 1

    return text


def final_gate(text: str, audit: Audit, sura, verse) -> None:
    if text == "":
        raise ValueError(f"Empty normalized text at sura={sura}, verse={verse}")

    for ch in text:
        if ch in FORBIDDEN_CHARS or ch not in ALLOWED_TEXT_CHARS:
            audit.final_fail_counts[ulabel(ch)] += 1
            add_fail_example(audit, ch, sura, verse, text)

    if audit.final_fail_counts:
        raise ValueError(f"Final gate failure in file {audit.file}")


def normalize_workbook(path: Path, out_dir: Path, audit_dir: Path, error_dir: Path) -> None:
    wb = load_workbook(path, read_only=True, data_only=True)
    ws = wb[wb.sheetnames[0]]

    header = next(ws.iter_rows(min_row=1, max_row=1, values_only=True))
    h0, h1, h2 = (header[0] if len(header) > 0 else None,
                  header[1] if len(header) > 1 else None,
                  header[2] if len(header) > 2 else None)
    if not header_is_valid(h0, h1, h2):
        raise ValueError(f"{path.name}: unexpected header {header[:3]}")

    audit = Audit(file=path.name)
    rows_out: List[Tuple[int, int, str]] = []

    for row_idx, row in enumerate(ws.iter_rows(min_row=2, values_only=True), start=2):
        sura = row[0] if len(row) > 0 else None
        verse = row[1] if len(row) > 1 else None
        text = row[2] if len(row) > 2 else None

        if sura is None and verse is None and (text is None or str(text).strip() == ""):
            continue

        audit.rows_read += 1

        # Structural validation
        if sura is None or verse is None:
            raise ValueError(f"{path.name}: missing sura/verse at row {row_idx}")
        if text is None:
            raise ValueError(f"{path.name}: missing text at row {row_idx}")

        try:
            sura_int = int(sura)
            verse_int = int(verse)
        except Exception:
            raise ValueError(f"{path.name}: non-integer sura/verse at row {row_idx}: {sura!r}, {verse!r}")

        normalized = normalize_text(str(text), audit)
        final_gate(normalized, audit, sura_int, verse_int)

        rows_out.append((sura_int, verse_int, normalized))
        audit.rows_written += 1

    stem = output_stem_from_path(path)
    xlsx_out = out_dir / "excel" / f"{stem}_normalized.xlsx"
    csv_out  = out_dir / "csv"   / f"{stem}_normalized.csv"
    json_out = out_dir / "json"  / f"{stem}_normalized.json"
    audit_out = audit_dir / f"{stem}_normalized.audit.json"

    (out_dir / "excel").mkdir(parents=True, exist_ok=True)
    (out_dir / "csv").mkdir(parents=True, exist_ok=True)
    (out_dir / "json").mkdir(parents=True, exist_ok=True)
    audit_dir.mkdir(parents=True, exist_ok=True)
    error_dir.mkdir(parents=True, exist_ok=True)

    # Write normalized Excel
    out_wb = Workbook()
    out_ws = out_wb.active
    out_ws.title = "Sheet1"
    out_ws.append(["sura", "verse", "text"])
    for sura_int, verse_int, normalized in rows_out:
        out_ws.append([sura_int, verse_int, normalized])
    out_wb.save(xlsx_out)

    # Write normalized CSV
    with open(csv_out, "w", encoding="utf-8-sig", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["sura", "verse", "text"])
        writer.writerows(rows_out)

    # Write normalized JSON
    with open(json_out, "w", encoding="utf-8") as f:
        json.dump([
            {"sura": sura_int, "verse": verse_int, "text": normalized}
            for sura_int, verse_int, normalized in rows_out
        ], f, ensure_ascii=False, indent=2)

    # Write audit JSON
    with open(audit_out, "w", encoding="utf-8") as f:
        json.dump({
            "file": audit.file,
            "rows_read": audit.rows_read,
            "rows_written": audit.rows_written,
            "mapped_counts": dict(audit.mapped_counts),
            "removed_invisible_counts": dict(audit.removed_invisible_counts),
            "removed_combining_counts": dict(audit.removed_combining_counts),
            "removed_trailing_numerals": audit.removed_trailing_numerals,
            "whitespace_normalizations": audit.whitespace_normalizations,
            "outputs": {
                "xlsx": str(xlsx_out),
                "csv": str(csv_out),
                "json": str(json_out),
            }
        }, f, ensure_ascii=False, indent=2)

    print(f"[OK] normalized {path.name}")


DEFAULT_EXCLUDED_DIRS = {
    ".git", ".venv", "venv", "__pycache__",
    "unicode_reports", "normalized_outputs_all_with_json",
}


def is_excluded_excel_path(path: Path, root: Path, extra_excluded_dirs: set[str] | None = None) -> bool:
    """Return True for generated/temp files that should not be treated as inputs."""
    excluded = set(DEFAULT_EXCLUDED_DIRS)
    if extra_excluded_dirs:
        excluded.update(extra_excluded_dirs)

    # Skip Excel temp/lock files.
    if path.name.startswith("~$"):
        return True

    try:
        rel_parts = path.relative_to(root).parts
    except ValueError:
        rel_parts = path.parts

    return any(part in excluded for part in rel_parts[:-1])


def collect_input_files(in_file: str | None, in_dir: str | None, extra_excluded_dirs: set[str] | None = None) -> List[Path]:
    files: List[Path] = []
    root = Path(in_dir).resolve() if in_dir else Path.cwd().resolve()

    if in_file:
        files.append(Path(in_file))

    if in_dir:
        # Recursive discovery: find Excel files anywhere under the input directory.
        files.extend(sorted(root.rglob("*.xlsx")))

    files = [p for p in files if p.exists() and not is_excluded_excel_path(p.resolve(), root, extra_excluded_dirs)]

    if not files:
        raise RuntimeError("No input .xlsx files found.")

    # deduplicate while preserving order
    seen = set()
    uniq = []
    for p in files:
        resolved = p.resolve()
        if resolved not in seen:
            uniq.append(resolved)
            seen.add(resolved)
    return uniq


def write_error(path: Path, file_name: str, exc: Exception) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump({
            "file": file_name,
            "error": str(exc),
        }, f, ensure_ascii=False, indent=2)


def main() -> None:
    ap = argparse.ArgumentParser(description="Normalize ornamentation Excel workbooks into canonical Excel + CSV only.")
    ap.add_argument("--in-file", help="Single input .xlsx file")
    ap.add_argument("--in-dir", help="Directory of input .xlsx files")
    ap.add_argument("--out-dir", required=True, help="Directory for normalized outputs")
    ap.add_argument("--audit-dir", help="Directory for audit JSON files")
    ap.add_argument("--error-dir", help="Directory for error JSON files")
    ap.add_argument("--fail-fast", action="store_true", help="Stop on first failure")
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    audit_dir = Path(args.audit_dir) if args.audit_dir else out_dir / "_audit"
    error_dir = Path(args.error_dir) if args.error_dir else out_dir / "_errors"

    extra_excluded_dirs = {out_dir.name, audit_dir.name, error_dir.name}
    files = collect_input_files(args.in_file, args.in_dir, extra_excluded_dirs=extra_excluded_dirs)

    failures = 0
    for path in files:
        try:
            normalize_workbook(path, out_dir, audit_dir, error_dir)
        except Exception as exc:
            failures += 1
            stem = output_stem_from_path(path)
            write_error(error_dir / f"{stem}_normalized.error.json", path.name, exc)
            print(f"[FAIL] {path.name} :: {exc}")
            if args.fail_fast:
                raise

    if failures:
        raise SystemExit(2)


if __name__ == "__main__":
    main()
