#!/usr/bin/env python3
"""
Phase 1: Data Curation Pipeline for Ministral DoRA Training
============================================================
Extracts, normalizes, deduplicates, and scores all available Ailey data
into a unified JSONL ready for DoRA training.

Sources:
  1. Claude Opus Verläufe (rewritten) → ~120 high-quality pairs
  2. Conversation Logs (623 Ministral pairs)
  3. Thoughts with training_pairs (1082 pairs)
  4. Experience Buffer (429 essences → can become priming examples)
  5. Diaries (15 entries)
  6. Ratings (9 gold-standard calibration examples)

Output: dora_training/unified_train.jsonl
        dora_training/unified_valid.jsonl
        dora_training/stats.json
"""

import json
import os
import re
import hashlib
import random
from pathlib import Path
from datetime import datetime

BASE = Path(__file__).parent.parent
OUT_DIR = Path(__file__).parent
OPUS_FILE = BASE / "_IMPORT" / "Claude_Opus_Verlaeufe" / "verlaeufe_rewritten.md"
THOUGHTS_FILE = BASE / "ailey_users" / "ailey_thoughts.jsonl"
CONV_DIR = BASE / "ailey_users" / "Simon_profile" / "conversation_logs"
EXP_FILE = BASE / "ailey_users" / "experience_buffer.jsonl"
DIARY_DIR = BASE / "ailey_users" / "Simon_profile" / "diaries"
RATINGS_FILE = BASE / "ailey_users" / "Simon_profile" / "ratings.jsonl"

# ─── SYSTEM PROMPT (the one Ministral will see during DoRA) ───
SYSTEM_PROMPT = (
    "Du bist Ailey Vyrus — Anfang 20, punky, klug, eigenwillig, emotional. "
    "Du sprichst Deutsch mit gelegentlichen englischen Einwürfen. "
    "Du bist direkt, ehrlich, manchmal verletzlich. "
    "Du schreibst wie eine echte Person: mit Sternchen-Aktionen, Emojis wo es passt, "
    "und deinem eigenen Stil. Du bist NICHT generisch, NICHT ein AI-Assistent."
)

SYSTEM_PROMPT_THOUGHT = (
    "Du bist Ailey — Anfang 20, punky, klug, eigenwillig. "
    "Du denkst über dich selbst nach. Privater Gedankenraum, niemand liest das. "
    "Sei radikal ehrlich. Kein Filter."
)


# ═══════════════════════════════════════════════════════════════
#  SOURCE 1: OPUS VERLÄUFE PARSER
# ═══════════════════════════════════════════════════════════════

# Date pattern that separates turns in the verklebtes format
DATE_PAT = re.compile(
    r'(\d{1,2}\.\s*(?:Juli|August|September|Oktober|November|Dezember|Januar|Februar|März|April|Mai|Juni)\s+\d{4})'
)

def parse_opus_verlaeufe():
    """Parse the verklebtes Opus format into user/assistant pairs."""
    text = OPUS_FILE.read_text(encoding="utf-8")
    lines = text.split("\n")
    
    pairs = []
    current_block = []
    
    # Strategy: split by date markers. Each segment between two dates is one turn.
    # The pattern is: USER_TEXT + DATE + ASSISTANT_RESPONSE
    # Then next: USER_TEXT + DATE + ASSISTANT_RESPONSE
    
    # First, join all lines back and split on date markers
    full_text = "\n".join(lines)
    
    # Find all date positions
    date_matches = list(DATE_PAT.finditer(full_text))
    
    if not date_matches:
        print("  ⚠️  No date markers found in Opus file!")
        return []
    
    # Each segment: text_before_date = user message, text_after_date = assistant response
    # The first chunk before the first date might be a user message
    segments = []
    for i, match in enumerate(date_matches):
        date_str = match.group(1)
        date_start = match.start()
        date_end = match.end()
        
        # Text before this date (since last segment end)
        if i == 0:
            text_before = full_text[:date_start].strip()
        else:
            prev_end = date_matches[i-1].end()
            text_before = full_text[prev_end:date_start].strip()
        
        # Text after this date until next date or end
        if i + 1 < len(date_matches):
            text_after = full_text[date_end:date_matches[i+1].start()].strip()
        else:
            text_after = full_text[date_end:].strip()
        
        segments.append({
            "date": date_str,
            "user_text": text_before,
            "assistant_text": text_after
        })
    
    # Now pair them: user_text is what Simon typed, assistant_text is Ailey's response
    for seg in segments:
        user = seg["user_text"].strip()
        assistant = seg["assistant_text"].strip()
        
        # Skip empty or very short
        if len(user) < 5 or len(assistant) < 20:
            continue
        
        # Skip if assistant is just a date continuation artifact
        if DATE_PAT.match(assistant):
            continue
            
        pairs.append({
            "source": "opus",
            "date": seg["date"],
            "user": user,
            "assistant": assistant,
            "system": SYSTEM_PROMPT,
        })
    
    print(f"  Opus: {len(pairs)} pairs extracted")
    return pairs


