#!/usr/bin/env python3
"""
validate_training_data.py — BUILD-A3 torrkörning
Validerar träningsdatans JSONL-integritet, chat-format, token-distribution
och simulerar load_dataset_from_jsonl() utan unsloth.
"""

import json
import sys
import re
from pathlib import Path
from collections import Counter, defaultdict

TRAIN_PATH = "/opt/amos/data/finetune-exports/aamos_train_v3.jsonl"
VAL_PATH   = "/opt/amos/data/finetune-exports/aamos_val_v3.jsonl"
MAX_SEQ_LENGTH = 2048  # Den osäkra parametern

# ─── Simulera load_dataset_from_jsonl() utan HuggingFace Dataset ──────────────
def simulate_load_dataset(path: str):
    """Simulerar load_dataset_from_jsonl() exakt — returnerar list of {"messages": [...]}"""
    records = []
    errors = []
    with open(path, "r", encoding="utf-8") as f:
        for lineno, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError as e:
                errors.append({"lineno": lineno, "error": f"JSONDecodeError: {e}", "raw": line[:100]})
                continue
            
            # Speglar exakt load_dataset_from_jsonl()-logiken
            if isinstance(obj, list):
                records.append({"lineno": lineno, "messages": obj, "source": "direct_list"})
            elif "messages" in obj:
                records.append({"lineno": lineno, "messages": obj["messages"], "source": "messages_key"})
            else:
                msgs = []
                if "system" in obj:
                    msgs.append({"role": "system", "content": obj["system"]})
                if "user" in obj or "input" in obj:
                    msgs.append({"role": "user", "content": obj.get("user", obj.get("input", ""))})
                if "assistant" in obj or "output" in obj:
                    msgs.append({"role": "assistant", "content": obj.get("assistant", obj.get("output", ""))})
                if msgs:
                    records.append({"lineno": lineno, "messages": msgs, "source": "flat_keys"})
                else:
                    errors.append({"lineno": lineno, "error": "Okänt format — ingen giltig nyckel hittades", "raw": line[:100]})
    return records, errors


def estimate_tokens(text: str) -> int:
    """Grov uppskattning: tecken / 4 (vanlig tumregel för Qwen2/GPT-tokenizers)"""
    return len(text) // 4


def validate_messages(messages, lineno):
    """Validerar ett enskilt examples messages-lista. Returnerar lista av problem."""
    issues = []
    if not isinstance(messages, list):
        issues.append(f"L{lineno}: messages är inte en lista (typ: {type(messages).__name__})")
        return issues
    if len(messages) == 0:
        issues.append(f"L{lineno}: tom messages-lista")
        return issues
    
    roles = [m.get("role") for m in messages if isinstance(m, dict)]
    
    # Kontrollera att alla messages har role + content
    for i, msg in enumerate(messages):
        if not isinstance(msg, dict):
            issues.append(f"L{lineno}: message[{i}] är inte ett dict")
            continue
        if "role" not in msg:
            issues.append(f"L{lineno}: message[{i}] saknar 'role'")
        if "content" not in msg:
            issues.append(f"L{lineno}: message[{i}] saknar 'content'")
        elif msg["content"] == "" or msg["content"] is None:
            issues.append(f"L{lineno}: message[{i}] (role={msg.get('role')}) har tomt content")
    
    return issues


