"""PUBLIC INTERFACE of Data_Module. Other modules may import ONLY this file.

run_clean(cfg) -> dict
    1. surgically cleans chat.jsonl -> cleaned_data/sft_clean.jsonl (+ post-write verification)
    2. streams every SOURCE_REPOS entry (dataset first, model fallback), blacklist -> garbage filter
       -> context-aware dedup (seeded with chat.jsonl for cross-deduplication)
    3. mixes PRETRAIN_CONV_RATIO of chat samples into cleaned_data/pretrain_clean.jsonl
"""
from __future__ import annotations

import gc
import json
from pathlib import Path
from typing import Dict, List, Optional

from .internal.blacklist_gatekeeper import BlacklistGatekeeper
from .internal.context_deduplicator import ContextDeduplicator
from .internal.hf_universal_streamer import HFUniversalStreamer, RepoTelemetry
from .internal.pretrain_garbage_filter import PretrainGarbageFilter
from .internal.sft_surgical_cleaner import (
    build_pretrain_chat_samples, clean_chat_file, verify_pretrain_jsonl, verify_sft_jsonl, write_starter_chat,
)


def cleaned_paths(cfg: Dict) -> Dict[str, Path]:
    out = Path(cfg["root"]) / "cleaned_data"
    return {"dir": out, "sft": out / "sft_clean.jsonl", "pretrain": out / "pretrain_clean.jsonl",
            "rejected": out / "sft_rejected.jsonl"}


def ensure_chat_file(cfg: Dict) -> Path:
    """Create the 40-line starter chat.jsonl (37 clean + 3 intentionally broken lines) if it does not exist."""
    chat_file = Path(cfg["chat_file"])
    if write_starter_chat(chat_file):
        print(f"[Data] {chat_file.name} not found -> generated the 40-line starter dataset at {chat_file}")
        print("       (3 lines are broken/duplicate on purpose: the cleaner below will excise them and re-verify)")
    return chat_file


def cleaned_data_ready(cfg: Dict) -> bool:
    p = cleaned_paths(cfg)
    return p["sft"].exists() and p["pretrain"].exists() and p["sft"].stat().st_size > 0


def _parse_limit(value) -> Optional[int]:
    if value is None:
        return None
    if str(value).strip().lower() in ("full", "all", "none", ""):
        return None
    try:
        return max(1, int(value))
    except ValueError:
        return None


def _print_table(rows: List[RepoTelemetry]) -> None:
    header = ["Repo", "Type", "[Total Discovered]", "[Skipped by Blacklist]",
              "[Destroyed Garbage/Dup]", "[Kept Clean Samples]"]
    body = [[r.repo, r.repo_type, str(r.total_discovered), str(r.skipped_blacklist),
             str(r.destroyed), str(r.kept)] for r in rows]
    widths = [max(len(header[i]), *(len(row[i]) for row in body)) if body else len(header[i])
              for i in range(len(header))]
    line = "+" + "+".join("-" * (w + 2) for w in widths) + "+"
    print("\n   TELEMETRY  [Total Discovered] -> [Skipped by Blacklist] -> [Destroyed Garbage/Dup] -> [Kept]")
    print("   " + line)
    print("   | " + " | ".join(h.ljust(w) for h, w in zip(header, widths)) + " |")
    print("   " + line)
    for row in body:
        print("   | " + " | ".join(c.ljust(w) for c, w in zip(row, widths)) + " |")
    print("   " + line)
    for r in rows:
        print(f"   {r.repo}: garbage-gate={r.garbage_gate} too-large={r.too_large} read-err={r.read_errors} "
              f"filter-destroyed={r.filter_destroyed} dup-files={r.duplicate_files} "
              f"partial-dedup(sub-sections excised)={r.partial_dedup}"
              + (f"  ERROR: {r.error}" if r.error else ""))


