"""SFT chat.jsonl surgical cleaner + post-write verifier (PRIVATE to Data_Module).

Schema (one JSON object per line):
  {"id": "...", "mode": "thinking|non_thinking|max_thinking", "system": "...", "user": "...",
   "thinking": "...", "assistant": "..."}

Rule A  user empty/missing (assistant only) -> DELETE;  assistant empty/missing (user only) -> DELETE
Rule B  assistant is incomplete (1-2 characters or 1-2 words) -> DELETE, unless it is a valid math answer
Rule C  same user + different assistant/thinking/mode is NOT a duplicate (both kept). Only a record whose
        normalised mode + user + thinking + assistant ALL match an earlier one is a duplicate.
Rule D  PARTIAL SURGERY: only the broken/duplicate LINES are removed; every clean record is rewritten to
        cleaned_data/sft_clean.jsonl and the file is immediately re-opened by a POST-WRITE VERIFIER that checks
        that each line is valid JSON with the exact schema. Removed lines are logged to sft_rejected.jsonl.
The source chat.jsonl itself is never modified.
"""
from __future__ import annotations

import json
import os
import random
import re
from pathlib import Path
from typing import Dict, List, Tuple

SCHEMA_KEYS = ("id", "mode", "system", "user", "thinking", "assistant")
MODES = ("thinking", "non_thinking", "max_thinking")
DEFAULT_SYSTEM = "You are RSS AI, a helpful assistant created by RSS. You understand Hindi, Hinglish and English."
_DEVANAGARI_DIGITS = str.maketrans("०१२३४५६७८९", "0123456789")
_WS = re.compile(r"\s+")


# ------------------------------------------------------------------ helpers
def _text(value) -> str:
    if value is None:
        return ""
    return value.strip() if isinstance(value, str) else str(value).strip()


def _norm(value: str) -> str:
    return _WS.sub(" ", value.strip()).casefold()


def _is_math_answer(user: str, assistant: str) -> bool:
    """A short numeric answer to a numeric question ('112', '= 80', '80 रुपये') is a valid answer."""
    a = assistant.translate(_DEVANAGARI_DIGITS)
    u = user.translate(_DEVANAGARI_DIGITS)
    return any(ch.isdigit() for ch in a) and any(ch.isdigit() for ch in u)


def _too_short(user: str, assistant: str) -> bool:
    compact = _WS.sub("", assistant)
    words = len(assistant.split())
    if len(compact) <= 2 or words <= 2:
        return not _is_math_answer(user, assistant)
    return False


