"""Human feedback loop (RLHF-style data + live translation memory).

Every rating or correction a visitor makes on a translated line is stored as one JSON line:
  {time, job, seg, src, tgt, engine, original, machine, rating, correction, voter, trusted}

That data is used three ways:
1. Translation memory - a correction confirmed by enough different people (or by the owner)
   replaces the machine output for that exact sentence in every future job.
2. Few-shot examples - the Claude engine gets recent approved corrections for the language pair.
3. Training data - /api/feedback/export gives preference pairs (machine vs. human) for
   fine-tuning NLLB / Whisper or training a reward model later.

Persistence: the Space's disk is wiped on restart, so when OXT_FEEDBACK_REPO and HF_TOKEN are set,
the feedback folder is pushed to a private Hugging Face dataset every few minutes and restored on boot.
"""
from __future__ import annotations

import hashlib
import json
import os
import re
import threading
import time
import uuid
from collections import defaultdict
from pathlib import Path

from .engines import DATA

FB_DIR = DATA / "feedback"
FB_DIR.mkdir(parents=True, exist_ok=True)
REPO = os.environ.get("OXT_FEEDBACK_REPO", "")          # e.g. MsThai/switchlingo-feedback
MIN_VOTES = int(os.environ.get("OXT_TM_MIN_VOTES", "2"))  # distinct people needed to trust a fix
ADMIN_TOKEN = os.environ.get("OXT_ADMIN_TOKEN", "")
_FILE = FB_DIR / f"fb-{time.strftime('%Y%m%d')}-{uuid.uuid4().hex[:6]}.jsonl"
_lock = threading.Lock()
_scheduler = None
_sync_error = ""

# (src, tgt, normalized original) -> {correction text -> set(voters)} ; trusted fixes win outright
_votes: dict[tuple, dict[str, set]] = defaultdict(lambda: defaultdict(set))
_trusted: dict[tuple, str] = {}


def _norm(s: str) -> str:
    return re.sub(r"\s+", " ", s.strip().lower())


def _key(src: str, tgt: str, original: str) -> tuple:
    return (src, tgt, _norm(original))


def voter_id(ip: str) -> str:
    return hashlib.sha256(("oxt" + ip).encode()).hexdigest()[:16]  # never store raw IPs


def _learn(rec: dict) -> None:
    if rec.get("kind") == "pair":  # contributed sentence pair counts as a vote both ways
        if rec.get("dialect"):  # dialects are kept for training only, not mixed into the parent language
            return
        a, b = rec["en"].strip(), rec["text"].strip()
        lang = rec["lang"]
        if a and b:
            _votes[_key("en", lang, a)][b].add(rec.get("voter", "?"))
            _votes[_key(lang, "en", b)][a].add(rec.get("voter", "?"))
            if rec.get("trusted"):
                _trusted[_key("en", lang, a)] = b
                _trusted[_key(lang, "en", b)] = a
        return
    fix = (rec.get("correction") or "").strip()
    if not fix:
        return
    k = _key(rec["src"], rec["tgt"], rec["original"])
    if rec.get("trusted"):
        _trusted[k] = fix
    else:
        _votes[k][fix].add(rec.get("voter", "?"))


def _restore() -> None:
    if REPO and os.environ.get("HF_TOKEN"):
        try:
            from huggingface_hub import snapshot_download
            snapshot_download(REPO, repo_type="dataset", local_dir=FB_DIR, allow_patterns=["*.jsonl"])
        except Exception as e:  # noqa: BLE001 - first run: repo may be empty
            print("feedback restore skipped:", e)
    for p in sorted(FB_DIR.glob("fb-*.jsonl")):
        for line in p.read_text(encoding="utf-8").splitlines():
            try:
                _learn(json.loads(line))
            except ValueError:
                pass


def start() -> None:
    """Call once at server start."""
    global _scheduler, _sync_error
    _restore()
    if REPO and not os.environ.get("HF_TOKEN"):
        _sync_error = "HF_TOKEN secret is missing in the Space"
    if REPO and os.environ.get("HF_TOKEN"):
        try:
            from huggingface_hub import CommitScheduler
            _scheduler = CommitScheduler(repo_id=REPO, repo_type="dataset", folder_path=FB_DIR,
                                         path_in_repo=".", every=5, private=True,
                                         allow_patterns=["*.jsonl", "audio/*"])
        except Exception as e:  # noqa: BLE001
            _sync_error = f"{type(e).__name__}: {str(e)[:200]}"
            print("feedback sync disabled:", e)


def add(rec: dict) -> None:
    rec["time"] = time.time()
    lock = _scheduler.lock if _scheduler else _lock
    with lock:
        with open(_FILE, "a", encoding="utf-8") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    _learn(rec)


def lookup(src: str, tgt: str, original: str) -> str | None:
    """Human-approved translation for this exact sentence, if there is one."""
    k = _key(src, tgt, original)
    if k in _trusted:
        return _trusted[k]
    best = max(_votes.get(k, {}).items(), key=lambda kv: len(kv[1]), default=None)
    return best[0] if best and len(best[1]) >= MIN_VOTES else None


def examples(src: str, tgt: str, n: int = 15) -> list[dict]:
    """Approved corrections for this pair, for few-shot prompting."""
    out = [{"original": k[2], "translation": v} for k, v in _trusted.items() if k[:2] == (src, tgt)]
    for k in _votes:
        if k[:2] == (src, tgt) and k not in _trusted:
            fix = lookup(*k)
            if fix:
                out.append({"original": k[2], "translation": fix})
    return out[-n:]


AUDIO_DIR = FB_DIR / "audio"
AUDIO_DIR.mkdir(exist_ok=True)


def add_pair(lang: str, en: str, text: str, voter: str, trusted: bool, dialect: str = "",
             audio: bytes | None = None, audio_ext: str = "webm") -> None:
    rec = {"kind": "pair", "lang": lang, "dialect": dialect, "en": en, "text": text,
           "voter": voter, "trusted": trusted}
    if audio:
        name = f"{uuid.uuid4().hex[:12]}.{audio_ext}"
        lock = _scheduler.lock if _scheduler else _lock
        with lock:
            (AUDIO_DIR / name).write_bytes(audio)
        rec["audio"] = f"audio/{name}"
    add(rec)


def stats() -> dict:
    n = good = bad = fixes = pairs = 0
    by_lang: dict[str, int] = defaultdict(int)
    for p in FB_DIR.glob("fb-*.jsonl"):
        for line in p.read_text(encoding="utf-8").splitlines():
            try:
                r = json.loads(line)
            except ValueError:
                continue
            if r.get("kind") == "pair":
                pairs += 1
                by_lang[r.get("dialect") or r.get("lang")] += 1
                continue
            n += 1
            good += r.get("rating") == 1
            bad += r.get("rating") == -1
            fixes += bool(r.get("correction"))
    return {"ratings": n, "sentence_pairs": pairs, "pairs_by_language": dict(by_lang), "thumbs_up": good, "thumbs_down": bad, "corrections": fixes,
            "approved_fixes": len(_trusted) + sum(1 for k in _votes if lookup(*k)),
            "synced_to": REPO or None, "sync_active": _scheduler is not None,
            "sync_error": _sync_error or None}


def export_lines():
    for p in sorted(FB_DIR.glob("fb-*.jsonl")):
        yield from (line + "\n" for line in p.read_text(encoding="utf-8").splitlines() if line.strip())