import argparse
import csv
import json
import shutil
from collections import OrderedDict
from pathlib import Path

from openpyxl import load_workbook, Workbook


# 28-letter normalized alphabet 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())

NORMALIZED_CSV_SUFFIX = "_normalized.csv"
NORMALIZED_XLSX_SUFFIX = "_normalized.xlsx"

DEFAULT_EXCLUDE_DIRS = {
    ".git",
    ".venv",
    "venv",
    "__pycache__",
    ".mypy_cache",
    ".pytest_cache",
    ".ruff_cache",
    ".vscode",
    "out",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build per-verse Arabic letter tally matrices from normalized Quran sources."
    )
    parser.add_argument(
        "--in-dir",
        type=Path,
        default=Path("."),
        help="Input workspace directory. Recursively scanned by default.",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("out"),
        help="Output root directory. Default: out",
    )
    parser.add_argument(
        "--no-recursive",
        action="store_true",
        help="Only scan the input folder, not subfolders.",
    )
    parser.add_argument(
        "--exclude-dir",
        action="append",
        default=[],
        help="Directory name to exclude from input discovery. Can be used multiple times.",
    )
    parser.add_argument(
        "--keep-existing",
        action="store_true",
        help="Do not delete/rebuild the output directory before running.",
    )
    return parser.parse_args()


def ensure_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def is_excluded(path: Path, root: Path, exclude_dirs: set[str]) -> bool:
    try:
        rel = path.relative_to(root)
    except ValueError:
        rel = path
    return any(part in exclude_dirs for part in rel.parts)


def slug_from_filename(path: Path) -> str:
    return path.stem


def sync_output_dirs(out_dir: Path, keep_existing: bool) -> tuple[Path, Path, Path, Path]:
    if out_dir.exists() and not keep_existing:
        shutil.rmtree(out_dir)

    csv_dir = out_dir / "csv"
    json_dir = out_dir / "json"
    excel_dir = out_dir / "excel"
    txt_dir = out_dir / "txt"

    for folder in (csv_dir, json_dir, excel_dir, txt_dir):
        ensure_dir(folder)

    return csv_dir, json_dir, excel_dir, txt_dir


def discover_source_files(root: Path, recursive: bool, exclude_dirs: set[str]) -> list[Path]:
    """
    Discover normalized source inputs.

    Rules:
    - Accept *_normalized.csv and *_normalized.xlsx.
    - Ignore Excel lock/temp files starting with ~$.
    - If both CSV and XLSX exist for the same stem, prefer CSV.
    - Recursive by default.
    """
    pattern = "**/*" if recursive else "*"
    csv_by_stem: dict[str, Path] = {}
    xlsx_by_stem: dict[str, Path] = {}

    for path in root.glob(pattern):
        if not path.is_file():
            continue
        if path.name.startswith("~$"):
            continue
        if is_excluded(path, root, exclude_dirs):
            continue

        name_lower = path.name.lower()
        if name_lower.endswith(NORMALIZED_CSV_SUFFIX):
            stem = path.name[:-len(NORMALIZED_CSV_SUFFIX)]
            csv_by_stem[stem.lower()] = path
        elif name_lower.endswith(NORMALIZED_XLSX_SUFFIX):
            stem = path.name[:-len(NORMALIZED_XLSX_SUFFIX)]
            xlsx_by_stem[stem.lower()] = path

    chosen = dict(xlsx_by_stem)
    chosen.update(csv_by_stem)  # CSV wins when same stem exists

    return sorted(chosen.values(), key=lambda p: str(p).lower())


def build_letter_matrix_for_csv(csv_path: Path):
    rows = []
    totals = OrderedDict((ltr, 0) for ltr in ARABIC_LETTERS)

    with csv_path.open(encoding="utf-8-sig", newline="") as f:
        reader = csv.reader(f)
        headers = next(reader, None)
        if not headers:
            return rows, totals

        normalized_headers = [h.strip().lower() 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 {csv_path}. Expected sura, verse, text.") from exc

        for record in reader:
            if not record:
                continue
            try:
                sura_raw = record[sura_idx]
                verse_raw = record[verse_idx]
                text_raw = record[text_idx]
            except IndexError:
                continue

            if not sura_raw or not verse_raw:
                continue

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

            text = "" if text_raw is None else str(text_raw)
            verse_counts = {ltr: 0 for ltr in ARABIC_LETTERS}
            for ch in text:
                if ch in verse_counts:
                    verse_counts[ch] += 1
                    totals[ch] += 1

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

    return rows, totals


def build_letter_matrix_for_xlsx(xlsx_path: Path):
    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)
        verse_counts = {ltr: 0 for ltr in ARABIC_LETTERS}
        for ch in text:
            if ch in verse_counts:
                verse_counts[ch] += 1
                totals[ch] += 1

        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, csv_dir: Path) -> Path:
    letter_headers = [f"{arabic} ({trans})" for arabic, trans in LETTER_MAP.items()]
    csv_path = csv_dir / f"{slug}_LETTER_MATRIX.csv"

    with csv_path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["sura", "verse", "text"] + letter_headers)
        for row in rows:
            writer.writerow([row["sura"], row["verse"], row["text"]] + [row[a] for a in ARABIC_LETTERS])
        writer.writerow(["", "TOTAL", ""] + [totals[a] for a in ARABIC_LETTERS])

    return csv_path


