#!/usr/bin/env python3
"""Fetch חוק הביטוח הלאומי and its תקנות from Hebrew Wikisource and convert to
plain text with the headers the shira-hermes chunker expects (חלק / פרק /
סימן / סעיף).

Modes:
  (default)         fetch the parent law (חוק הביטוח הלאומי) → /tmp/national_insurance_law.txt
  --regulations     fetch all תקנות הביטוח הלאומי canonical pages (~92, redirects skipped)
                    → /tmp/regulations/<slug>.txt + <slug>.txt.meta.json
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
import urllib.request
import urllib.parse

WIKI_PAGE = "חוק_הביטוח_הלאומי"
API = "https://he.wikisource.org/w/api.php"
REGULATIONS_PREFIX = "תקנות הביטוח הלאומי"


def fetch_wikitext(page: str) -> str:
    url = f"{API}?action=parse&page={urllib.parse.quote(page)}&prop=wikitext&format=json"
    req = urllib.request.Request(url, headers={"User-Agent": "shira-kb/1.0"})
    with urllib.request.urlopen(req) as r:
        data = json.load(r)
    return data["parse"]["wikitext"]["*"]


# ---- template handlers ----------------------------------------------------

# {{ח:קטע2|סימולא|שם החלק}} → "חלק <number> — <title>"
# Wikisource numbers parts/chapters in their own way; sometimes the number is
# embedded in the title. We emit the raw text of the second param and let the
# chunker pick up the prefix from there.
def _section_header(prefix: str, args: list[str]) -> str:
    args = _positional_args(args)
    if not args:
        return ""
    # The title is usually the last positional arg; for קטע templates the
    # title already contains its own prefix ("פרק א׳: ...") so we don't
    # add our own. Otherwise prepend the given prefix.
    title = (args[-1] if args[-1].strip() else (args[0] if args else "")).strip()
    if not title:
        return ""
    if prefix:
        return f"\n\n{prefix} {title}\n\n"
    return f"\n\n{title}\n\n"


# Templates that produce a statute heading line.
# Note: {{ח:קטע2}} / {{ח:קטע3}} don't directly encode the level name —
# the actual level is embedded in the *title* argument (e.g. "פרק א׳: ...").
# So we emit the title verbatim and let the chunker's existing regexes
# pick it up (which they do, for פרק / סימן prefixes).
_HEADER_TEMPLATES = {
    "ח:קטע2": "",
    "ח:קטע3": "",
    "ח:קטע4": "",
}


def _positional_args(args: list[str]) -> list[str]:
    """Filter out named arguments (e.g. 'אחר=[פרק א׳]') from a template arg list."""
    return [a for a in args if "=" not in a.split("\n", 1)[0][:40]]


def _render_section(args: list[str]) -> str:
    """{{ח:סעיף|number|title|anchor?}} → "סעיף <number>. <title>" """
    args = _positional_args(args)
    if not args:
        return ""
    number = args[0].strip()
    title = (args[1] if len(args) > 1 else "").strip()
    if title:
        return f"\n\nסעיף {number}. {title}\n\n"
    return f"\n\nסעיף {number}.\n\n"


def _render_subsection(args: list[str]) -> str:
    """{{ח:ת|א|text}} → "(א) text" on its own line"""
    args = _positional_args(args)
    if not args:
        return ""
    letter = args[0].strip()
    body = (args[1] if len(args) > 1 else "").strip()
    return f"\n({letter}) {body}"


def _render_linkish(args: list[str]) -> str:
    """{{ח:פנימי|anchor|label}} / {{ח:חיצוני|url|label}} / {{ח:תיבה|...}} →
    just the label text, dropping the link itself."""
    if not args:
        return ""
    # label is typically the second arg; if only one, return that.
    return args[1].strip() if len(args) > 1 else args[0].strip()


def _drop(args: list[str]) -> str:
    return ""


# Map template name → handler. Unknown templates are silently stripped.
_HANDLERS = {
    "ח:סעיף": _render_section,
    "ח:סעיף*": _render_section,
    "ח:ת": _render_subsection,
    "ח:תת": _render_subsection,
    "ח:תתת": _render_subsection,
    "ח:תתתת": _render_subsection,
    "ח:תתתתת": _render_subsection,
    "ח:תתתתתת": _render_subsection,
    "ח:פנימי": _render_linkish,
    "ח:חיצוני": _render_linkish,
    "ח:תיבה": _render_linkish,
    # No body — drop silently
    "ח:התחלה": _drop,
    "ח:סוף": _drop,
    "ח:כותרת": _drop,
    "ח:מאגר": _drop,
    "ח:פתיח-התחלה": _drop,
    "ח:מפריד": _drop,
    "ח:סוגר": _drop,
    "ח:מבוא": _drop,
    "ח:חתימות": _drop,
    "ח:הערה": _drop,   # footnotes — not needed for RAG
}


def _parse_template_args(raw: str) -> list[str]:
    """Split a template body by '|' at depth 0 (respecting nested templates)."""
    args: list[str] = []
    buf: list[str] = []
    depth = 0
    i = 0
    while i < len(raw):
        c = raw[i]
        if c == "{" and i + 1 < len(raw) and raw[i + 1] == "{":
            depth += 1
            buf.append("{{")
            i += 2
            continue
        if c == "}" and i + 1 < len(raw) and raw[i + 1] == "}":
            depth -= 1
            buf.append("}}")
            i += 2
            continue
        if c == "[" and i + 1 < len(raw) and raw[i + 1] == "[":
            depth += 1
            buf.append("[[")
            i += 2
            continue
        if c == "]" and i + 1 < len(raw) and raw[i + 1] == "]":
            depth -= 1
            buf.append("]]")
            i += 2
            continue
        if c == "|" and depth == 0:
            args.append("".join(buf))
            buf = []
            i += 1
            continue
        buf.append(c)
        i += 1
    args.append("".join(buf))
    return args


def _expand_templates(text: str) -> str:
    """Recursively resolve {{ח:...}} templates, depth-first."""
    out: list[str] = []
    i = 0
    while i < len(text):
        if text[i : i + 2] == "{{":
            # find matching }}
            depth = 1
            j = i + 2
            while j < len(text) and depth > 0:
                if text[j : j + 2] == "{{":
                    depth += 1
                    j += 2
                elif text[j : j + 2] == "}}":
                    depth -= 1
                    j += 2
                else:
                    j += 1
            body = text[i + 2 : j - 2]  # inside {{ ... }}
            # recurse into body first
            expanded_body = _expand_templates(body)
            args = _parse_template_args(expanded_body)
            name = args[0].strip() if args else ""
            if name in _HEADER_TEMPLATES:
                out.append(_section_header(_HEADER_TEMPLATES[name], args[1:]))
            elif name in _HANDLERS:
                out.append(_HANDLERS[name](args[1:]))
            else:
                # Unknown template — keep the last arg if it looks like text.
                if len(args) > 1:
                    candidate = args[-1].strip()
                    if candidate and len(candidate) < 200:
                        out.append(candidate)
            i = j
        else:
            out.append(text[i])
            i += 1
    return "".join(out)


def _cleanup(text: str) -> str:
    # Drop wiki tables/images we don't care about.
    text = re.sub(r"\[\[קובץ:[^\]]*\]\]", "", text)
    text = re.sub(r"\[\[(?:[^\]|]*\|)?([^\]]*)\]\]", r"\1", text)  # [[link|label]] → label
    # Drop HTML-ish refs (footnotes, comments).
    text = re.sub(r"<ref[^>]*>.*?</ref>", "", text, flags=re.DOTALL)
    text = re.sub(r"<!--.*?-->", "", text, flags=re.DOTALL)
    text = re.sub(r"</?[a-zA-Z][^>]*>", "", text)
    # Collapse 3+ newlines to 2; strip trailing spaces.
    text = re.sub(r"[ \t]+\n", "\n", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    # Drop bold/italic wiki markup.
    text = text.replace("'''", "").replace("''", "")
    return text.strip()


def fetch_regulation_titles() -> list[str]:
    """Query MediaWiki API for canonical (non-redirect) תקנות הביטוח הלאומי pages.

    Wikisource keeps duplicate-spelling pages (e.g. קיצבה/קצבה) as redirects to
    the modern canonical title. Filtering apfilterredir=nonredirects gives us
    the deduped set without having to maintain a hand-curated list.
    """
    titles: list[str] = []
    apcontinue: str | None = None
    while True:
        params = {
            "action": "query",
            "list": "allpages",
            "apprefix": REGULATIONS_PREFIX,
            "apnamespace": "0",
            "apfilterredir": "nonredirects",
            "aplimit": "500",
            "format": "json",
        }
        if apcontinue:
            params["apcontinue"] = apcontinue
        url = API + "?" + urllib.parse.urlencode(params)
        req = urllib.request.Request(url, headers={"User-Agent": "shira-kb/1.0"})
        with urllib.request.urlopen(req) as r:
            data = json.load(r)
        titles.extend(p["title"] for p in data["query"]["allpages"])
        cont = data.get("continue", {}).get("apcontinue")
        if not cont:
            break
        apcontinue = cont
    return titles


_SLUG_DROP = re.compile(r"[()\[\]{}'\"־–—]")  # parens, brackets, quotes, hyphens
_SLUG_WS = re.compile(r"[\s\-]+")


def _slug(title: str) -> str:
    """Convert a Hebrew title into a safe filename — Hebrew letters preserved.

    Strips parens/brackets/quotes (problematic in shells and S3 console URLs),
    collapses whitespace and hyphens to a single hyphen.
    """
    s = _SLUG_DROP.sub(" ", title)
    s = _SLUG_WS.sub("-", s.strip())
    return s


def fetch_one_regulation(title: str, out_dir: str) -> dict:
    """Fetch a single regulation, expand templates, clean, write text + sidecar."""
    wt = fetch_wikitext(title)
    expanded = _expand_templates(wt)
    clean = _cleanup(expanded)
    slug = _slug(title)
    txt_path = os.path.join(out_dir, f"{slug}.txt")
    meta_path = txt_path + ".meta.json"
    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(clean)
    sidecar = {
        "title": title,
        "identifier": title,
        "source_url": f"https://he.wikisource.org/wiki/{urllib.parse.quote(title.replace(' ', '_'))}",
    }
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(sidecar, f, ensure_ascii=False, indent=2)
    return {"title": title, "slug": slug, "raw_chars": len(wt), "clean_chars": len(clean)}


def fetch_all_regulations(out_dir: str) -> int:
    """Fetch every canonical תקנות הביטוח הלאומי page into out_dir."""
    os.makedirs(out_dir, exist_ok=True)
    titles = fetch_regulation_titles()
    print(f"[regulations] {len(titles)} canonical pages", file=sys.stderr)
    total_clean = 0
    for i, title in enumerate(titles, 1):
        try:
            r = fetch_one_regulation(title, out_dir)
            total_clean += r["clean_chars"]
            print(
                f"[{i:3d}/{len(titles)}] {r['slug']:.<70s} "
                f"raw={r['raw_chars']:>7,}  clean={r['clean_chars']:>7,}",
                file=sys.stderr,
            )
        except Exception as e:
            print(f"[{i:3d}/{len(titles)}] FAILED {title}: {e}", file=sys.stderr)
        # Be polite to Wikisource — 50ms between requests = ~20 req/s, well under any limit.
        time.sleep(0.05)
    print(f"[done] total clean text: {total_clean:,} chars across {len(titles)} regulations",
          file=sys.stderr)
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--regulations", action="store_true",
                        help="fetch all תקנות (default: fetch the parent law only)")
    parser.add_argument("out", nargs="?", default=None,
                        help="output path (default: /tmp/national_insurance_law.txt for law, /tmp/regulations/ for --regulations)")
    args = parser.parse_args()

    if args.regulations:
        out_dir = args.out or "/tmp/regulations"
        return fetch_all_regulations(out_dir)

    out_path = args.out or "/tmp/national_insurance_law.txt"
    wt = fetch_wikitext(WIKI_PAGE)
    print(f"[fetch] raw wikitext: {len(wt):,} chars", file=sys.stderr)
    expanded = _expand_templates(wt)
    print(f"[expand] after templates: {len(expanded):,} chars", file=sys.stderr)
    clean = _cleanup(expanded)
    print(f"[cleanup] final: {len(clean):,} chars", file=sys.stderr)
    with open(out_path, "w", encoding="utf-8") as f:
        f.write(clean)
    print(f"[write] {out_path}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())