#!/usr/bin/env python3
"""Convert a wiktextract dump (kaikki.org) to pack-builder JSONL.

Source: https://kaikki.org/<edition>/raw-wiktextract-data.jsonl.gz
Licence: CC BY-SA (same as Wiktionary) — the resulting pack MUST carry
attribution and stay under CC BY-SA. See docs/mobile/PACK_FORMAT.md.

    python wiktextract_to_jsonl.py --in ruwikt.jsonl.gz --out ru.jsonl \
        --lang-code ru

Streams the input: the uncompressed dump is several GB and must never be
held in memory. One output row per (word, part-of-speech) record — the app
groups same-headword articles at display time, exactly as it does for the
several dictionaries of the main corpus.
"""
from __future__ import annotations

import argparse
import gzip
import html
import io
import json
import re
import sys
from pathlib import Path

# wiktextract POS code -> short label shown before the senses.
POS_RU = {
    "noun": "сущ.", "verb": "гл.", "adj": "прил.", "adv": "нареч.",
    "pron": "мест.", "num": "числ.", "prep": "предл.", "conj": "союз",
    "particle": "частица", "intj": "межд.", "phrase": "выражение",
    "proverb": "посл.", "prefix": "приставка", "suffix": "суффикс",
    "abbrev": "сокр.", "name": "имя собств.", "participle": "прич.",
    "adv_phrase": "нареч. выражение", "character": "символ",
}
POS_EN = {
    "noun": "n.", "verb": "v.", "adj": "adj.", "adv": "adv.",
    "pron": "pron.", "num": "num.", "prep": "prep.", "conj": "conj.",
    "particle": "part.", "intj": "interj.", "phrase": "phrase",
    "proverb": "proverb", "prefix": "prefix", "suffix": "suffix",
    "abbrev": "abbr.", "name": "prop. n.",
}


def esc(text: str) -> str:
    """Escape for embedding in the pack's restricted markup."""
    return html.escape(str(text), quote=False)


def build_article(rec: dict, pos_labels: dict, ui_ru: bool) -> str:
    out: list[str] = []

    pos = rec.get("pos") or ""
    label = pos_labels.get(pos, pos)
    if label:
        out.append("<s-n>%s</s-n>" % esc(label))

    ipa = ""
    for s in rec.get("sounds") or []:
        if s.get("ipa"):
            ipa = s["ipa"]
            break
    if ipa:
        out.append(" <s-n>%s</s-n>" % esc(ipa))

    senses = [s for s in (rec.get("senses") or []) if s.get("glosses")]
    for i, sense in enumerate(senses, 1):
        gloss = " ".join(sense["glosses"])
        num = ("<b>%d.</b> " % i) if len(senses) > 1 else ""
        out.append("<b-q>%s%s</b-q>" % (num, esc(gloss)))
        for ex in (sense.get("examples") or [])[:3]:
            text = (ex.get("text") or "").strip()
            if not text:
                continue
            out.append("<e-x>%s</e-x>" % esc(text))
            ref = (ex.get("ref") or "").strip()
            if ref:
                out.append("<r-f>%s</r-f>" % esc(ref))

    if not senses:
        return ""

    def related(field: str, title_ru: str, title_en: str) -> None:
        words = []
        for item in (rec.get(field) or []):
            w = (item.get("word") or "").strip()
            if w and w not in words:
                words.append(w)
        if not words:
            return
        title = title_ru if ui_ru else title_en
        links = ", ".join("<r>%s</r>" % esc(w) for w in words[:20])
        out.append("<f-r><b>%s:</b> %s</f-r>" % (title, links))

    related("synonyms", "Синонимы", "Synonyms")
    related("antonyms", "Антонимы", "Antonyms")
    related("hypernyms", "Гиперонимы", "Hypernyms")

    ety = rec.get("etymology_texts") or []
    if ety:
        joined = " ".join(t.strip() for t in ety if t and t.strip())
        if joined:
            out.append("<d-t>%s</d-t>" % esc(joined))

    return "".join(out)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--in", dest="src", required=True,
                    help="raw-wiktextract-data.jsonl(.gz)")
    ap.add_argument("--out", required=True)
    ap.add_argument("--lang-code", required=True,
                    help="keep only entries of this language, e.g. ru")
    ap.add_argument("--min-gloss", type=int, default=8,
                    help="skip articles shorter than this many chars")
    args = ap.parse_args()

    src = Path(args.src)
    opener = gzip.open if src.suffix == ".gz" else open
    kept = skipped = 0
    ui_ru = args.lang_code == "ru"
    pos_labels = POS_RU if ui_ru else POS_EN

    with opener(str(src), "rt", encoding="utf-8", errors="replace") as fh, \
            Path(args.out).open("w", encoding="utf-8") as out:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                skipped += 1
                continue
            if rec.get("lang_code") != args.lang_code:
                continue
            word = (rec.get("word") or "").strip()
            if not word or len(word) > 120:
                skipped += 1
                continue
            body = build_article(rec, pos_labels, ui_ru)
            if len(body) < args.min_gloss:
                skipped += 1
                continue
            out.write(json.dumps(
                {"term": word, "definition": body, "lang": args.lang_code},
                ensure_ascii=False) + "\n")
            kept += 1
            if kept % 50000 == 0:
                print("  %d kept…" % kept, file=sys.stderr)

    print("wrote %d entries, skipped %d" % (kept, skipped))
    return 0


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