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

import argparse
import json
import re
import sys
import unicodedata
from collections import OrderedDict, defaultdict
from pathlib import Path

# must exist in same folder
try:
    from gretil_iast_to_deva import transliterate_text
except Exception as e:
    print("ERROR: cannot import transliterate_text from gretil_iast_to_deva.py")
    print(f"Import error: {e}")
    sys.exit(1)


# Marker examples handled:
# /AP291.001ab/
# //AP291.001cd//
# AP291.001ab
MARKER_RE = re.compile(
    r'(?P<full>/+)?(?P<work>[A-Za-z]{1,10})(?P<chapter>\d{1,4})\.(?P<verse>\d{3})(?P<pada>[a-z]{0,2})(?P<full2>/+)?'
)

SPACE_RE = re.compile(r'\s+')

# crude boilerplate filters for roman files
SKIP_LINE_PATTERNS = [
    re.compile(r'^\s*$'),
    re.compile(r'^\s*GRETIL\b', re.I),
    re.compile(r'^\s*Göttingen\b', re.I),
    re.compile(r'^\s*THIS GRETIL TEXT FILE\b', re.I),
    re.compile(r'^\s*COPYRIGHT\b', re.I),
    re.compile(r'^\s*Read .*license', re.I),
    re.compile(r'^\s*Input by\b', re.I),
    re.compile(r'^\s*Proof[- ]?reading\b', re.I),
    re.compile(r'^\s*Structural markers\b', re.I),
    re.compile(r'^\s*Text converted to Unicode', re.I),
    re.compile(r'^\s*DOCUMENTATION', re.I),
    re.compile(r'^\s*C O N T E N T S', re.I),
    re.compile(r'^\s*Chapter [0-9IVXLCDM]+', re.I),
    re.compile(r'^\s*BOOK [0-9IVXLCDM]+', re.I),
    re.compile(r'^\s*APPENDIX\b', re.I),
    re.compile(r'^\s*Index\b', re.I),
    re.compile(r'^\s*Page\b', re.I),
    re.compile(r'^\s*\[[0-9]+\]\s*$'),
    re.compile(r'^\s*[0-9]+\s*$'),
]

# order padas sensibly
PADA_ORDER = {
    "a": 1,
    "ab": 1,
    "b": 2,
    "bc": 2,
    "c": 3,
    "cd": 3,
    "d": 4,
    "de": 4,
    "e": 5,
    "ef": 5,
    "f": 6,
    "fg": 6,
    "g": 7,
    "gh": 7,
    "h": 8,
    "uv": 9,
    "": 99,
}


def norm(s: str) -> str:
    s = unicodedata.normalize("NFC", s)
    s = s.replace("\r\n", "\n").replace("\r", "\n")
    s = s.replace("\xa0", " ")
    return s


def is_skip_line(line: str) -> bool:
    for pat in SKIP_LINE_PATTERNS:
        if pat.search(line):
            return True
    return False


def clean_line_text(line: str) -> str:
    line = norm(line)
    line = line.strip()

    # remove trailing slash clusters before marker area
    line = re.sub(r'[/|।॥\s]+$', '', line)

    # collapse whitespace
    line = SPACE_RE.sub(" ", line).strip()
    return line


def parse_marker_from_line(line: str):
    """
    Return (text_before_marker, work, chapter, verse, pada) using the LAST marker on the line.
    """
    matches = list(MARKER_RE.finditer(line))
    if not matches:
        return None

    m = matches[-1]
    text = line[:m.start()]
    text = clean_line_text(text)

    work = m.group("work")
    chapter = int(m.group("chapter"))
    verse = int(m.group("verse"))
    pada = m.group("pada") or ""

    if not text:
        return None

    return text, work, chapter, verse, pada


def pada_sort_key(pada: str):
    return (PADA_ORDER.get(pada, 50), pada)


def join_padas_plain(pada_map: dict) -> str:
    parts = [pada_map[k] for k in sorted(pada_map.keys(), key=pada_sort_key) if pada_map[k].strip()]
    return " ".join(parts).strip()


def join_padas_display_dev(pada_map: dict) -> str:
    """
    Human-readable Devanagari display line.
    For common ab/cd case:
        ab । cd ॥
    Otherwise join all padas with danda separators and end with double danda if >1 segment.
    """
    ordered_keys = [k for k in sorted(pada_map.keys(), key=pada_sort_key) if pada_map[k].strip()]
    ordered_vals = [pada_map[k].strip() for k in ordered_keys]

    if not ordered_vals:
        return ""

    if len(ordered_vals) == 1:
        return ordered_vals[0]

    if len(ordered_vals) == 2 and ordered_keys == ["ab", "cd"]:
        return f"{ordered_vals[0]} । {ordered_vals[1]} ॥"

    return " । ".join(ordered_vals) + " ॥"


