#!/usr/bin/env python3
"""
14 Initial Engine

Read-only extractor and counter for normalized mushaf CSV/XLSX source files.
See README.md for full usage, input format, and output schema documentation.

Outputs (written to 14_initial_engine_outputs/):
    json/<mushaf>_initial_counts.json
    csv/<mushaf>_initial_counts.csv
    excel/<mushaf>_initial_counts.xlsx

Usage:
    python 14_initial_engine.py                            # scans workspace root
    python 14_initial_engine.py --input Hafs_normalized.xlsx
"""

import argparse
import csv
import json
import re
import zipfile
from datetime import datetime
from pathlib import Path

from openpyxl import Workbook, load_workbook

# -----------------------------
# Config
# -----------------------------


# Set BASE_DIR to the script's directory
BASE_DIR = Path(__file__).resolve().parent
OUTPUT_DIR = BASE_DIR / "14_initial_engine_outputs"
JSON_OUTPUT_DIR = OUTPUT_DIR / "json"
CSV_OUTPUT_DIR = OUTPUT_DIR / "csv"
EXCEL_OUTPUT_DIR = OUTPUT_DIR / "excel"

# 28-letter alphabet in Clear Proof normalization order
ARABIC_LETTERS = [
    "ا","ب","ت","ث","ج","ح","خ","د","ذ","ر","ز","س","ش","ص",
    "ض","ط","ظ","ع","غ","ف","ق","ك","ل","م","ن","ه","و","ي"
]

# Quranic initials by sura (rasm forms), as used in your project
INITIALS_BY_SURA = {
    2:  ["ا","ل","م"],
    3:  ["ا","ل","م"],
    7:  ["ا","ل","م","ص"],
    10: ["ا","ل","ر"],
    11: ["ا","ل","ر"],
    12: ["ا","ل","ر"],
    13: ["ا","ل","م","ر"],
    14: ["ا","ل","ر"],
    15: ["ا","ل","ر"],
    19: ["ك","ه","ي","ع","ص"],
    20: ["ط","ه"],
    26: ["ط","س","م"],
    27: ["ط","س"],
    28: ["ط","س","م"],
    29: ["ا","ل","م"],
    30: ["ا","ل","م"],
    31: ["ا","ل","م"],
    32: ["ا","ل","م"],
    36: ["ي","س"],
    38: ["ص"],
    40: ["ح","م"],
    41: ["ح","م"],
    42: ["ح","م","ع","س","ق"],
    43: ["ح","م"],
    44: ["ح","م"],
    45: ["ح","م"],
    46: ["ح","م"],
    50: ["ق"],
    68: ["ن"],
}

# -----------------------------
# Helpers
# -----------------------------


def slug_from_filename(name: str) -> str:
    """
    Bazzi_Ornamentation__normalized_LETTER_MATRIX.csv -> bazzi
    Submission.org_Ornamentation__normalized_LETTER_MATRIX.csv -> submission.org
    """
    base = Path(name).name
    lower = base.lower()
    if lower.endswith("_letter_matrix.csv"):
        base = base[: -len("_LETTER_MATRIX.csv")]
    elif lower.endswith(".csv"):
        base = base[:-len(".csv")]
    elif lower.endswith(".xlsx"):
        base = base[:-len(".xlsx")]

    lower = base.lower()
    if "_ornamentation__normalized" in lower:
        base = base[: lower.index("_ornamentation__normalized")]
        lower = base.lower()

    # Support names like "The Criterion INITIALS.xlsx".
    if lower.endswith("_initials"):
        base = base[: -len("_initials")]
    elif lower.endswith(" initials"):
        base = base[: -len(" initials")]

    return canonicalize_mushaf_slug(base)


def canonicalize_mushaf_slug(value: str) -> str:
    """Normalize mushaf slug and remove trailing normalized marker for output names."""
    slug = str(value).strip().lower().replace(" ", "")
    if slug.endswith("_normalized"):
        slug = slug[: -len("_normalized")]
    elif slug.endswith("normalized"):
        slug = slug[: -len("normalized")]
    return slug.rstrip("_-")


