import json
import hashlib
from collections import Counter

path = "18_overnight_pipeline/final_train_candidates_20000.jsonl"
bad_path = "19_validation_bad_lines.txt"

rows = 0
bad_json = 0
empty_lines = 0
empty_required = 0

keys = Counter()
categories = Counter()
sources = Counter()
exact_hashes = Counter()
bad_records = []

def stable_hash(obj):
    s = json.dumps(obj, ensure_ascii=False, sort_keys=True)
    return hashlib.sha256(s.encode("utf-8")).hexdigest()

with open(path, "r", encoding="utf-8") as f:
    for line_no, line in enumerate(f, 1):
        raw = line
        line = line.strip()

        if not line:
            empty_lines += 1
            bad_records.append((line_no, "EMPTY_LINE", raw))
            continue

        try:
            obj = json.loads(line)
        except Exception as e:
            bad_json += 1
            bad_records.append((line_no, f"BAD_JSON: {e}", raw))
            continue

        rows += 1
        keys.update(obj.keys())
        exact_hashes[stable_hash(obj)] += 1

        category = obj.get("category")
        source = obj.get("source")

        if category:
            categories[category] += 1
        if source:
            sources[source] += 1

        has_instruction_output = "instruction" in obj and "output" in obj
        has_messages = "messages" in obj

        if has_instruction_output:
            if not str(obj.get("instruction", "")).strip():
                empty_required += 1
                bad_records.append((line_no, "EMPTY_INSTRUCTION", raw))
            if not str(obj.get("output", "")).strip():
                empty_required += 1
                bad_records.append((line_no, "EMPTY_OUTPUT", raw))
        elif has_messages:
            if not isinstance(obj.get("messages"), list) or len(obj.get("messages")) == 0:
                empty_required += 1
                bad_records.append((line_no, "BAD_MESSAGES", raw))
        else:
            empty_required += 1
            bad_records.append((line_no, "UNKNOWN_SCHEMA", raw))

duplicate_exact = sum(v - 1 for v in exact_hashes.values() if v > 1)

with open(bad_path, "w", encoding="utf-8") as out:
    for line_no, reason, raw in bad_records[:200]:
        out.write(f"LINE {line_no} | {reason}\n")
        out.write(raw)
        if not raw.endswith("\n"):
            out.write("\n")
        out.write("\n")

print("============================================================")
print("FINAL DATASET VALIDATION")
print("============================================================")
print("path =", path)
print("rows =", rows)
print("empty_lines =", empty_lines)
print("bad_json =", bad_json)
print("empty_required =", empty_required)
print("duplicate_exact =", duplicate_exact)
print("bad_records_saved =", bad_path)
print()
print("[keys]")
for k, v in keys.most_common():
    print(k, v)
print()
print("[category counts]")
for k, v in categories.most_common():
    print(k, v)
print()
print("[source counts]")
for k, v in sources.most_common():
    print(k, v)
print("============================================================")