#!/usr/bin/env python3
"""Script to prepare MLX chat format datasets from rank_log_*.jsonl files."""

import json
import random
import pathlib

# Set random seed for reproducibility
random.seed(42)

def prepare_filter_data(log_file="data/rank_log_filter.jsonl"):
    p = pathlib.Path(log_file)
    if not p.exists():
        print(f"File {log_file} does not exist.")
        return []

    lines = p.read_text().strip().split("\n")
    dataset = []
    
    for line in lines:
        if not line:
            continue
        data = json.loads(line)
        
        # Build the user prompt identically to how BAML does or provide sufficient context
        user_msg = f"Title: {data.get('title')}\nSource: {data.get('source')}\nSummary: {data.get('summary')}\n\nShould we KEEP or DROP this item?"
        
        reasoning = data.get("reasoning", "")
        verdict = data.get("verdict", "DROP")
        assistant_msg = f"<think>\n{reasoning}\n</think>\nVerdict: {verdict}"

        dataset.append({
            "messages": [
                {"role": "user", "content": user_msg},
                {"role": "assistant", "content": assistant_msg}
            ]
        })
        
    return dataset

def prepare_score_data(log_file="data/rank_log_score.jsonl"):
    p = pathlib.Path(log_file)
    if not p.exists():
        print(f"File {log_file} does not exist.")
        return []

    lines = p.read_text().strip().split("\n")
    dataset = []
    
    for line in lines:
        if not line:
            continue
        data = json.loads(line)
        
        user_msg = f"Title: {data.get('title')}\nSource: {data.get('source')}\nSummary: {data.get('summary')}\n\nScore this item from 0 to 10 and provide a reason."
        score = data.get("score")
        reason = data.get("reason", "")
        reasoning = data.get("reasoning", "")
        
        assistant_msg = f"<think>\n{reasoning}\n</think>\nScore: {score}\nReason: {reason}"

        dataset.append({
            "messages": [
                {"role": "user", "content": user_msg},
                {"role": "assistant", "content": assistant_msg}
            ]
        })
        
    return dataset

def main():
    filter_data = prepare_filter_data()
    score_data = prepare_score_data()
    
    # For now, we can combine them or train two separate LoRAs.
    # We will combine them into a single distillation dataset.
    all_data = filter_data + score_data
    if not all_data:
        print("No data found. Exiting.")
        return

    random.shuffle(all_data)
    
    # 90% train, 10% valid
    split_idx = int(len(all_data) * 0.9)
    train_data = all_data[:split_idx]
    valid_data = all_data[split_idx:]
    
    out_dir = pathlib.Path("data/mlx_dataset")
    out_dir.mkdir(parents=True, exist_ok=True)
    
    with open(out_dir / "train.jsonl", "w") as f:
        for item in train_data:
            f.write(json.dumps(item) + "\n")
            
    with open(out_dir / "valid.jsonl", "w") as f:
        for item in valid_data:
            f.write(json.dumps(item) + "\n")
            
    print(f"Created MLX dataset with {len(train_data)} train and {len(valid_data)} valid examples.")

if __name__ == "__main__":
    main()