def read_matrix_rows(path: Path) -> tuple[list[str], list[dict]]:
    """Read tabular rows from either CSV or XLSX matrix files."""
    suffix = path.suffix.lower()

    if suffix == ".csv":
        with open(path, "r", encoding="utf-8-sig", newline="") as csv_file:
            reader = csv.DictReader(csv_file)
            rows = list(reader)
            headers = reader.fieldnames or []
        return headers, rows

    if suffix == ".xlsx":
        wb = load_workbook(path, read_only=True, data_only=True)

        # Criterion-style workbook: one sheet per initialed sura plus summary sheet.
        sura_sheets = [s for s in wb.sheetnames if s.lower().startswith("sura ")]
        if sura_sheets:
            rows = []
            all_headers = {"Sura"}

            for sheet_name in sura_sheets:
                ws = wb[sheet_name]
                try:
                    sura_num = int(sheet_name.split()[1])
                except (ValueError, IndexError):
                    continue

                header_row = next(ws.iter_rows(min_row=1, max_row=1, values_only=True), None)
                if header_row is None:
                    continue
                headers = ["" if h is None else str(h).strip() for h in header_row]

                for row in ws.iter_rows(min_row=2, values_only=True):
                    if row is None:
                        continue
                    row_dict = {"Sura": sura_num}
                    has_value = False
                    for idx, header in enumerate(headers):
                        if not header:
                            continue
                        value = row[idx] if idx < len(row) else None
                        row_dict[header] = value
                        if value not in (None, ""):
                            has_value = True
                            all_headers.add(header)
                    if has_value:
                        rows.append(row_dict)

            wb.close()
            headers = ["Sura", *sorted(h for h in all_headers if h != "Sura")]
            return headers, rows

        # Generic xlsx fallback: active sheet with normal header row.
        ws = wb.active
        header_row = next(ws.iter_rows(min_row=1, max_row=1, values_only=True), None)
        if header_row is None:
            wb.close()
            return [], []

        headers = ["" if h is None else str(h).strip() for h in header_row]
        rows = []
        for row in ws.iter_rows(min_row=2, values_only=True):
            if row is None:
                continue
            row_dict = {}
            has_value = False
            for idx, header in enumerate(headers):
                if not header:
                    continue
                value = row[idx] if idx < len(row) else None
                if value not in (None, ""):
                    has_value = True
                row_dict[header] = value
            if has_value:
                rows.append(row_dict)

        wb.close()
        return headers, rows

    raise RuntimeError(f"Unsupported file type: {path.name}")


def find_column(headers: list[str], candidates: list[str], fallback_index: int | None = None):
    """Find a column name by case-insensitive header lookup; optionally by index."""
    normalized = {str(c).strip().lower(): c for c in headers}
    for cand in candidates:
        key = cand.strip().lower()
        if key in normalized:
            return normalized[key]
    if fallback_index is not None and 0 <= fallback_index < len(headers):
        return headers[fallback_index]
    return None


def to_int(value, default: int = 0) -> int:
    """Safely coerce numeric-looking values into int for stable output."""
    if value is None or value == "":
        return default
    try:
        return int(float(value))
    except (ValueError, TypeError):
        return default


def write_flat_xlsx(headers: list[str], rows: list[dict], output_path: Path) -> None:
    """Write flat rows to a single-sheet xlsx without pandas."""
    wb = Workbook()
    fixed_dt = datetime(2000, 1, 1, 0, 0, 0)
    # Fix workbook timestamps to avoid run-dependent metadata drift.
    wb.properties.created = fixed_dt
    wb.properties.modified = fixed_dt
    ws = wb.active
    ws.title = "initial_engine_flat"
    ws.append(headers)
    for row in rows:
        ws.append([row.get(h, "") for h in headers])
    wb.save(output_path)
    normalize_zip_for_reproducibility(output_path)