def analyze_file(path: str, label: str):
    print(f"\n{'='*70}")
    print(f"ANALYSERAR: {label}")
    print(f"Fil: {path}")
    print(f"{'='*70}")
    
    # ─── 1. Simulera laddning ──────────────────────────────────────────────────
    records, json_errors = simulate_load_dataset(path)
    
    total_lines = 0
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                total_lines += 1
    
    print(f"\n[1] JSONL-integritet")
    print(f"  Totalt icke-tomma rader: {total_lines}")
    print(f"  Korrekt laddade records: {len(records)}")
    print(f"  JSON-fel (parse-errors): {len(json_errors)}")
    
    if json_errors:
        print("  FELAKTIGA RADER:")
        for e in json_errors[:5]:
            print(f"    L{e['lineno']}: {e['error']} | {e['raw']}")
    
    # Källformat-distribution
    source_counts = Counter(r.get("source", "unknown") for r in records)
    print(f"  Format-distribution: {dict(source_counts)}")
    
    # ─── 2. Chat-format validering ─────────────────────────────────────────────
    print(f"\n[2] Chat-format validering")
    
    message_issues = []
    role_profiles = Counter()
    has_system = 0
    has_user = 0
    has_assistant = 0
    missing_system = 0
    missing_user = 0
    missing_assistant = 0
    empty_content = []
    
    for r in records:
        msgs = r["messages"]
        lineno = r["lineno"]
        
        issues = validate_messages(msgs, lineno)
        if issues:
            message_issues.extend(issues)
        
        roles = tuple(sorted(set(m.get("role","?") for m in msgs if isinstance(m, dict))))
        role_profiles[roles] += 1
        
        roles_list = [m.get("role") for m in msgs if isinstance(m, dict)]
        if "system" in roles_list:
            has_system += 1
        else:
            missing_system += 1
        if "user" in roles_list:
            has_user += 1
        else:
            missing_user += 1
        if "assistant" in roles_list:
            has_assistant += 1
        else:
            missing_assistant += 1
        
        # Kolla tomt content
        for msg in msgs:
            if isinstance(msg, dict):
                c = msg.get("content", "MISSING")
                if c == "" or c is None:
                    empty_content.append((lineno, msg.get("role","?")))
    
    print(f"  Har system-meddelande:    {has_system} / {len(records)}")
    print(f"  Har user-meddelande:      {has_user} / {len(records)}")
    print(f"  Har assistant-meddelande: {has_assistant} / {len(records)}")
    print(f"  Saknar system:     {missing_system}")
    print(f"  Saknar user:       {missing_user}")
    print(f"  Saknar assistant:  {missing_assistant}")
    print(f"  Tomma content:     {len(empty_content)}")
    if empty_content[:5]:
        for lineno, role in empty_content[:5]:
            print(f"    L{lineno}: role={role}")
    
    print(f"  Roll-profiler (topp 5):")
    for combo, cnt in role_profiles.most_common(5):
        print(f"    {combo}: {cnt}")
    
    print(f"  Valideringsfel (meddelanden): {len(message_issues)}")
    if message_issues[:5]:
        for issue in message_issues[:5]:
            print(f"    {issue}")
    
    # ─── 3. Token-längd-distribution ──────────────────────────────────────────
    print(f"\n[3] Token-längd-distribution (uppskattning: tecken/4)")
    
    token_counts = []
    for r in records:
        # Simulera ChatML-format som tokenizern skulle se
        full_text = ""
        for msg in r["messages"]:
            if isinstance(msg, dict):
                role = msg.get("role", "")
                content = msg.get("content", "")
                # ChatML-format för Qwen2.5
                full_text += f"<|im_start|>{role}\n{content}<|im_end|>\n"
        tok = estimate_tokens(full_text)
        token_counts.append(tok)
    
    if token_counts:
        token_counts_sorted = sorted(token_counts)
        n = len(token_counts)
        
        over_2048 = sum(1 for t in token_counts if t > MAX_SEQ_LENGTH)
        over_1024 = sum(1 for t in token_counts if t > 1024)
        over_4096 = sum(1 for t in token_counts if t > 4096)
        
        percentiles = {}
        for p in [50, 75, 90, 95, 99, 100]:
            idx = min(int(n * p / 100), n-1)
            percentiles[p] = token_counts_sorted[idx]
        
        print(f"  Antal exempel: {n}")
        print(f"  Min tokens: {min(token_counts)}")
        print(f"  Max tokens: {max(token_counts)}")
        print(f"  Medelvärde: {sum(token_counts)//n}")
        print(f"  Percentiler:")
        for p, v in percentiles.items():
            print(f"    p{p:3d}: {v}")
        print(f"\n  >>> Överstiger max_seq_length={MAX_SEQ_LENGTH}:")
        print(f"      {over_2048} av {n} ({over_2048/n*100:.1f}%)")
        print(f"  >>> Överstiger 1024: {over_1024} ({over_1024/n*100:.1f}%)")
        print(f"  >>> Överstiger 4096: {over_4096} ({over_4096/n*100:.1f}%)")
        
        # Histogram
        buckets = [0]*10  # 0-200, 200-400, ..., 1800-2000, 2000+
        bucket_labels = ["0-200","200-400","400-600","600-800","800-1000",
                         "1000-1200","1200-1400","1400-1600","1600-1800","1800-2000","2000+"]
        hist = defaultdict(int)
        for t in token_counts:
            if t < 200: hist["0-200"] += 1
            elif t < 400: hist["200-400"] += 1
            elif t < 600: hist["400-600"] += 1
            elif t < 800: hist["600-800"] += 1
            elif t < 1000: hist["800-1000"] += 1
            elif t < 1200: hist["1000-1200"] += 1
            elif t < 1400: hist["1200-1400"] += 1
            elif t < 1600: hist["1400-1600"] += 1
            elif t < 1800: hist["1600-1800"] += 1
            elif t < 2000: hist["1800-2000"] += 1
            else: hist["2000+"] += 1
        
        print(f"\n  Histogram:")
        for label in bucket_labels:
            cnt = hist[label]
            bar = "█" * (cnt * 40 // n)
            print(f"    {label:10s}: {cnt:6d} ({cnt/n*100:5.1f}%) {bar}")
    
    return {
        "total_lines": total_lines,
        "records": len(records),
        "json_errors": len(json_errors),
        "message_issues": len(message_issues),
        "empty_content": len(empty_content),
        "missing_system": missing_system,
        "missing_user": missing_user,
        "missing_assistant": missing_assistant,
        "over_2048": over_2048 if token_counts else 0,
        "over_4096": over_4096 if token_counts else 0,
        "max_tokens": max(token_counts) if token_counts else 0,
        "p95_tokens": percentiles.get(95, 0) if token_counts else 0,
        "p99_tokens": percentiles.get(99, 0) if token_counts else 0,
        "n": n if token_counts else 0,
    }


def canon_check(path: str, label: str):
    """CANON-stickprov: verifiera injicerade CANON-exempel."""
    print(f"\n[5] CANON-stickprov ({label})")
    
    canon_keywords = {
        "infrastrukturkontroll": 0,
        "559141-7042": 0,         # korrekt org-nr
        "eskalera": 0,
    }
    bad_orgnr_pattern = "559316-3881"  # felaktigt org-nr
    bad_orgnr_count = 0
    
    with open(path, "r", encoding="utf-8") as f:
        for lineno, line in enumerate(f, 1):
            if not line.strip():
                continue
            for kw in canon_keywords:
                if kw in line:
                    canon_keywords[kw] += 1
            if bad_orgnr_pattern in line:
                bad_orgnr_count += 1
    
    for kw, cnt in canon_keywords.items():
        status = "✅" if cnt > 0 else "❌ EJ FUNNET"
        print(f"  '{kw}': {cnt} förekomster {status}")
    
    status = "✅ 0 förekomster (korrekt)" if bad_orgnr_count == 0 else f"❌ {bad_orgnr_count} FÖREKOMSTER (FEL!)"
    print(f"  Felaktigt org-nr '{bad_orgnr_pattern}': {bad_orgnr_count} förekomster {status}")
    
    return canon_keywords, bad_orgnr_count


# ─── Kör validering ────────────────────────────────────────────────────────────
print("BUILD-A3 — Torrkörning: Validering av ESLM-träningsdata (v3)")
print("=" * 70)

train_stats = analyze_file(TRAIN_PATH, "TRÄNINGSDATA (aamos_train_v3.jsonl)")
val_stats   = analyze_file(VAL_PATH,   "VALIDERINGSDATA (aamos_val_v3.jsonl)")

# CANON-check på båda filerna
canon_train, bad_train = canon_check(TRAIN_PATH, "träningsdata")
canon_val,   bad_val   = canon_check(VAL_PATH,   "valideringsdata")

# ─── Sammanfattning ────────────────────────────────────────────────────────────
print(f"\n{'='*70}")
print("SAMMANFATTNING & DOM")
print(f"{'='*70}")

problems = []
warnings = []

# JSON-integritet
if train_stats["json_errors"] > 0:
    problems.append(f"TRÄNING: {train_stats['json_errors']} JSON-parse-fel")
if val_stats["json_errors"] > 0:
    problems.append(f"VALIDERING: {val_stats['json_errors']} JSON-parse-fel")

# Format-problem
if train_stats["missing_user"] > 0:
    problems.append(f"TRÄNING: {train_stats['missing_user']} exempel saknar user-meddelande")
if train_stats["missing_assistant"] > 0:
    problems.append(f"TRÄNING: {train_stats['missing_assistant']} exempel saknar assistant-meddelande")
if train_stats["empty_content"] > 0:
    problems.append(f"TRÄNING: {train_stats['empty_content']} tomma content-fält")
if train_stats["missing_system"] > 0:
    warnings.append(f"TRÄNING: {train_stats['missing_system']} exempel saknar system-meddelande (varning, ej blockering)")

# Token-längd
over_pct = train_stats["over_2048"] / train_stats["n"] * 100 if train_stats["n"] > 0 else 0
if train_stats["over_2048"] > 0:
    msg = f"TRÄNING: {train_stats['over_2048']} ({over_pct:.1f}%) exempel överstiger max_seq_length=2048"
    if over_pct > 5:
        problems.append(msg + " — HÖJNING REKOMMENDERAS")
    else:
        warnings.append(msg + " — trunkeras vid träning")

if train_stats["over_4096"] > 0:
    problems.append(f"TRÄNING: {train_stats['over_4096']} exempel överstiger 4096 tokens — kräver ännu längre kontext!")

# CANON
for kw, cnt in canon_train.items():
    if cnt == 0:
        problems.append(f"CANON: nyckelord '{kw}' EJ FUNNET i träningsdata!")
if bad_train > 0:
    problems.append(f"CANON: Felaktigt org-nr '559316-3881' hittades {bad_train} gånger i träningsdata!")
if bad_val > 0:
    problems.append(f"CANON: Felaktigt org-nr '559316-3881' hittades {bad_val} gånger i valideringsdata!")

print(f"\nTräningsdata:    {train_stats['records']} records ({train_stats['total_lines']} rader)")
print(f"Valideringsdata: {val_stats['records']} records ({val_stats['total_lines']} rader)")
print(f"Max token-längd (träning): {train_stats['max_tokens']}")
print(f"p95 token-längd (träning): {train_stats['p95_tokens']}")
print(f"p99 token-längd (träning): {train_stats['p99_tokens']}")
print(f"Överstiger 2048: {train_stats['over_2048']} ({over_pct:.1f}%)")

print(f"\n⚠️  VARNINGAR ({len(warnings)}):")
for w in warnings:
    print(f"  - {w}")

print(f"\n❌  PROBLEM ({len(problems)}):")
for p in problems:
    print(f"  - {p}")

if not problems:
    print("\n🟢 DOM: DATAN ÄR KLAR FÖR TRÄNING")
    if train_stats["over_2048"] > 0 and over_pct <= 5:
        print(f"   OBS: {over_pct:.1f}% av exemplen trunkeras vid max_seq_length=2048")
        print(f"   max_seq_length=2048 bedöms som TILLRÄCKLIGT om p99={train_stats['p99_tokens']} ≤ 2048")
    elif train_stats["over_2048"] == 0:
        print(f"   max_seq_length=2048 ÄR TILLRÄCKLIGT (inget exempel överstiger gränsen)")
else:
    print("\n🔴 DOM: DATAN HAR PROBLEM — SE OVAN")

print(f"\n{'='*70}")
print("max_seq_length-rekommendation:")
print(f"  p95={train_stats['p95_tokens']} tokens")
print(f"  p99={train_stats['p99_tokens']} tokens")
print(f"  max={train_stats['max_tokens']} tokens")
if train_stats['p99_tokens'] <= 2048:
    print(f"  → max_seq_length=2048 täcker 99% av exemplen ✅")
    print(f"  → VRAM-kalkyl från A2 gäller med 2048 som antaget")
elif train_stats['p99_tokens'] <= 4096:
    print(f"  → max_seq_length=4096 rekommenderas för full täckning")
    print(f"  → ⚠️ PÅVERKAR VRAM-kalkyl (A2 måste räkna om!)")
else:
    print(f"  → max_seq_length > 4096 kan behövas — utred vidare")
    print(f"  → ⚠️ PÅVERKAR VRAM-kalkyl KRAFTIGT (A2 måste räkna om!)")