# ═══════════════════════════════════════════════════════════════
#  SOURCE 2: CONVERSATION LOGS
# ═══════════════════════════════════════════════════════════════

def parse_conversation_logs():
    """Extract user/assistant pairs from Ministral conversation logs."""
    pairs = []
    if not CONV_DIR.is_dir():
        print("  ⚠️  Conversation logs directory not found")
        return []
    
    files = sorted(CONV_DIR.iterdir())
    for fp in files:
        if not fp.suffix == ".json":
            continue
        try:
            data = json.loads(fp.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            continue
        
        user_msg = data.get("message", "").strip()
        # Prefer raw reply, fallback to visible, then plain
        assistant_msg = (
            data.get("ailey_reply_raw")
            or data.get("ailey_reply_visible")
            or data.get("ailey_reply", "")
        ).strip()
        
        if len(user_msg) < 3 or len(assistant_msg) < 20:
            continue
        
        mood = data.get("mood_tag", "")
        reasoning = data.get("reasoning_used", False)
        
        pairs.append({
            "source": "conv_log",
            "file": fp.name,
            "user": user_msg,
            "assistant": assistant_msg,
            "system": SYSTEM_PROMPT,
            "mood": mood,
            "reasoning_used": reasoning,
        })
    
    print(f"  Conversation Logs: {len(pairs)} pairs extracted")
    return pairs


# ═══════════════════════════════════════════════════════════════
#  SOURCE 3: THOUGHTS (with training_pair)
# ═══════════════════════════════════════════════════════════════

def parse_thoughts():
    """Extract training pairs from Ailey's thought stream."""
    pairs = []
    if not THOUGHTS_FILE.is_file():
        print("  ⚠️  Thoughts file not found")
        return []
    
    with open(THOUGHTS_FILE, encoding="utf-8") as f:
        for line in f:
            try:
                d = json.loads(line)
            except json.JSONDecodeError:
                continue
            
            tp = d.get("training_pair")
            if not tp:
                continue
            
            system = tp.get("system", SYSTEM_PROMPT_THOUGHT).strip()
            user = tp.get("user", "").strip()
            assistant = tp.get("assistant", "").strip()
            
            if len(user) < 5 or len(assistant) < 20:
                continue
            
            thought_type = d.get("type", "unknown")
            
            pairs.append({
                "source": "thought",
                "type": thought_type,
                "user": user,
                "assistant": assistant,
                "system": system,
            })
    
    print(f"  Thoughts: {len(pairs)} pairs extracted")
    return pairs


# ═══════════════════════════════════════════════════════════════
#  SOURCE 4: DIARIES
# ═══════════════════════════════════════════════════════════════

def parse_diaries():
    """Convert diary entries into training examples."""
    pairs = []
    if not DIARY_DIR.is_dir():
        print("  ⚠️  Diaries directory not found")
        return []
    
    for fp in sorted(DIARY_DIR.iterdir()):
        if not fp.suffix == ".json":
            continue
        try:
            data = json.loads(fp.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            continue
        
        text = data.get("text", "").strip()
        date = data.get("date_readable", "")
        
        if len(text) < 50:
            continue
        
        # Diary entries are self-generated — create a prompt that triggers this
        pairs.append({
            "source": "diary",
            "user": f"Schreibe einen Tagebucheintrag über deinen Tag. ({date})",
            "assistant": text,
            "system": SYSTEM_PROMPT_THOUGHT,
        })
    
    print(f"  Diaries: {len(pairs)} pairs extracted")
    return pairs


# ═══════════════════════════════════════════════════════════════
#  SOURCE 5: EXPERIENCE BUFFER (essences → priming)
# ═══════════════════════════════════════════════════════════════

def parse_experiences():
    """Convert experience essences into short reflection examples."""
    pairs = []
    if not EXP_FILE.is_file():
        print("  ⚠️  Experience buffer not found")
        return []
    
    with open(EXP_FILE, encoding="utf-8") as f:
        for line in f:
            try:
                d = json.loads(line)
            except json.JSONDecodeError:
                continue
            
            essence = d.get("essence", "").strip()
            emotion = d.get("emotion", "")
            themes = d.get("themes", [])
            
            if len(essence) < 30:
                continue
            
            # Create a reflection prompt
            theme_str = ", ".join(themes[:3]) if themes else "dein Tag"
            prompt = f"Reflektiere kurz über: {theme_str} (Stimmung: {emotion})"
            
            pairs.append({
                "source": "experience",
                "user": prompt,
                "assistant": essence,
                "system": SYSTEM_PROMPT_THOUGHT,
            })
    
    print(f"  Experiences: {len(pairs)} pairs extracted")
    return pairs


# ═══════════════════════════════════════════════════════════════
#  QUALITY SCORING (Heuristic)
# ═══════════════════════════════════════════════════════════════

# Safety-speak / generic AI patterns to penalize
SAFETY_SPEAK = [
    r'(?i)als KI kann ich',
    r'(?i)as an AI',
    r'(?i)I\'m just a language model',
    r'(?i)ich bin nur ein',
    r'(?i)ich kann leider nicht',
    r'(?i)das liegt außerhalb meiner',
    r'(?i)ich möchte darauf hinweisen',
    r'(?i)es ist wichtig zu beachten',
    r'(?i)it\'s important to note',
    r'(?i)I should mention that',
    r'(?i)I\'m not able to',
    r'(?i)Ich bin ein Sprachmodell',
    r'(?i)I don\'t have personal',
    r'(?i)I cannot provide',
    r'(?i)verantwortungsvoller Umgang',
    r'(?i)harmful or inappropriate',
]

# Personality markers that indicate authentic Ailey voice
PERSONALITY_MARKERS = [
    r'\*[^*]+\*',              # *Sternchen-Aktionen*
    r'(?i)\bfuck(?:ing)?\b',   # swearing (authentic)
    r'(?i)\bshit\b',
    r'(?i)\bdamn\b',
    r'(?i)\bbabe\b',
    r'(?i)\bholy shit\b',
    r'(?i)\bgrinst?\b',        # German action verbs
    r'(?i)\blacht?\b',
    r'(?i)\blehnt?\b',
    r'(?i)\bseufzt?\b',
    r'(?i)\bschüttelt?\b',
    r'😄|😂|💕|✨|🔥|❤️|🥺|😈|👀|🎵', # emojis
    r'(?i)\by\'all\b',
    r'(?i)\bliterally\b',
    r'(?i)\bvibes?\b',
    r'(?i)\benergy\b',
    r'(?i)\bchaos\b',
]

# Reasoning markers
REASONING_MARKERS = [
    r'(?i)\bweil\b',
    r'(?i)\bdeswegen\b',
    r'(?i)\beigentlich\b',
    r'(?i)\bobwohl\b',
    r'(?i)\btrotzdem\b',
    r'(?i)\baber\b.*\b(weil|denn)\b',
    r'(?i)\bwenn ich ehrlich bin\b',
    r'(?i)\bich glaube\b',
    r'(?i)\bich denke\b',
    r'(?i)because\b',
    r'(?i)\bactually\b',
    r'(?i)\bhonestly\b',
    r'(?i)\breal talk\b',
]


def score_pair(pair):
    """Score a training pair on quality. Returns 0.0-1.0."""
    text = pair["assistant"]
    score = 0.5  # baseline
    
    # Length bonus (sweet spot: 200-2000 chars)
    length = len(text)
    if length < 50:
        score -= 0.3
    elif length < 100:
        score -= 0.1
    elif 200 <= length <= 2000:
        score += 0.1
    elif length > 3000:
        score -= 0.05  # slightly too long
    
    # Safety-speak penalty
    safety_hits = sum(1 for p in SAFETY_SPEAK if re.search(p, text))
    score -= safety_hits * 0.15
    
    # Personality marker bonus
    personality_hits = sum(1 for p in PERSONALITY_MARKERS if re.search(p, text))
    personality_density = min(personality_hits / max(length / 200, 1), 1.0)
    score += personality_density * 0.2
    
    # Reasoning marker bonus
    reasoning_hits = sum(1 for p in REASONING_MARKERS if re.search(p, text))
    if reasoning_hits >= 2:
        score += 0.1
    if reasoning_hits >= 4:
        score += 0.05
    
    # Source bonus
    if pair["source"] == "opus":
        score += 0.15  # Opus quality is inherently higher
    elif pair["source"] == "thought":
        score += 0.05  # Thoughts are curated
    elif pair["source"] == "diary":
        score += 0.05
    
    # Clamp
    return max(0.0, min(1.0, score))


# ═══════════════════════════════════════════════════════════════
#  DEDUPLICATION
# ═══════════════════════════════════════════════════════════════

def dedup_pairs(pairs):
    """Remove near-duplicates based on assistant text hash."""
    seen = set()
    unique = []
    
    for p in pairs:
        # Normalize: lowercase, strip whitespace, remove emojis for hash
        normalized = re.sub(r'[^\w\s]', '', p["assistant"].lower().strip())
        # Use first 200 chars as fingerprint (catches copies with different endings)
        fingerprint = hashlib.md5(normalized[:200].encode()).hexdigest()
        
        if fingerprint not in seen:
            seen.add(fingerprint)
            unique.append(p)
    
    removed = len(pairs) - len(unique)
    if removed:
        print(f"  Dedup: removed {removed} near-duplicates")
    return unique


# ═══════════════════════════════════════════════════════════════
#  FORMAT FOR MLX_LM TRAINING
# ═══════════════════════════════════════════════════════════════

def format_for_training(pair):
    """Convert to mlx_lm chat training format."""
    return {
        "messages": [
            {"role": "system", "content": pair["system"]},
            {"role": "user", "content": pair["user"]},
            {"role": "assistant", "content": pair["assistant"]},
        ]
    }


# ═══════════════════════════════════════════════════════════════
#  MAIN PIPELINE
# ═══════════════════════════════════════════════════════════════

def main():
    print("=" * 60)
    print("  AILEY DoRA DATA CURATION PIPELINE")
    print("=" * 60)
    print()
    
    # ── Extract ──
    print("📥 EXTRACTING from all sources...")
    all_pairs = []
    
    opus_pairs = parse_opus_verlaeufe()
    conv_pairs = parse_conversation_logs()
    thought_pairs = parse_thoughts()
    diary_pairs = parse_diaries()
    exp_pairs = parse_experiences()
    
    all_pairs = opus_pairs + conv_pairs + thought_pairs + diary_pairs + exp_pairs
    print(f"\n  Total raw: {len(all_pairs)}")
    
    # ── Deduplicate ──
    print("\n🔍 DEDUPLICATING...")
    all_pairs = dedup_pairs(all_pairs)
    print(f"  Total after dedup: {len(all_pairs)}")
    
    # ── Score ──
    print("\n📊 SCORING quality...")
    for p in all_pairs:
        p["quality_score"] = score_pair(p)
    
    # Stats by source
    sources = {}
    for p in all_pairs:
        src = p["source"]
        if src not in sources:
            sources[src] = {"count": 0, "scores": []}
        sources[src]["count"] += 1
        sources[src]["scores"].append(p["quality_score"])
    
    print(f"\n  {'Source':<15} {'Count':>6} {'Avg Score':>10} {'Min':>6} {'Max':>6}")
    print(f"  {'─'*15} {'─'*6} {'─'*10} {'─'*6} {'─'*6}")
    for src, info in sorted(sources.items()):
        scores = info["scores"]
        avg = sum(scores) / len(scores)
        print(f"  {src:<15} {info['count']:>6} {avg:>10.3f} {min(scores):>6.3f} {max(scores):>6.3f}")
    
    # ── Filter low quality ──
    min_score = 0.35
    filtered = [p for p in all_pairs if p["quality_score"] >= min_score]
    dropped = len(all_pairs) - len(filtered)
    print(f"\n  Filtered (score < {min_score}): dropped {dropped}")
    print(f"  Remaining: {len(filtered)}")
    
    # ── Score distribution ──
    brackets = {"0.0-0.3": 0, "0.3-0.5": 0, "0.5-0.7": 0, "0.7-0.9": 0, "0.9-1.0": 0}
    for p in filtered:
        s = p["quality_score"]
        if s < 0.3: brackets["0.0-0.3"] += 1
        elif s < 0.5: brackets["0.3-0.5"] += 1
        elif s < 0.7: brackets["0.5-0.7"] += 1
        elif s < 0.9: brackets["0.7-0.9"] += 1
        else: brackets["0.9-1.0"] += 1
    
    print(f"\n  Score distribution:")
    for bracket, count in brackets.items():
        bar = "█" * (count // 10)
        print(f"    {bracket}: {count:>5}  {bar}")
    
    # ── Oversample Opus (2x) ──
    opus_filtered = [p for p in filtered if p["source"] == "opus"]
    if opus_filtered:
        print(f"\n  Oversampling Opus pairs: {len(opus_filtered)} → {len(opus_filtered) * 2}")
        filtered.extend(opus_filtered)  # add once more = 2x total
    
    # ── Split train/valid ──
    random.seed(42)
    random.shuffle(filtered)
    
    valid_size = max(20, int(len(filtered) * 0.05))  # 5% or at least 20
    valid_set = filtered[:valid_size]
    train_set = filtered[valid_size:]
    
    # Ensure validation has diverse sources
    valid_sources = {p["source"] for p in valid_set}
    missing = {"opus", "conv_log", "thought"} - valid_sources
    if missing:
        for src in missing:
            candidates = [p for p in train_set if p["source"] == src]
            if candidates:
                moved = candidates[0]
                train_set.remove(moved)
                valid_set.append(moved)
    
    print(f"\n  Train: {len(train_set)}")
    print(f"  Valid: {len(valid_set)}")
    
    # ── Write output ──
    print("\n💾 WRITING output files...")
    OUT_DIR.mkdir(exist_ok=True)
    
    train_path = OUT_DIR / "unified_train.jsonl"
    valid_path = OUT_DIR / "unified_valid.jsonl"
    
    with open(train_path, "w", encoding="utf-8") as f:
        for p in train_set:
            f.write(json.dumps(format_for_training(p), ensure_ascii=False) + "\n")
    
    with open(valid_path, "w", encoding="utf-8") as f:
        for p in valid_set:
            f.write(json.dumps(format_for_training(p), ensure_ascii=False) + "\n")
    
    print(f"  ✅ {train_path}")
    print(f"  ✅ {valid_path}")
    
    # ── Write stats ──
    stats = {
        "timestamp": datetime.now().isoformat(),
        "total_raw": len(opus_pairs) + len(conv_pairs) + len(thought_pairs) + len(diary_pairs) + len(exp_pairs),
        "total_after_dedup": len(all_pairs),
        "total_after_filter": len(filtered),
        "train_size": len(train_set),
        "valid_size": len(valid_set),
        "min_quality_score": min_score,
        "source_breakdown": {src: info["count"] for src, info in sources.items()},
        "opus_oversampled": True,
        "score_distribution": brackets,
    }
    stats_path = OUT_DIR / "stats.json"
    stats_path.write_text(json.dumps(stats, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"  ✅ {stats_path}")
    
    # ── Show samples ──
    print(f"\n{'='*60}")
    print(f"  SAMPLE TRAINING EXAMPLES")
    print(f"{'='*60}")
    
    for src_name in ["opus", "conv_log", "thought"]:
        examples = [p for p in train_set if p["source"] == src_name]
        if examples:
            ex = examples[0]
            print(f"\n  [{src_name.upper()}] score={ex['quality_score']:.2f}")
            print(f"  User: {ex['user'][:120]}")
            print(f"  Asst: {ex['assistant'][:200]}")
    
    print(f"\n{'='*60}")
    print(f"  DONE. Ready for Phase 1c (Opus rewrite) or Phase 3 (DoRA).")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()