def process_file(src_file: Path):
    """
    Read one .roman.txt file and return verse records in source order.
    """
    text = src_file.read_text(encoding="utf-8", errors="replace")
    text = norm(text)

    records = OrderedDict()

    for raw_line in text.splitlines():
        line = raw_line.strip()
        if is_skip_line(line):
            continue

        parsed = parse_marker_from_line(line)
        if not parsed:
            continue

        verse_text, work, chapter, verse, pada = parsed
        key = (work, chapter, verse)

        if key not in records:
            records[key] = {
                "work_code": work,
                "chapter": chapter,
                "verse": verse,
                "padas": OrderedDict(),
            }

        # keep first occurrence if duplicate pada appears
        if pada not in records[key]["padas"]:
            records[key]["padas"][pada] = verse_text

    out = []
    for key, rec in records.items():
        plain_roman = join_padas_plain(rec["padas"])
        display_roman = plain_roman  # optional; Roman display kept plain

        plain_dev = transliterate_text(
            plain_roman,
            strip_gretil_tags=True,
            devanagari_digits=False
        )

        display_dev = transliterate_text(
            join_padas_display_dev(rec["padas"]),
            strip_gretil_tags=True,
            devanagari_digits=False
        )

        unified_id = f"{rec['work_code']}{rec['chapter']}.{rec['verse']:03d}"

        out.append({
            "id": unified_id,
            "work_code": rec["work_code"],
            "chapter": rec["chapter"],
            "verse": rec["verse"],
            "source_file": src_file.name,
            "text_roman": plain_roman,
            "text_dev": plain_dev,
            "display_dev": display_dev,
            "padas": rec["padas"],
        })

    return out


def write_grouped_clean_text(records_by_file, out_root: Path):
    out_root.mkdir(parents=True, exist_ok=True)

    for src_name, recs in records_by_file.items():
        stem = Path(src_name).stem.replace(".roman", "")
        roman_out = out_root / f"{stem}.clean.roman.txt"
        dev_out = out_root / f"{stem}.clean.dev.txt"

        roman_lines = []
        dev_lines = []

        for r in recs:
            roman_lines.append(f"{r['id']}\t{r['text_roman']}")
            dev_lines.append(f"{r['id']}\t{r['display_dev']}")

        roman_out.write_text("\n".join(roman_lines) + "\n", encoding="utf-8", newline="\n")
        dev_out.write_text("\n".join(dev_lines) + "\n", encoding="utf-8", newline="\n")


def main():
    ap = argparse.ArgumentParser(
        description="Clean GRETIL roman corpus, merge padas, and emit JSONL + clean Roman/Devanagari text."
    )
    ap.add_argument("src_root", help="Folder containing .roman.txt files")
    ap.add_argument("out_jsonl", help="Output JSONL path")
    ap.add_argument(
        "--clean-text-root",
        help="Optional folder to write per-book .clean.roman.txt and .clean.dev.txt files"
    )
    args = ap.parse_args()

    src_root = Path(args.src_root).resolve()
    out_jsonl = Path(args.out_jsonl).resolve()

    if not src_root.exists():
        print(f"ERROR: source folder does not exist: {src_root}")
        sys.exit(1)

    files = sorted(src_root.rglob("*.roman.txt"))
    if not files:
        print(f"ERROR: no .roman.txt files found under {src_root}")
        sys.exit(1)

    all_records = []
    by_file = OrderedDict()

    for f in files:
        recs = process_file(f)
        by_file[f.name] = recs
        all_records.extend(recs)
        print(f"OK  {f.name} -> {len(recs)} verses")

    out_jsonl.parent.mkdir(parents=True, exist_ok=True)
    with out_jsonl.open("w", encoding="utf-8", newline="\n") as out:
        for rec in all_records:
            out.write(json.dumps(rec, ensure_ascii=False) + "\n")

    print(f"JSONL written: {out_jsonl}")
    print(f"Total verses : {len(all_records)}")

    if args.clean_text_root:
        clean_root = Path(args.clean_text_root).resolve()
        write_grouped_clean_text(by_file, clean_root)
        print(f"Clean text written under: {clean_root}")


if __name__ == "__main__":
    main()