def normalize_zip_for_reproducibility(zip_path: Path) -> None:
    """Rewrite a zip container with fixed timestamps, stable member order,
    and pinned compression so output is byte-identical across runs.

    All variable metadata sourced from openpyxl (flag_bits, create_system,
    compress_type, compression level) is replaced with fixed values so that
    the resulting ZIP is fully deterministic regardless of openpyxl internals.
    """
    fixed_zip_dt = (2000, 1, 1, 0, 0, 0)
    fixed_modified = (
        b'<dcterms:modified xsi:type="dcterms:W3CDTF">2000-01-01T00:00:00Z</dcterms:modified>'
    )
    temp_path = zip_path.with_suffix(zip_path.suffix + ".tmp")

    with zipfile.ZipFile(zip_path, "r") as zin, zipfile.ZipFile(temp_path, "w") as zout:
        for info in sorted(zin.infolist(), key=lambda i: i.filename):
            data = zin.read(info.filename)
            if info.filename == "docProps/core.xml":
                data = re.sub(
                    rb"<dcterms:modified xsi:type=\"dcterms:W3CDTF\">.*?</dcterms:modified>",
                    fixed_modified,
                    data,
                    flags=re.DOTALL,
                )
            new_info = zipfile.ZipInfo(info.filename)
            new_info.date_time = fixed_zip_dt
            # Pin all variable ZIP metadata to fixed values so the container
            # is byte-identical regardless of openpyxl's internal flag choices.
            new_info.compress_type = zipfile.ZIP_DEFLATED
            new_info.flag_bits = 0
            new_info.create_system = 0
            new_info.external_attr = 0
            new_info.comment = b""
            new_info.extra = b""
            zout.writestr(new_info, data, compresslevel=6)

    temp_path.replace(zip_path)


def is_data_row(verse_val) -> bool:
    """
    True if this row is a verse row, not the totals row.
    Verse column should be numeric for verse rows.
    """
    if verse_val is None or str(verse_val).strip() == "":
        return False
    try:
        int(verse_val)
        return True
    except (ValueError, TypeError):
        return False


def count_letters_in_text(text: str) -> dict[str, int]:
    """Count occurrences of ARABIC_LETTERS directly from normalized verse text."""
    counts = {letter: 0 for letter in ARABIC_LETTERS}
    for ch in text:
        if ch in counts:
            counts[ch] += 1
    return counts


# -----------------------------
# Core logic
# -----------------------------


