import os
import shutil
import csv
import json
from collections import OrderedDict
from openpyxl import load_workbook, Workbook

# --- CONFIG ---

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Input: all normalized mushaf Excel files in workspace root
INPUT_XLSX_DIR = BASE_DIR

# Only process canonical normalized sources (prevents accidental extra outputs)
NORMALIZED_SUFFIX = "_normalized.xlsx"

# Output root with fixed format subfolders under out/: json, csv, excel
OUT_DIR = os.path.join(BASE_DIR, "out")
JSON_OUT_DIR = os.path.join(OUT_DIR, "json")
CSV_OUT_DIR = os.path.join(OUT_DIR, "csv")
EXCEL_OUT_DIR = os.path.join(OUT_DIR, "excel")

# 28-letter normalized alphabet (Clear Proof spec) WITH transliterations
LETTER_MAP = OrderedDict([
    ("ا", "alif"),
    ("ب", "ba"),
    ("ت", "ta"),
    ("ث", "tha"),
    ("ج", "jim"),
    ("ح", "ha"),
    ("خ", "kha"),
    ("د", "dal"),
    ("ذ", "dhal"),
    ("ر", "ra"),
    ("ز", "zay"),
    ("س", "sin"),
    ("ش", "shin"),
    ("ص", "sad"),
    ("ض", "dad"),
    ("ط", "ta'"),
    ("ظ", "za'"),
    ("ع", "ayn"),
    ("غ", "ghayn"),
    ("ف", "fa"),
    ("ق", "qaf"),
    ("ك", "kaf"),
    ("ل", "lam"),
    ("م", "mim"),
    ("ن", "nun"),
    ("ه", "ha'"),
    ("و", "waw"),
    ("ي", "ya"),
])

ARABIC_LETTERS = list(LETTER_MAP.keys())

README_PATH = os.path.join(BASE_DIR, "README.md")


def slug_from_filename(filename: str) -> str:
    """
    Keep output names aligned with source file names.
    Example: 'Hafs_normalized.xlsx' -> 'Hafs_normalized'
    """
    return os.path.splitext(os.path.basename(filename))[0]


def ensure_dir(path: str):
    if not os.path.exists(path):
        os.makedirs(path)


def sync_output_dirs():
    """
    Keep OUT_DIR reproducible by allowing only fixed format folders:
    out/json, out/csv, out/excel.
    """
    ensure_dir(OUT_DIR)
    ensure_dir(JSON_OUT_DIR)
    ensure_dir(CSV_OUT_DIR)
    ensure_dir(EXCEL_OUT_DIR)

    expected = {"json", "csv", "excel"}
    existing = {
        name
        for name in os.listdir(OUT_DIR)
        if os.path.isdir(os.path.join(OUT_DIR, name))
    }

    stale = sorted(existing - expected)
    for name in stale:
        stale_path = os.path.join(OUT_DIR, name)
        shutil.rmtree(stale_path)
        print(f"[CLEAN] Removed stale output folder: {stale_path}")


def discover_root_source_xlsx_files():
    """
    Discover reproducible source inputs from workspace root.
    Rules:
      - filename ends with *_normalized.xlsx
      - exclude Excel lock/temp files that start with '~$'
    """
    candidates = []
    for fn in os.listdir(INPUT_XLSX_DIR):
        if fn.startswith("~$"):
            continue
        if not fn.lower().endswith(NORMALIZED_SUFFIX):
            continue
        full_path = os.path.join(INPUT_XLSX_DIR, fn)
        if os.path.isfile(full_path):
            candidates.append(full_path)
    return sorted(candidates)