def run_clean(cfg: Dict) -> Dict:
    paths = cleaned_paths(cfg)
    paths["dir"].mkdir(parents=True, exist_ok=True)
    token = cfg.get("hf_token") or None
    chat_file = Path(cfg["chat_file"])

    print("\n[Data] === STEP 1/3: SFT surgical cleaning of chat.jsonl ===")
    records, sft_report = clean_chat_file(chat_file, paths["sft"], paths["rejected"], verbose=True)
    print(f"   chat.jsonl: {sft_report['raw_lines']} lines -> {sft_report['kept']} clean "
          f"({sft_report['removed']} removed, {sft_report['repaired_modes']} mode repairs)")
    if not records:
        raise RuntimeError("No usable SFT records after cleaning chat.jsonl")

    print("\n[Data] === STEP 2/3: streaming + filtering source repos ===")
    gate = BlacklistGatekeeper(cfg.get("extra_blocked_extensions", []))
    garbage_filter = PretrainGarbageFilter(json_quota=cfg.get("json_quota", 0.15))
    dedup = ContextDeduplicator()
    for rec in records:                                    # cross-dedup: repo vs chat.jsonl
        dedup.seed_text("\n".join([rec["user"], rec["thinking"], rec["assistant"]]))

    telemetry: List[RepoTelemetry] = []
    repo_tmp = paths["dir"] / "_repo_text.tmp.jsonl"
    repo_docs = 0
    try:
        streamer = HFUniversalStreamer(token, gate)
    except Exception as exc:
        streamer = None
        print(f"   !! could not start HF streamer: {exc}")

    with open(repo_tmp, "w", encoding="utf-8") as sink:
        for entry in cfg.get("source_repos", []):
            repo = entry.get("repo") if isinstance(entry, dict) else str(entry)
            override = cfg.get("limit_override")
            limit = _parse_limit(override if override not in (None, "") else
                                 (entry.get("limit", 150) if isinstance(entry, dict) else 150))
            tel = RepoTelemetry(repo=repo)
            telemetry.append(tel)
            if streamer is None:
                tel.error = "streamer unavailable"
                continue
            try:
                for item in streamer.stream(repo, limit, tel):
                    cleaned = garbage_filter.process(item["path"], item["ext"], item["text"])
                    if cleaned is None:
                        tel.filter_destroyed += 1
                        continue
                    doc, info = dedup.process(item["path"], cleaned)
                    if doc is None:
                        tel.duplicate_files += 1
                        continue
                    if info["status"] == "partial":
                        tel.partial_dedup += 1
                    sink.write(json.dumps({"text": doc, "source": repo, "path": item["path"]},
                                          ensure_ascii=False) + "\n")
                    tel.kept += 1
                    repo_docs += 1
            except Exception as exc:
                tel.error = f"{type(exc).__name__}: {exc}"
                print(f"   !! stream error on {repo}: {tel.error}")
    _print_table(telemetry)
    if garbage_filter.stats:
        print("   filter stats: " + ", ".join(f"{k}={v}" for k, v in sorted(garbage_filter.stats.items())))
    if repo_docs == 0:
        print("   !! WARNING: no repo documents survived - pre-training will use chat.jsonl data only.")

    print("\n[Data] === STEP 3/3: building pretrain_clean.jsonl (repo text + chat mix) ===")
    ratio = float(cfg.get("conv_ratio", 0.25))
    chat_samples = build_pretrain_chat_samples(records, repo_docs, ratio)
    total = repo_docs + len(chat_samples)
    gap = max(1, repo_docs // len(chat_samples)) if chat_samples else 0
    written_chat = 0
    with open(paths["pretrain"], "w", encoding="utf-8") as out, open(repo_tmp, "r", encoding="utf-8") as src:
        for i, line in enumerate(src, 1):
            out.write(line)
            if chat_samples and written_chat < len(chat_samples) and i % gap == 0:
                out.write(json.dumps({"text": chat_samples[written_chat], "source": "chat.jsonl"},
                                     ensure_ascii=False) + "\n")
                written_chat += 1
        while written_chat < len(chat_samples):
            out.write(json.dumps({"text": chat_samples[written_chat], "source": "chat.jsonl"},
                                 ensure_ascii=False) + "\n")
            written_chat += 1
    try:
        repo_tmp.unlink()
    except OSError:
        pass
    check = verify_pretrain_jsonl(paths["pretrain"])
    print(f"   [PRETRAIN-VERIFY] {'PASSED' if check['ok'] else 'FAILED'}: {check['lines']} lines "
          f"({repo_docs} repo docs + {written_chat} chat samples = "
          f"{(100.0 * written_chat / max(total, 1)):.0f}% conversation share)")
    if not check["ok"]:
        raise RuntimeError("pretrain_clean.jsonl failed verification: " + "; ".join(check["errors"][:5]))
    sft_check = verify_sft_jsonl(paths["sft"])
    if not sft_check["ok"]:
        raise RuntimeError("sft_clean.jsonl failed final verification")

    summary = {"sft_records": len(records), "repo_docs": repo_docs, "chat_pretrain_samples": written_chat,
               "sft_path": str(paths["sft"]), "pretrain_path": str(paths["pretrain"])}
    del records, chat_samples, dedup, garbage_filter, telemetry
    gc.collect()
    print(f"[Data] done -> {paths['sft'].name}, {paths['pretrain'].name} (RAM released)")
    return summary