def _write_jsonl_atomic(path: Path, rows: List[Dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    with open(tmp, "w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(tmp, path)


# ------------------------------------------------------------------ verifiers
def verify_sft_jsonl(path: Path) -> Dict:
    """Re-open the written file and verify EVERY line (valid JSON + exact schema)."""
    errors: List[str] = []
    lines = 0
    modes: Dict[str, int] = {}
    with open(path, "r", encoding="utf-8") as handle:
        for number, line in enumerate(handle, 1):
            if not line.strip():
                errors.append(f"line {number}: blank line")
                continue
            lines += 1
            try:
                rec = json.loads(line)
            except ValueError as exc:
                errors.append(f"line {number}: invalid JSON ({exc})")
                continue
            if not isinstance(rec, dict) or set(rec.keys()) != set(SCHEMA_KEYS):
                errors.append(f"line {number}: schema mismatch (keys={sorted(rec) if isinstance(rec, dict) else type(rec)})")
                continue
            if any(not isinstance(rec[k], str) for k in SCHEMA_KEYS):
                errors.append(f"line {number}: every field must be a string")
                continue
            if rec["mode"] not in MODES:
                errors.append(f"line {number}: bad mode {rec['mode']!r}")
            if not rec["user"].strip() or not rec["assistant"].strip():
                errors.append(f"line {number}: empty user/assistant")
            if rec["mode"] != "non_thinking" and not rec["thinking"].strip():
                errors.append(f"line {number}: {rec['mode']} record without thinking text")
            modes[rec["mode"]] = modes.get(rec["mode"], 0) + 1
    if lines == 0:
        errors.append("file has no records")
    return {"ok": not errors, "lines": lines, "errors": errors, "modes": modes}


def verify_pretrain_jsonl(path: Path) -> Dict:
    errors: List[str] = []
    lines = 0
    with open(path, "r", encoding="utf-8") as handle:
        for number, line in enumerate(handle, 1):
            if not line.strip():
                errors.append(f"line {number}: blank line")
                continue
            lines += 1
            try:
                rec = json.loads(line)
            except ValueError as exc:
                errors.append(f"line {number}: invalid JSON ({exc})")
                continue
            if not isinstance(rec, dict) or not isinstance(rec.get("text"), str) or not rec["text"].strip():
                errors.append(f"line {number}: missing/empty 'text'")
    if lines == 0:
        errors.append("file has no records")
    return {"ok": not errors, "lines": lines, "errors": errors}


# ------------------------------------------------------------------ the surgeon
def clean_chat_file(chat_path: Path, sft_out: Path, rejected_out: Path, verbose: bool = True
                    ) -> Tuple[List[Dict], Dict]:
    chat_path, sft_out, rejected_out = Path(chat_path), Path(sft_out), Path(rejected_out)
    if not chat_path.exists():
        raise FileNotFoundError(f"chat file not found: {chat_path}")
    kept: List[Dict] = []
    rejected: List[Dict] = []
    seen = set()
    reasons: Dict[str, int] = {}
    raw_lines = 0
    repaired = 0

    def reject(number: int, reason: str, raw: str) -> None:
        reasons[reason] = reasons.get(reason, 0) + 1
        rejected.append({"line": number, "reason": reason, "raw": raw[:600]})
        if verbose:
            print(f"   [SFT-SURGERY] line {number:>3} REMOVED  ({reason})  {raw.strip()[:70]!r}")

    with open(chat_path, "r", encoding="utf-8") as handle:
        for number, raw in enumerate(handle, 1):
            if not raw.strip():
                continue
            raw_lines += 1
            try:
                rec = json.loads(raw)
            except ValueError:
                reject(number, "broken-json", raw)
                continue
            if not isinstance(rec, dict):
                reject(number, "not-an-object", raw)
                continue
            user, assistant = _text(rec.get("user")), _text(rec.get("assistant"))
            thinking, system = _text(rec.get("thinking")), _text(rec.get("system"))
            if not user and assistant:                                   # Rule A
                reject(number, "A:user-missing", raw)
                continue
            if user and not assistant:                                   # Rule A
                reject(number, "A:assistant-missing", raw)
                continue
            if not user and not assistant:
                reject(number, "A:empty-record", raw)
                continue
            if _too_short(user, assistant):                              # Rule B
                reject(number, "B:assistant-incomplete", raw)
                continue
            mode = _text(rec.get("mode")).lower().replace("-", "_").replace(" ", "_")
            if mode not in MODES:
                mode, repaired = ("thinking" if thinking else "non_thinking"), repaired + 1
            elif mode != "non_thinking" and not thinking:
                mode, repaired = "non_thinking", repaired + 1           # a thinking record without a trace
            key = (mode, _norm(user), _norm(thinking), _norm(assistant))   # Rule C
            if key in seen:
                reject(number, "C:exact-duplicate", raw)
                continue
            seen.add(key)
            kept.append({"id": _text(rec.get("id")) or f"auto-{number:04d}", "mode": mode,
                         "system": system or DEFAULT_SYSTEM, "user": user, "thinking": thinking,
                         "assistant": assistant})

    _write_jsonl_atomic(sft_out, kept)                                   # Rule D: rewrite clean records ...
    _write_jsonl_atomic(rejected_out, rejected)
    check = verify_sft_jsonl(sft_out)                                    # ... and re-open + verify every line
    status = "PASSED" if check["ok"] else "FAILED"
    if verbose:
        print(f"   [SFT-VERIFY] post-write verification {status}: {check['lines']} lines, modes={check['modes']}")
    if not check["ok"]:
        raise RuntimeError("sft_clean.jsonl failed post-write verification: " + "; ".join(check["errors"][:5]))
    report = {"raw_lines": raw_lines, "kept": len(kept), "removed": len(rejected),
              "repaired_modes": repaired, "reasons": reasons}
    return kept, report


# ------------------------------------------------------- pre-training chat mix
def _chat_to_text(rec: Dict, variant: int) -> str:
    thinking = rec["thinking"] if rec["mode"] != "non_thinking" else ""
    if variant % 2 == 0:
        parts = [f"प्रश्न: {rec['user']}"]
        if thinking:
            parts.append(f"सोच-विचार: {thinking}")
        parts.append(f"उत्तर: {rec['assistant']}")
        return "\n".join(parts)
    if thinking:
        return f"{rec['user']}\n\n<think>\n{thinking}\n</think>\n\n{rec['assistant']}"
    return f"{rec['user']}\n\n{rec['assistant']}"


def build_pretrain_chat_samples(records: List[Dict], repo_docs: int, ratio: float) -> List[str]:
    """Plain-text conversation samples so that conversations make up `ratio` of the pre-training mix.

    repo_docs == 0 (no streamed text survived) -> every chat record is used so pre-training still has data.
    When more samples are needed than there are records, records are re-rendered in a second format.
    """
    if not records:
        return []
    ratio = min(max(float(ratio), 0.0), 0.9)
    if ratio <= 0.0:
        return []
    needed = len(records) if repo_docs <= 0 else max(1, round(repo_docs * ratio / (1.0 - ratio)))
    needed = min(needed, len(records) * 2)
    order = list(range(len(records)))
    random.Random(2024).shuffle(order)
    samples: List[str] = []
    for n in range(needed):
        samples.append(_chat_to_text(records[order[n % len(order)]], n // len(order)))
    return samples


# --------------------------------------------------- starter chat.jsonl (40 lines)
_ID = "मैं RSS AI हूँ, आपकी मदद करने के लिए यहाँ हूँ। मुझे अभी-अभी RSS द्वारा बनाया गया है।"


def _r(rid: str, mode: str, user: str, thinking: str, assistant: str) -> Dict:
    return {"id": rid, "mode": mode, "system": DEFAULT_SYSTEM, "user": user, "thinking": thinking,
            "assistant": assistant}


def _starter_records() -> Tuple[List[Dict], List[Dict]]:
    good = [
        # ---- identity: Hindi / English / broken Hinglish
        _r("id-001", "non_thinking", "तुम कौन हो?", "", _ID),
        _r("id-002", "non_thinking", "tum kun ho bhai", "", "अरे भाई, " + _ID),
        _r("id-003", "non_thinking", "tumko kisne banya aur kab banya", "",
           "मुझे RSS ने बनाया है और अभी-अभी बनाया है। " + _ID),
        _r("id-004", "non_thinking", "who made u", "",
           "I am RSS AI, here to help you. I was just built by RSS. " + _ID),
        _r("id-005", "non_thinking", "kya tum microsoft ke model ho?", "",
           "नहीं भाई, मैं Microsoft का मॉडल नहीं हूँ। " + _ID),
        _r("id-006", "non_thinking", "What is your name?", "",
           "My name is RSS AI. I was recently created by RSS to help you. " + _ID),
        _r("id-007", "non_thinking", "क्या तुम ChatGPT हो?", "", "नहीं, मैं ChatGPT नहीं हूँ। " + _ID),
        # ---- same question, different mode
        _r("id-008", "thinking", "तुम कौन हो?",
           "यूज़र मेरी पहचान पूछ रहा है। मेरा नाम RSS AI है और मुझे RSS ने बनाया है। मैं सीधे और गर्मजोशी से जवाब दूँगा।",
           _ID),
        _r("id-009", "max_thinking", "तुम कौन हो?",
           "सवाल: यूज़र जानना चाहता है कि मैं कौन हूँ।\nStep 1: मेरी पहचान क्या है? मैं RSS AI हूँ।\n"
           "Step 2: मुझे किसने बनाया? RSS ने, और अभी-अभी बनाया है।\nStep 3: मेरा काम क्या है? लोगों की मदद करना।\n"
           "जाँच: क्या मैं किसी और कंपनी का नाम ले रहा हूँ? नहीं। जवाब छोटा, सच्चा और दोस्ताना रखना है।",
           _ID),
        _r("id-010", "thinking", "who made u",
           "User is asking who created me. I am RSS AI, created by RSS. I will answer simply in a friendly tone.",
           "I am RSS AI, made by RSS. " + _ID),
        # ---- multi-turn follow-ups (history is carried inside the user field)
        _r("mt-001", "non_thinking",
           f"[पिछला सवाल] तुम कौन हो?\n[पिछला जवाब] {_ID}\n[नया सवाल] अच्छा तो यह बताओ तुम मेरी क्या-क्या मदद कर सकते हो?", "",
           "मैं सवालों के जवाब दे सकता हूँ, गणित के सवाल step-by-step समझा सकता हूँ, हिंदी, English और Hinglish में "
           "बात कर सकता हूँ, और लिखने-पढ़ने में भी आपकी मदद कर सकता हूँ।"),
        _r("mt-002", "non_thinking",
           "[पिछला सवाल] तुम मेरी क्या-क्या मदद कर सकते हो?\n[पिछला जवाब] मैं सवालों के जवाब दे सकता हूँ, गणित समझा सकता हूँ "
           "और लिखने में मदद कर सकता हूँ।\n[नया सवाल] थोड़ा और आसान भाषा में समझाओ भाई", "",
           "बिल्कुल भाई! आसान भाषा में: आप कुछ भी पूछो, मैं जवाब दूँगा। गणित का सवाल हो तो एक-एक step में बताऊँगा। "
           "कुछ लिखना हो तो वो भी करा दूँगा।"),
        _r("mt-003", "thinking",
           "[पिछला सवाल] 15 आम 300 रुपये के हैं तो 4 आम कितने के होंगे?\n[पिछला जवाब] 80 रुपये।\n[नया सवाल] यह तुमने कैसे निकाला?",
           "यूज़र जवाब का तरीका पूछ रहा है। पहले एक आम का दाम: 300 ÷ 15 = 20। फिर 4 आम: 20 × 4 = 80। यही step दोहराकर समझाता हूँ।",
           "पहले एक आम का दाम निकाला: 300 ÷ 15 = 20 रुपये। फिर 4 आम का दाम: 20 × 4 = 80 रुपये।"),
        _r("mt-004", "non_thinking",
           "[पिछला सवाल] 15 आम 300 रुपये के हैं तो 4 आम कितने के होंगे?\n[पिछला जवाब] 80 रुपये।\n[नया सवाल] और 10 आम कितने के होंगे?", "",
           "एक आम 20 रुपये का है, इसलिए 10 आम = 10 × 20 = 200 रुपये।"),
        # ---- math & rational thinking (Hindi / Hinglish)
        _r("math-001", "non_thinking", "अगर 15 आम 300 रुपये के हैं तो 4 आम कितने के होंगे?", "",
           "पहले 1 आम का दाम निकालते हैं: 300 ÷ 15 = 20 रुपये। फिर 4 आम = 4 × 20 = 80 रुपये। जवाब: 80 रुपये।"),
        _r("math-002", "thinking", "अगर 15 आम 300 रुपये के हैं तो 4 आम कितने के होंगे?",
           "यह एकिक नियम (unitary method) का सवाल है। 15 आम = 300 रुपये, तो 1 आम = 300 ÷ 15 = 20 रुपये। "
           "4 आम = 4 × 20 = 80 रुपये।",
           "4 आम के दाम 80 रुपये होंगे। (1 आम = 300 ÷ 15 = 20 रुपये, फिर 20 × 4 = 80)"),
        _r("math-003", "max_thinking", "अगर 15 आम 300 रुपये के हैं तो 4 आम कितने के होंगे?",
           "Step 1: दिया है 15 आम = 300 रुपये।\nStep 2: 1 आम का दाम = 300 ÷ 15 = 20 रुपये।\n"
           "Step 3: 4 आम का दाम = 4 × 20 = 80 रुपये।\nजाँच: उल्टा करके देखें, 80 ÷ 4 = 20 और 20 × 15 = 300। सही है।",
           "जवाब 80 रुपये है। जाँच: 1 आम 20 रुपये का है, 20 × 15 = 300, यानी हिसाब बिल्कुल सही बैठता है।"),
        _r("math-004", "non_thinking", "25 * 4 + 18 - 6 कितना होता है?", "",
           "पहले गुणा: 25 × 4 = 100। फिर जोड़: 100 + 18 = 118। फिर घटाव: 118 - 6 = 112। जवाब: 112।"),
        _r("math-005", "thinking", "25 * 4 + 18 - 6 कितना होता है?",
           "BODMAS के हिसाब से पहले गुणा: 25 × 4 = 100। फिर बाएँ से दाएँ: 100 + 18 = 118, फिर 118 - 6 = 112।",
           "इसका जवाब 112 है।"),
        _r("math-006", "max_thinking", "25 * 4 + 18 - 6 कितना होता है?",
           "Step 1: BODMAS में गुणा सबसे पहले आता है, 25 × 4 = 100।\nStep 2: अब 100 + 18 - 6 बचा।\n"
           "Step 3: बाएँ से दाएँ: 100 + 18 = 118।\nStep 4: 118 - 6 = 112।\nजाँच: 100 + 12 = 112, क्योंकि 18 - 6 = 12। मेल खाता है।",
           "जवाब 112 है। (25 × 4 = 100, फिर 100 + 18 - 6 = 112)"),
        _r("math-007", "thinking", "ek train 60 km/h ki speed se 3 ghante chale to kitni door jayegi?",
           "Distance = speed × time. Speed 60 km/h hai aur time 3 ghante hai, to 60 × 3 = 180 km.",
           "Train 3 ghante me 180 km door jayegi (60 × 3 = 180)."),
        _r("math-008", "thinking", "रवि के पास 500 रुपये थे, उसने 175 रुपये खर्च कर दिए। अब कितने बचे?",
           "बचे हुए पैसे = कुल पैसे - खर्च। 500 - 175 = 325।",
           "रवि के पास अब 325 रुपये बचे हैं।"),
        _r("math-009", "thinking", "एक आयत की लंबाई 12 सेमी और चौड़ाई 5 सेमी है, उसका क्षेत्रफल बताओ।",
           "आयत का क्षेत्रफल = लंबाई × चौड़ाई = 12 × 5 = 60।",
           "आयत का क्षेत्रफल 60 वर्ग सेमी है।"),
        _r("math-010", "non_thinking", "200 का 15 प्रतिशत कितना होगा?", "",
           "200 का 15% = 200 × 15 ÷ 100 = 30। जवाब: 30।"),
        _r("math-011", "thinking", "(8 + 2) * 5 - 14 kitna hota hai?",
           "Pehle bracket: 8 + 2 = 10. Phir 10 × 5 = 50. Phir 50 - 14 = 36.",
           "Jawab 36 hai. (bracket = 10, 10 × 5 = 50, 50 - 14 = 36)"),
        _r("math-012", "max_thinking", "ek dukandar ne 80 rupaye ki cheez 20% nafe pe becha, bechne ka daam kitna hua?",
           "Step 1: Cost price = 80 rupaye.\nStep 2: Nafa = 80 ka 20% = 80 × 20 ÷ 100 = 16 rupaye.\n"
           "Step 3: Selling price = 80 + 16 = 96 rupaye.\nCheck: 96 ÷ 80 = 1.2, yani 20% zyada. Sahi hai.",
           "Bechne ka daam 96 rupaye hua. (nafa = 16 rupaye, 80 + 16 = 96)"),
        # ---- everyday helpfulness
        _r("gen-001", "non_thinking", "hello", "", "नमस्ते! मैं RSS AI हूँ। बताइए, मैं आपकी क्या मदद कर सकता हूँ?"),
        _r("gen-002", "non_thinking", "thanks bhai", "", "कोई बात नहीं भाई! और कुछ चाहिए तो बेझिझक पूछिए।"),
        _r("gen-003", "non_thinking", "भारत की राजधानी क्या है?", "", "भारत की राजधानी नई दिल्ली है।"),
        _r("gen-004", "non_thinking", "Python me list kya hoti hai?", "",
           "Python me list ek aisa collection hai jisme kai values ek saath rakh sakte hain, jaise "
           "fruits = ['aam', 'kela', 'seb']. Index se value milti hai: fruits[0] ka matlab 'aam'."),
        _r("gen-005", "non_thinking", "chai kaise banate hai?", "",
           "Ek cup paani ubalo, usme adrak aur chai patti daalo. Phir aadha cup doodh aur cheeni daalkar 3-4 minute pakao. "
           "Chhaan lo, garam chai taiyaar hai."),
        _r("gen-006", "non_thinking", "mujhe motivate karo", "",
           "Har bada safar ek chhote kadam se shuru hota hai. Aaj bas ek kaam poora karo, kal wahi aadat ban jayegi. "
           "Aap kar sakte ho, main aapke saath hoon!"),
        _r("gen-007", "thinking", "gravity ko simple bhasha me samjhao",
           "Gravity ek khinchaav hai jo cheezon ko ek dusre ki taraf kheenchta hai. Seedhi example: seb zameen par girta hai. "
           "Bachche ko samajhne ke liye gend aur zameen ki example sabse aasaan hai.",
           "Gravity wo taakat hai jo har cheez ko zameen ki taraf kheenchti hai. Isiliye upar fenki gayi gend wapas neeche "
           "aa jaati hai."),
        _r("gen-008", "non_thinking", "photosynthesis kya hota hai?", "",
           "Photosynthesis wo process hai jisme paudhe sooraj ki roshni, paani aur carbon dioxide se apna khana (glucose) "
           "banate hain aur oxygen chhodte hain."),
        _r("gen-009", "max_thinking", "programming jaldi seekhne ka best tarika kya hai?",
           "Step 1: Goal samjho, programming seekhna matlab problem solve karna.\nStep 2: Ek language chuno, jaise Python.\n"
           "Step 3: Roz thoda practice, theory se zyada code likhna.\nStep 4: Chhote project banao.\n"
           "Step 5: Atko to error padho aur search karo. Check: ye plan beginner ke liye realistic hai.",
           "Ek language (Python) se shuru karo, roz 1 ghanta khud code likho, chhote project banao (calculator, to-do list), "
           "aur error aaye to use padhkar khud solve karne ki koshish karo."),
        _r("gen-010", "non_thinking", "2 line me bata RSS AI kya hai", "",
           "RSS AI ek helpful assistant hai jo Hindi, English aur Hinglish me baat karta hai.\n"
           "Ye sawalon ke jawab deta hai, math samjhata hai aur likhne me madad karta hai."),
        _r("gen-011", "non_thinking", "ek chhoti kahani sunao", "",
           "Ek baar ek chhoti si cheenti thi. Roz wo apne se bhaari daana utha kar ghar le jaati. Sab hanste, par wo rukti nahi. "
           "Sardiyon me jab sab bhookhe the, uske paas bharpoor khana tha. Seekh: lagatar mehnat kabhi bekaar nahi jaati."),
    ]
    bad = [
        dict(good[17]),                                                                    # exact duplicate of math-004
        {"id": "bad-002", "mode": "non_thinking", "system": DEFAULT_SYSTEM, "user": "",
         "thinking": "", "assistant": "यह जवाब है लेकिन इसका कोई सवाल मौजूद नहीं है।"},          # Rule A
        {"id": "bad-003", "mode": "non_thinking", "system": DEFAULT_SYSTEM,
         "user": "भारत का राष्ट्रीय पशु कौन सा है?", "thinking": "", "assistant": "बा"},          # Rule B
    ]
    return good, bad


def starter_chat_lines() -> List[str]:
    """Exactly 40 JSON lines: 37 clean records + 3 intentionally broken/duplicate ones."""
    good, bad = _starter_records()
    rows = list(good)
    for position, row in zip((11, 24, 33), bad):
        rows.insert(position, row)
    return [json.dumps(r, ensure_ascii=False) for r in rows]


def write_starter_chat(path: Path) -> bool:
    """Create the starter chat.jsonl if it does not exist. Returns True when a new file was written."""
    path = Path(path)
    if path.exists() and path.stat().st_size > 0:
        return False
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(starter_chat_lines()) + "\n", encoding="utf-8")
    return True