def build_letter_matrix_for_mushaf(xlsx_path: str):
    """
    Read one normalized mushaf xlsx file and build:
      - rows: list of { 'sura', 'verse', 'text', <letter>:count, ... }
      - totals: OrderedDict(letter -> total count across all verses)

    Expected columns: sura, verse, text (case-insensitive).
    """
    rows = []
    totals = OrderedDict((ltr, 0) for ltr in ARABIC_LETTERS)

    wb = load_workbook(xlsx_path, read_only=True, data_only=True)
    ws = wb[wb.sheetnames[0]]

    headers = next(ws.iter_rows(min_row=1, max_row=1, values_only=True), None)
    if not headers:
        return rows, totals

    normalized_headers = [str(h).strip().lower() if h is not None else "" for h in headers]
    try:
        sura_idx = normalized_headers.index("sura")
        verse_idx = normalized_headers.index("verse")
        text_idx = normalized_headers.index("text")
    except ValueError as exc:
        raise RuntimeError(
            f"Missing required columns in {xlsx_path}. Expected sura, verse, text."
        ) from exc

    for record in ws.iter_rows(min_row=2, values_only=True):
        if record is None:
            continue

        try:
            sura_raw = record[sura_idx]
            verse_raw = record[verse_idx]
            text_raw = record[text_idx]
        except IndexError:
            continue

        if sura_raw is None or verse_raw is None:
            continue

        try:
            sura = int(sura_raw)
            verse = int(verse_raw)
        except (TypeError, ValueError):
            continue

        text = "" if text_raw is None else str(text_raw)

        # Initialize counts for this verse
        verse_counts = {ltr: 0 for ltr in ARABIC_LETTERS}

        # Count the letters in the text
        for ch in text:
            if ch in verse_counts:
                verse_counts[ch] += 1
                totals[ch] += 1

        # Build row dict
        row = {
            "sura": sura,
            "verse": verse,
            "text": text,
        }
        row.update(verse_counts)
        rows.append(row)

    return rows, totals


def write_letter_matrix_csv(slug: str, rows, totals):
    """
    Write per-verse matrix for one mushaf:
    out/csv/<slug>_LETTER_MATRIX.csv

    Columns: sura, verse, text, then 28 letters as "ا (alif)", etc.
    Last row: TOTAL row with column sums for each letter.
    """
    ensure_dir(CSV_OUT_DIR)

    # Build the letter-column headers as "ا (alif)", etc.
    letter_headers = [f"{arabic} ({trans})" for arabic, trans in LETTER_MAP.items()]

    fieldnames = ["sura", "verse", "text"] + letter_headers

    csv_path = os.path.join(CSV_OUT_DIR, f"{slug}_LETTER_MATRIX.csv")

    with open(csv_path, "w", encoding="utf-8", newline="") as f:
        writer = csv.writer(f)
        # Write header row manually (DictWriter + mapped keys would be messier here)
        writer.writerow(fieldnames)

        # For each verse row, emit counts in the same order
        for row in rows:
            base = [row["sura"], row["verse"], row["text"]]
            counts = [row[arabic] for arabic in ARABIC_LETTERS]
            writer.writerow(base + counts)

        # Write totals row
        total_base = ["", "TOTAL", ""]
        total_counts = [totals[arabic] for arabic in ARABIC_LETTERS]
        writer.writerow(total_base + total_counts)

    return csv_path


def write_letter_matrix_json(slug: str, rows, totals):
    """
    Write per-verse matrix for one mushaf as JSON:
      out/json/<slug>_LETTER_MATRIX.json
    """
    ensure_dir(JSON_OUT_DIR)
    json_path = os.path.join(JSON_OUT_DIR, f"{slug}_LETTER_MATRIX.json")

    payload = {
        "letters": [
            {"arabic": arabic, "transliteration": trans}
            for arabic, trans in LETTER_MAP.items()
        ],
        "rows": rows,
        "totals": dict(totals),
        "total_letters": int(sum(totals.values())),
    }

    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)

    return json_path