def build_14_initial_engine(input_paths: list[Path] | None = None):
    """
    Read matrix files in BASE_DIR and write separated outputs per mushaf
    as JSON, CSV, and XLSX.
    """
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    JSON_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    CSV_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    EXCEL_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # Find input files in the workspace root.
    # By default, use root csv/xlsx source files and skip generated flat files.
    if input_paths:
        matrix_files = [p for p in input_paths if p.is_file()]
    else:
        excluded_names = {
            "14_initial_engine_flat.txt",
            "14_initial_engine_flat.xlsx",
            "14_initial_engine_flat.csv",
        }
        matrix_files = sorted(
            f for f in BASE_DIR.iterdir()
            if f.is_file() and (
                f.suffix.lower() in {".csv", ".xlsx"}
                and f.name.lower() not in excluded_names
            )
        )

    if not matrix_files:
        raise RuntimeError("No supported matrix input files found.")

    # When using default root scan, prune stale outputs not produced in this run.
    prune_stale_outputs = input_paths is None
    produced_json_paths: set[Path] = set()
    produced_csv_paths: set[Path] = set()
    produced_xlsx_paths: set[Path] = set()

    for path in matrix_files:
        headers, rows = read_matrix_rows(path)

        sura_col = find_column(headers, ["sura", "surah"], fallback_index=0)
        verse_col = find_column(headers, ["verse", "ayah"], fallback_index=1)
        text_col = find_column(headers, ["text", "normtext", "rasm"], fallback_index=2)
        if sura_col is None or verse_col is None or text_col is None:
            raise RuntimeError(f"Could not identify Sura/Verse/Text columns in {path.name}")

        mushaf_col = find_column(headers, ["mushaf", "mushaf_slug"])
        if mushaf_col is not None:
            grouped_rows = {}
            for row in rows:
                mushaf_raw = row.get(mushaf_col)
                if mushaf_raw is None or str(mushaf_raw).strip() == "":
                    continue
                mushaf_slug = canonicalize_mushaf_slug(mushaf_raw)
                grouped_rows.setdefault(mushaf_slug, []).append(row)
        else:
            mushaf_slug = slug_from_filename(path.name)
            grouped_rows = {mushaf_slug: rows}

        init_style = any(str(h).strip().lower().startswith("init_") for h in headers)

        # Resolve the 28 letter columns by header first; if unavailable, count from text.
        letter_cols = []
        has_explicit_letter_columns = True
        for idx, letter in enumerate(ARABIC_LETTERS):
            candidates = [letter]
            if init_style:
                candidates.append(f"Init_{letter}")
            col = find_column(headers, candidates, fallback_index=3 + idx)
            if col is None:
                has_explicit_letter_columns = False
            letter_cols.append(col)

        for mushaf_slug, mushaf_rows in sorted(grouped_rows.items()):
            print(f"Processing mushaf: {mushaf_slug} ({path.name})")

            mushaf_data = {
                "alphabet": ARABIC_LETTERS,
                "initials_by_sura": {str(k): v for k, v in INITIALS_BY_SURA.items()},
                "mushaf": mushaf_slug,
                "suras": {},
            }
            flat_rows = []

            for sura, initials in INITIALS_BY_SURA.items():
                verses_list = []

                for row in mushaf_rows:
                    if to_int(row.get(sura_col), default=-1) != sura:
                        continue

                    verse_val = row.get(verse_col)
                    if not is_data_row(verse_val):
                        continue

                    verse_num = to_int(verse_val, default=-1)
                    if verse_num < 0:
                        continue

                    text = row.get(text_col)
                    text = "" if text is None else str(text)

                    if has_explicit_letter_columns:
                        letter_counts = {}
                        for letter, col_name in zip(ARABIC_LETTERS, letter_cols):
                            if col_name is None:
                                letter_counts[letter] = 0
                            else:
                                letter_counts[letter] = to_int(row.get(col_name), default=0)
                    else:
                        letter_counts = count_letters_in_text(text)

                    verses_list.append({
                        "verse": verse_num,
                        "text": text,
                        "letters": letter_counts,
                    })

                    flat_row = {
                        "mushaf": mushaf_slug,
                        "sura": sura,
                        "verse": verse_num,
                        "text": text,
                        "initials": "".join(initials),
                    }
                    for letter in ARABIC_LETTERS:
                        flat_row[letter] = letter_counts[letter]
                    flat_rows.append(flat_row)

                mushaf_data["suras"][str(sura)] = {"initials": initials, "verses": verses_list}

            # Write per-mushaf outputs.
            json_path = JSON_OUTPUT_DIR / f"{mushaf_slug}_initial_counts.json"
            csv_path = CSV_OUTPUT_DIR / f"{mushaf_slug}_initial_counts.csv"
            xlsx_path = EXCEL_OUTPUT_DIR / f"{mushaf_slug}_initial_counts.xlsx"

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

            flat_rows.sort(key=lambda r: (r["sura"], r["verse"]))
            flat_headers = ["mushaf", "sura", "verse", "text", "initials", *ARABIC_LETTERS]
            with open(csv_path, "w", encoding="utf-8-sig", newline="") as f:
                writer = csv.DictWriter(f, fieldnames=flat_headers)
                writer.writeheader()
                writer.writerows(flat_rows)

            write_flat_xlsx(flat_headers, flat_rows, xlsx_path)

            produced_json_paths.add(json_path)
            produced_csv_paths.add(csv_path)
            produced_xlsx_paths.add(xlsx_path)

            print("  Wrote files:")
            print(f"  - JSON: {json_path}")
            print(f"  - CSV:  {csv_path}")
            print(f"  - XLSX: {xlsx_path}")

    if prune_stale_outputs:
        for existing in JSON_OUTPUT_DIR.glob("*_initial_counts.json"):
            if existing not in produced_json_paths:
                existing.unlink(missing_ok=True)
        for existing in CSV_OUTPUT_DIR.glob("*_initial_counts.csv"):
            if existing not in produced_csv_paths:
                existing.unlink(missing_ok=True)
        for existing in EXCEL_OUTPUT_DIR.glob("*_initial_counts.xlsx"):
            if existing not in produced_xlsx_paths:
                existing.unlink(missing_ok=True)

    print(f"\nAll separated outputs written to: {OUTPUT_DIR}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Build 14-initial outputs from matrix CSV/XLSX files."
    )
    parser.add_argument(
        "--input",
        nargs="+",
        default=None,
        help="Optional one or more input files. If omitted, scans root CSV/XLSX files and excludes 14_initial_engine_flat.(txt|xlsx|csv).",
    )
    args = parser.parse_args()

    input_paths = None
    if args.input:
        input_paths = [Path(p).resolve() for p in args.input]

    build_14_initial_engine(input_paths=input_paths)
