#!/usr/bin/env python3
# DGX AI Factory - 14B Repair Data Problem Audit
#
# Purpose:
# - Do not train.
# - Do not move adapters.
# - Audit why operational_repair_v3 made the model worse.
# - Analyze:
#   1) repair dataset repetition / template collapse
#   2) failed smoke report
#   3) training script risk patterns
#   4) next safe direction
#
# Outputs:
#   reports/14b_repair_data_problem_audit_TIMESTAMP.json
#   reports/14b_repair_data_problem_audit_TIMESTAMP.md

import json
import glob
import hashlib
import statistics
import re
import time
from pathlib import Path
from collections import Counter, defaultdict

ROOT = Path("/home/harness_user_1/dgx_ai_factory")
REPORT_DIR = ROOT / "reports"
REPORT_DIR.mkdir(parents=True, exist_ok=True)

DATASET = ROOT / "data/14b_operational_repair_v3/train_operational_repair_v3.jsonl"
TRAIN_SCRIPT = ROOT / "scripts/117_train_14b_operational_repair_v3.py"

CHAMPION_POINTER = ROOT / "state/current_14b_champion_adapter.txt"
CHAMPION_STATE_GLOB = str(ROOT / "reports/14b_champion_state_*.json")
OP_SMOKE_GLOB = str(ROOT / "reports/operational_smoke_after_repair_v3_*.json")
TENK_FAILED_REPORT = ROOT / "reports/14b10k_auto_quick_benchmark_20260628_163530.json"
BASELINE_REPORT = ROOT / "reports/14b_canary_direct_benchmark_20260627_042507.json"