def write_letter_matrix_xlsx(slug: str, rows, totals):
    """
    Write per-verse matrix for one mushaf as Excel:
      out/excel/<slug>_LETTER_MATRIX.xlsx
    """
    ensure_dir(EXCEL_OUT_DIR)
    xlsx_path = os.path.join(EXCEL_OUT_DIR, f"{slug}_LETTER_MATRIX.xlsx")

    wb = Workbook()
    ws = wb.active
    ws.title = "LETTER_MATRIX"

    letter_headers = [f"{arabic} ({trans})" for arabic, trans in LETTER_MAP.items()]
    ws.append(["sura", "verse", "text"] + letter_headers)

    for row in rows:
        base = [row["sura"], row["verse"], row["text"]]
        counts = [row[arabic] for arabic in ARABIC_LETTERS]
        ws.append(base + counts)

    ws.append(["", "TOTAL", ""] + [totals[arabic] for arabic in ARABIC_LETTERS])
    wb.save(xlsx_path)

    return xlsx_path


def write_workspace_readme():
        """
        Keep README in sync with the pipeline and audit utilities.
        """
        readme_text = """# Letter TALLY PIPELINE

This workspace includes:

- A per-verse Arabic letter tally engine for normalized sources.
- A reproducible workspace audit utility (registry + hashes + TXT validation report).

## Letter Matrix Engine

- Script: `letter_tally_engine.py`
- Input discovery:
    - Scans workspace root for `*_normalized.xlsx`.
    - Ignores Excel lock files starting with `~$`.
- Outputs per source slug:
    - `out/csv/<slug>_LETTER_MATRIX.csv`
    - `out/json/<slug>_LETTER_MATRIX.json`
    - `out/excel/<slug>_LETTER_MATRIX.xlsx`

Run:

```powershell
c:/Users/wolfk/Desktop/Letter TALLY PIPELINE/.venv/Scripts/python.exe letter_tally_engine.py
```

## Reproducible Audit Utility

- Script: `workspace_registry_hash_audit.py`
- Default outputs:
    - `out/audit/workspace_registry.csv`
    - `out/audit/workspace_hash_manifest.sha256`
    - `out/audit/workspace_validation_report.txt`

Run:

```powershell
c:/Users/wolfk/Desktop/Letter TALLY PIPELINE/.venv/Scripts/python.exe workspace_registry_hash_audit.py
```

Optional arguments:

```text
--root <path>           Workspace root to audit (default: script directory)
--output-dir <path>     Output directory (default: out/audit)
--exclude-dir <name>    Additional directory name to exclude (repeatable)
```

## Notes

- `README.md` is auto-generated by `letter_tally_engine.py` each time it runs.
- If you manually edit `README.md`, rerunning the engine will overwrite it.
"""

        with open(README_PATH, "w", encoding="utf-8") as f:
                f.write(readme_text)


def main():
    if not os.path.isdir(INPUT_XLSX_DIR):
        raise RuntimeError(f"Input folder not found: {INPUT_XLSX_DIR}")

    xlsx_files = discover_root_source_xlsx_files()

    if not xlsx_files:
        raise RuntimeError(
            f"No source files matching *{NORMALIZED_SUFFIX} found in {INPUT_XLSX_DIR}"
        )

    sync_output_dirs()

    for xlsx_path in sorted(xlsx_files):
        slug = slug_from_filename(xlsx_path)
        print(f"[INFO] Building letter matrix for {slug} from {xlsx_path} ...")

        rows, totals = build_letter_matrix_for_mushaf(xlsx_path)
        csv_path = write_letter_matrix_csv(slug, rows, totals)
        json_path = write_letter_matrix_json(slug, rows, totals)
        out_xlsx_path = write_letter_matrix_xlsx(slug, rows, totals)

        total_letters = sum(totals.values())
        print(f"[OK] {slug}: verses = {len(rows)}, total letters = {total_letters}")
        print(f"[OK] CSV -> {csv_path}")
        print(f"[OK] JSON -> {json_path}")
        print(f"[OK] XLSX -> {out_xlsx_path}")

    write_workspace_readme()
    print(f"[OK] README -> {README_PATH}")

    print("Done. All per-verse letter matrices generated in csv/json/xlsx.")


if __name__ == "__main__":
    main()