def write_letter_matrix_json(slug: str, rows, totals, json_dir: Path) -> Path:
    json_path = json_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())),
    }
    json_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    return json_path


def write_letter_matrix_xlsx(slug: str, rows, totals, excel_dir: Path) -> Path:
    xlsx_path = excel_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:
        ws.append([row["sura"], row["verse"], row["text"]] + [row[a] for a in ARABIC_LETTERS])

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


def write_grand_summary_txt(all_results: list, txt_dir: Path) -> Path:
    txt_path = txt_dir / "GRAND_SUMMARY.txt"

    combined_totals = OrderedDict((ltr, 0) for ltr in ARABIC_LETTERS)
    total_verses = 0
    for _slug, rows, totals in all_results:
        total_verses += len(rows)
        for ltr in ARABIC_LETTERS:
            combined_totals[ltr] += totals[ltr]

    grand_total_letters = sum(combined_totals.values())
    bar = "=" * 68
    bar2 = "-" * 68

    lines = [
        bar,
        "GRAND SUMMARY — ALL NORMALIZED SOURCES",
        f"Mushafs        : {len(all_results)}",
        f"Total Verses   : {total_verses:>12,}",
        f"Total Letters  : {grand_total_letters:>12,}",
        bar,
        f"  {'Mushaf':<42} {'Verses':>8} {'Letters':>12}",
        bar2,
    ]

    for slug, rows, totals in all_results:
        lines.append(f"  {slug:<42} {len(rows):>8,} {sum(totals.values()):>12,}")

    lines.extend([
        bar,
        f"  {'Letter':<6} {'Transliteration':<16} {'Combined Count':>14} {'Share':>8}",
        bar2,
    ])

    for arabic, trans in LETTER_MAP.items():
        count = combined_totals[arabic]
        pct = (count / grand_total_letters * 100) if grand_total_letters else 0.0
        lines.append(f"  {arabic:<6} {trans:<16} {count:>14,} {pct:>7.3f}%")

    lines.extend([bar, "END OF GRAND SUMMARY", bar])
    txt_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return txt_path


def main() -> int:
    args = parse_args()
    root = args.in_dir.resolve()
    out_dir = args.out_dir.resolve() if args.out_dir.is_absolute() else (Path.cwd() / args.out_dir).resolve()

    exclude_dirs = set(DEFAULT_EXCLUDE_DIRS)
    exclude_dirs.update(args.exclude_dir)

    if not root.is_dir():
        raise RuntimeError(f"Input folder not found: {root}")

    source_files = discover_source_files(root, recursive=not args.no_recursive, exclude_dirs=exclude_dirs)
    if not source_files:
        raise RuntimeError(f"No source files matching *{NORMALIZED_CSV_SUFFIX} or *{NORMALIZED_XLSX_SUFFIX} found in {root}")

    csv_dir, json_dir, excel_dir, txt_dir = sync_output_dirs(out_dir, keep_existing=args.keep_existing)

    all_results = []

    print(f"[INFO] Source root: {root}")
    print(f"[INFO] Output root: {out_dir}")
    print(f"[INFO] Normalized sources discovered: {len(source_files)}")

    for src_path in source_files:
        slug = slug_from_filename(src_path)
        print(f"[INFO] Building letter matrix for {slug} from {src_path.relative_to(root) if src_path.is_relative_to(root) else src_path} ...")

        if src_path.suffix.lower() == ".csv":
            rows, totals = build_letter_matrix_for_csv(src_path)
        else:
            rows, totals = build_letter_matrix_for_xlsx(src_path)

        csv_path = write_letter_matrix_csv(slug, rows, totals, csv_dir)
        json_path = write_letter_matrix_json(slug, rows, totals, json_dir)
        xlsx_path = write_letter_matrix_xlsx(slug, rows, totals, excel_dir)

        all_results.append((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 -> {xlsx_path}")

    txt_path = write_grand_summary_txt(all_results, txt_dir)
    print(f"[OK] GRAND SUMMARY TXT -> {txt_path}")
    print("[OK] Done. Letter matrices generated in csv/json/xlsx; GRAND_SUMMARY.txt written.")

    return 0


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