def sha(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()

def read_json(path: Path):
    if not path or not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None

def latest_path(pattern: str):
    files = glob.glob(pattern)
    if not files:
        return None
    return Path(max(files, key=lambda x: Path(x).stat().st_mtime))

def load_jsonl(path: Path):
    rows = []
    bad = 0
    if not path.exists():
        return rows, 0
    with open(path, "r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
                obj["_line_no"] = line_no
                rows.append(obj)
            except Exception:
                bad += 1
    return rows, bad

def text_len_stats(values):
    if not values:
        return {}
    return {
        "min": min(values),
        "max": max(values),
        "mean": round(statistics.mean(values), 2),
        "median": round(statistics.median(values), 2),
    }

def audit_dataset(rows, bad_json):
    category_counts = Counter(str(r.get("category", "unknown")) for r in rows)

    instruction_hashes = Counter()
    output_hashes = Counter()
    pair_hashes = Counter()
    output_to_rows = defaultdict(list)

    output_lens = []
    instruction_lens = []

    required = {
        "lora": ["adapter_config.json", "adapter_model.safetensors", "PEFT", "smoke", "benchmark", "stable", "rejected"],
        "systemd": ["217/USER", "User=", "getent passwd", "id ", "systemctl", "journalctl", "daemon-reload"],
        "dataset": ["category", "quota", "중복", "CJK", "template", "balanced", "train"],
    }

    missing_by_category = defaultdict(Counter)

    for r in rows:
        ins = str(r.get("instruction", ""))
        out = str(r.get("output", ""))
        cat = str(r.get("category", "unknown"))

        ih = sha(ins)
        oh = sha(out)
        ph = sha(ins + "\n<<<SEP>>>\n" + out)

        instruction_hashes[ih] += 1
        output_hashes[oh] += 1
        pair_hashes[ph] += 1
        output_to_rows[oh].append({
            "line": r.get("_line_no"),
            "category": cat,
            "instruction": ins,
            "output_preview": out[:180].replace("\n", " "),
        })

        instruction_lens.append(len(ins))
        output_lens.append(len(out))

        for token in required.get(cat, []):
            if token.lower() not in out.lower():
                missing_by_category[cat][token] += 1

    top_reused_outputs = []
    for h, cnt in output_hashes.most_common(10):
        if cnt > 1:
            rows_for_hash = output_to_rows[h]
            cats = Counter(x["category"] for x in rows_for_hash)
            top_reused_outputs.append({
                "output_hash": h,
                "count": cnt,
                "categories": dict(cats),
                "sample_rows": rows_for_hash[:5],
            })

    unique_outputs = len(output_hashes)
    total = len(rows)
    output_reuse_ratio = round(1 - (unique_outputs / total), 4) if total else 0

    repeated_output_problem = (
        total >= 20 and unique_outputs <= max(5, total // 5)
    )

    return {
        "path": str(DATASET),
        "rows": total,
        "bad_json": bad_json,
        "category_counts": dict(category_counts),
        "unique_instructions": len(instruction_hashes),
        "unique_outputs": unique_outputs,
        "unique_pairs": len(pair_hashes),
        "duplicate_pairs": sum(cnt - 1 for cnt in pair_hashes.values() if cnt > 1),
        "output_reuse_ratio": output_reuse_ratio,
        "instruction_length": text_len_stats(instruction_lens),
        "output_length": text_len_stats(output_lens),
        "top_reused_outputs": top_reused_outputs,
        "missing_required_by_category": {k: dict(v) for k, v in missing_by_category.items()},
        "repeated_output_problem": repeated_output_problem,
        "diagnosis": (
            "DATASET_TEMPLATE_COLLAPSE_RISK"
            if repeated_output_problem
            else "DATASET_REPETITION_ACCEPTABLE"
        ),
    }

def audit_smoke(report):
    if not report:
        return {"found": False}

    results = report.get("results", [])
    category_scores = defaultdict(list)
    low_results = []
    flag_counts = Counter()

    for r in results:
        cat = str(r.get("category", "unknown"))
        score = r.get("score")
        if isinstance(score, (int, float)):
            category_scores[cat].append(score)

        for fl in r.get("flags", []):
            flag_counts[str(fl)] += 1

        if isinstance(score, (int, float)) and score < 70:
            low_results.append({
                "name": r.get("name"),
                "category": cat,
                "score": score,
                "flags": r.get("flags", []),
                "prompt": r.get("prompt"),
                "output_preview": str(r.get("output", ""))[:1200],
            })

    category_avg = {
        cat: round(sum(vals) / len(vals), 2)
        for cat, vals in sorted(category_scores.items())
        if vals
    }

    return {
        "found": True,
        "smoke_pass": report.get("smoke_pass"),
        "smoke_average": report.get("smoke_average"),
        "category_avg": category_avg,
        "decision": report.get("decision"),
        "low_results": low_results,
        "flag_counts": dict(flag_counts.most_common()),
    }

def audit_training_script(path: Path):
    if not path.exists():
        return {"found": False}

    text = path.read_text(encoding="utf-8", errors="replace")

    findings = {
        "found": True,
        "path": str(path),
        "labels_all_tokens": 'enc["labels"] = enc["input_ids"].copy()' in text or "enc['labels'] = enc['input_ids'].copy()" in text,
        "assistant_only_loss": ("-100" in text and "labels" in text and "assistant" in text.lower()),
        "gradient_checkpointing": "gradient_checkpointing=True" in text or "gradient_checkpointing_enable" in text,
        "max_steps_40": 'MAX_STEPS", "40"' in text or "MAX_STEPS = int(os.environ.get(\"MAX_STEPS\", \"40\"))" in text,
        "lr_8e_7": "8e-7" in text,
        "no_replay_detected": "replay" not in text.lower(),
        "uses_full_prompt_chat_template": "apply_chat_template" in text and "assistant" in text,
    }

    risks = []
    if findings["labels_all_tokens"] and not findings["assistant_only_loss"]:
        risks.append("FULL_SEQUENCE_LABELS_RISK: user/system tokens are also trained, not assistant-only loss.")
    if findings["no_replay_detected"]:
        risks.append("NO_REPLAY_RISK: repair-only training can overwrite useful champion behavior.")
    if findings["max_steps_40"]:
        risks.append("MAX_STEPS_40_ON_60_ROWS: small repetitive dataset may be over-emphasized.")
    if findings["lr_8e_7"]:
        risks.append("LR_8E_7_NOT_HUGE_BUT_STILL_ACTIVE: conservative, but dataset repetition can still dominate.")

    findings["risks"] = risks
    return findings

def summarize_reports():
    baseline = read_json(BASELINE_REPORT)
    failed10k = read_json(TENK_FAILED_REPORT)
    latest_smoke_path = latest_path(OP_SMOKE_GLOB)
    smoke = read_json(latest_smoke_path) if latest_smoke_path else None
    latest_state_path = latest_path(CHAMPION_STATE_GLOB)
    state = read_json(latest_state_path) if latest_state_path else None

    champion = None
    if CHAMPION_POINTER.exists():
        champion = CHAMPION_POINTER.read_text(encoding="utf-8").strip()

    return {
        "champion_adapter": champion,
        "champion_state_report": str(latest_state_path) if latest_state_path else None,
        "champion_state_decision": state.get("decision") if state else None,
        "baseline": {
            "path": str(BASELINE_REPORT),
            "average_score": baseline.get("average_score") if baseline else None,
            "decision": baseline.get("decision") if baseline else None,
            "category_avg": baseline.get("category_avg") if baseline else None,
        },
        "failed_10k": {
            "path": str(TENK_FAILED_REPORT),
            "average_score": failed10k.get("average_score") if failed10k else None,
            "decision": failed10k.get("decision") if failed10k else None,
            "category_avg": failed10k.get("category_avg") if failed10k else None,
        },
        "operational_smoke_path": str(latest_smoke_path) if latest_smoke_path else None,
        "operational_smoke": audit_smoke(smoke),
    }

def make_markdown(report):
    lines = []
    lines.append("# 14B Repair Data Problem Audit")
    lines.append("")
    lines.append(f"Timestamp: `{report['timestamp']}`")
    lines.append("")
    lines.append("## Decision")
    lines.append("")
    lines.append(f"**{report['decision']}**")
    lines.append("")
    lines.append("## Current champion")
    lines.append("")
    lines.append(f"`{report['context']['champion_adapter']}`")
    lines.append("")
    lines.append("## Main findings")
    lines.append("")
    for item in report["main_findings"]:
        lines.append(f"- {item}")
    lines.append("")
    lines.append("## Dataset audit")
    lines.append("")
    ds = report["dataset_audit"]
    lines.append(f"- Rows: `{ds['rows']}`")
    lines.append(f"- Category counts: `{ds['category_counts']}`")
    lines.append(f"- Unique instructions: `{ds['unique_instructions']}`")
    lines.append(f"- Unique outputs: `{ds['unique_outputs']}`")
    lines.append(f"- Output reuse ratio: `{ds['output_reuse_ratio']}`")
    lines.append(f"- Diagnosis: `{ds['diagnosis']}`")
    lines.append("")
    lines.append("## Training script risks")
    lines.append("")
    for risk in report["training_script_audit"].get("risks", []):
        lines.append(f"- {risk}")
    lines.append("")
    lines.append("## Operational smoke")
    lines.append("")
    sm = report["context"]["operational_smoke"]
    lines.append(f"- Smoke pass: `{sm.get('smoke_pass')}`")
    lines.append(f"- Smoke average: `{sm.get('smoke_average')}`")
    lines.append(f"- Category avg: `{sm.get('category_avg')}`")
    lines.append(f"- Decision: `{sm.get('decision')}`")
    lines.append("")
    lines.append("## Recommendation")
    lines.append("")
    for rec in report["recommendations"]:
        lines.append(f"- {rec}")
    lines.append("")
    return "\n".join(lines)

def main():
    rows, bad_json = load_jsonl(DATASET)
    ds_audit = audit_dataset(rows, bad_json)
    script_audit = audit_training_script(TRAIN_SCRIPT)
    context = summarize_reports()

    main_findings = []

    if ds_audit["repeated_output_problem"]:
        main_findings.append(
            f"Repair v3 has {ds_audit['rows']} rows but only {ds_audit['unique_outputs']} unique outputs. This is template collapse / answer repetition risk."
        )

    if script_audit.get("labels_all_tokens") and not script_audit.get("assistant_only_loss"):
        main_findings.append(
            "Training script appears to train labels on the full chat sequence, not assistant-only tokens. This can encourage prompt/role echo and worsen instruction following."
        )

    if script_audit.get("no_replay_detected"):
        main_findings.append(
            "Repair training used no replay set. Small targeted data can overwrite the champion's broader behavior."
        )

    op_smoke = context["operational_smoke"]
    if op_smoke.get("decision") == "REJECT_OR_REPAIR_NEEDED":
        main_findings.append(
            "Operational smoke confirms regression after repair training."
        )

    failed10k_avg = context["failed_10k"]["average_score"]
    baseline_avg = context["baseline"]["average_score"]
    if baseline_avg and failed10k_avg:
        main_findings.append(
            f"Previous 10K expansion also regressed from {baseline_avg} to {failed10k_avg}; repeated failures suggest data/training design issue, not lack of steps."
        )

    recommendations = [
        "Do not continue training from operational_repair_v3 adapter.",
        "Keep current 14B canary adapter as champion.",
        "Do not train again on datasets with many prompts sharing the same output template.",
        "Before any new repair training, build a dataset with diverse assistant answers and add a replay set from champion-safe categories.",
        "Modify future training scripts to use assistant-only loss masking instead of labels=input_ids for the full chat sequence.",
        "Use router/template for deterministic lora/systemd/dataset operational questions until a safer training recipe is validated.",
    ]

    decision = "DATA_PROBLEM_CONFIRMED_TRAINING_PAUSED"

    report = {
        "timestamp": time.strftime("%Y%m%d_%H%M%S"),
        "decision": decision,
        "dataset_audit": ds_audit,
        "training_script_audit": script_audit,
        "context": context,
        "main_findings": main_findings,
        "recommendations": recommendations,
    }

    ts = report["timestamp"]
    json_path = REPORT_DIR / f"14b_repair_data_problem_audit_{ts}.json"
    md_path = REPORT_DIR / f"14b_repair_data_problem_audit_{ts}.md"

    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    md_path.write_text(make_markdown(report), encoding="utf-8")

    print("==== 14B REPAIR DATA PROBLEM AUDIT SUMMARY ====")
    print("dataset:", DATASET)
    print("rows:", ds_audit["rows"])
    print("category_counts:", ds_audit["category_counts"])
    print("unique_outputs:", ds_audit["unique_outputs"])
    print("output_reuse_ratio:", ds_audit["output_reuse_ratio"])
    print("dataset_diagnosis:", ds_audit["diagnosis"])
    print("labels_all_tokens:", script_audit.get("labels_all_tokens"))
    print("assistant_only_loss:", script_audit.get("assistant_only_loss"))
    print("no_replay_detected:", script_audit.get("no_replay_detected"))
    print("operational_smoke_average:", op_smoke.get("smoke_average"))
    print("operational_smoke_decision:", op_smoke.get("decision"))
    print("champion_adapter:", context["champion_adapter"])
    print("decision:", decision)
    print("json_report:", json_path)
    print("md_report:", md_path)

    print("\n==== MAIN FINDINGS ====")
    for item in main_findings:
        print("-", item)

    print("\n==== RECOMMENDATIONS ====")
    for item in recommendations:
        print("-", item)

if __name__ == "__main__":
    main()