"""
homograph preprocessor
rewrites russian sentences containing stress homographs via local llm.
writes each sentence immediately — resume-safe if interrupted.
in --all mode: processes every sentence and reports discovered homographs.

usage:
  python preprocess_homographs.py data/neutral-voice-dataset-list.txt
  python preprocess_homographs.py data/*.txt --all
"""

import json
import re
import time
import argparse
import requests
from pathlib import Path
from collections import Counter
from tqdm import tqdm

# ── config ────────────────────────────────────────────────────────────────────

LLAMA_URL = "http://localhost:10000/completion"
SUFFIX    = "-clean"

KNOWN_HOMOGRAPHS = {
    "стоят", "стоит", "замок", "замки", "дома",
    "мою", "моют", "пили", "пила", "вина",
    "орган", "органы", "атлас", "полки", "полку",
    "белки", "белка", "угол", "лечу", "лечит",
}

# ── prompts ───────────────────────────────────────────────────────────────────

SYSTEM = (
    "Ты — редактор русских текстов. Ищешь омографы.\n\n"

    "Омограф — ТОЛЬКО слово, у которого:\n"
    "1) есть минимум 2 значения\n"
    "2) значения различаются УДАРЕНИЕМ\n"
    "3) написание одинаковое\n\n"

    "Если ты не можешь явно назвать ОБА значения с разным ударением — это НЕ омограф.\n"
    "Не угадывай. Если сомневаешься — не добавляй.\n\n"
    
    "НЕ являются омографами:\n"
    "- обычные глаголы (остыл, сделал, пошёл)\n"
    "- обычные существительные (ведро, тишина)\n"
    "- служебные слова (кто, это, в, и)\n"
    "- слова без известной пары с другим ударением\n\n"

    "Разрешённые ориентиры (примеры):\n"
    "замок, стоят, плачу, мука, атлас, орган\n\n"

    "Задача:\n"
    "1) Найди омографы в исходном предложении\n"
    "2) Перепиши предложение без них\n"
    "3) Можно полностью перефразировать\n"
    "4) В rewritten НЕ должно остаться омографов\n"
    "5) Если их нет — верни текст без изменений\n\n"

    "Перед добавлением слова в список задай себе вопрос: \"Какие ДВА значения и где ударение?\""
    "Если ответа нет — не добавляй."

    "Формат строго JSON:\n"
    "{\"rewritten\": \"...\", \"homographs\": [\"...\"]}\n"
)

def build_prompt(sentence: str) -> str:
    return (
        f"<|begin_of_text|>"
        f"<|start_header_id|>system<|end_header_id|>\n\n{SYSTEM}<|eot_id|>"
        f"<|start_header_id|>user<|end_header_id|>\n\n{sentence}<|eot_id|>"
        f"<|start_header_id|>assistant<|end_header_id|>\n\n"
    )

# ── llm ───────────────────────────────────────────────────────────────────────

def call_llm(sentence: str) -> dict | None:
    payload = {
        "prompt": build_prompt(sentence),
        "temperature": 0.2,
        "top_p": 0.9,
        "repeat_penalty": 1.1,
        "enable_thinking": True,
        "stop": ["<|eot_id|>", "<|end_of_text|>"],
        "session": None,
        "cache_prompt": False
    }
    try:
        resp = requests.post(LLAMA_URL, json=payload, timeout=60)
        resp.raise_for_status()
        raw   = resp.json().get("content", "").strip()
        clean = re.sub(r"```json|```", "", raw).strip()
        return json.loads(clean)
    except Exception:
        return None

# ── helpers ───────────────────────────────────────────────────────────────────

def has_known_homograph(sentence: str) -> bool:
    words = re.sub(r"[^\w\s]", "", sentence.lower()).split()
    return any(w in KNOWN_HOMOGRAPHS for w in words)

# ── main ──────────────────────────────────────────────────────────────────────

def process_file(path: Path, all_mode: bool):
    lines    = [l.strip() for l in path.read_text(encoding="utf-8").splitlines() if l.strip()]
    out_path = path.parent / (path.stem + SUFFIX + path.suffix)
    homo_out = path.parent / (path.stem + "-homographs.txt")

    # resume support — count already written lines
    done_count = 0
    if out_path.exists():
        done_count = sum(1 for l in out_path.read_text(encoding="utf-8").splitlines() if l.strip())
        print(f"  [*] resuming from line {done_count}/{len(lines)}")

    todo = lines[done_count:]
    if not todo:
        print(f"[=] {path.name} already complete")
        return

    print(f"\n[>] {path.name} — {len(todo)} remaining ({done_count} done)")

    changed    = 0
    failed     = 0
    discovered = Counter()

    out_f  = open(out_path,  "a", encoding="utf-8")
    homo_f = open(homo_out,  "a", encoding="utf-8")

    try:
        for line in tqdm(todo):
            if not all_mode and not has_known_homograph(line):
                out_f.write(line + "\n")
                out_f.flush()
                continue

            result = call_llm(line)

            if result is None:
                out_f.write(line + "\n")
                out_f.flush()
                failed += 1
                continue

            rewritten  = result.get("rewritten", line).strip()
            homographs = result.get("homographs", [])

            for h in homographs:
                if isinstance(h, str) and h.strip():
                    h = h.strip().lower()
                    discovered[h] += 1
                    homo_f.write(h + "\n")
                    homo_f.flush()

            final = rewritten if (rewritten and rewritten != line) else line
            if final != line:
                changed += 1

            out_f.write(final + "\n")
            out_f.flush()
            time.sleep(0.05)

    finally:
        out_f.close()
        homo_f.close()

    if discovered:
        print(f"\n  [*] {len(discovered)} unique homographs found → {homo_out}")
        print(f"      top 10: {', '.join(w for w, _ in discovered.most_common(10))}")
    else:
        print(f"  [*] no new homographs discovered")

    print(f"  [✓] {changed} rewritten, {failed} failed → {out_path}")


def main():
    parser = argparse.ArgumentParser(description="Rewrite homograph sentences via local LLM")
    parser.add_argument("files", nargs="+", help="txt files to process")
    parser.add_argument("--all", action="store_true",
                        help="process ALL sentences and report discovered homographs")
    args = parser.parse_args()

    for f in args.files:
        process_file(Path(f), all_mode=args.all)


if __name__ == "__main__":
    main()