#!/usr/bin/env python3
"""
Prepare Cyron summary dataset for LoRA training.

Reads combined_20k.jsonl and formats it for Hugging Face TRL training.
"""

import json
import argparse
from pathlib import Path


def prepare_dataset(input_file, output_file, test_size=0.05):
    """Prepare dataset for training."""
    
    with open(input_file) as f:
        examples = [json.loads(line) for line in f]
    
    print(f"Loaded {len(examples)} examples")
    
    # Format for training
    formatted = []
    for ex in examples:
        # Create instruction-input-output format
        instruction = f"Task: {ex['task']}"
        if ex['files_changed']:
            instruction += f"\nFiles: {', '.join(ex['files_changed'][:3])}"
        
        input_text = ""
        if ex['tests_run']:
            input_text += f"Tests run: {ex['test_count']}"
        if ex['commit']:
            input_text += f"\nCommit: {ex['git']['commit'] if ex.get('git') else 'yes'}"
        if ex['push']:
            input_text += "\nPush: yes"
        
        output_text = ex['output']
        
        # Create conversation format
        conversation = {
            "conversations": [
                {
                    "from": "human",
                    "value": f"Generate a Cyron summary for this task:\n\n{instruction}\n\n{input_text}"
                },
                {
                    "from": "gpt",
                    "value": output_text
                }
            ]
        }
        
        formatted.append(conversation)
    
    # Split into train/test
    import random
    random.seed(42)
    random.shuffle(formatted)
    
    split_point = int(len(formatted) * (1 - test_size))
    train_data = formatted[:split_point]
    test_data = formatted[split_point:]
    
    # Save
    with open(output_dir / "train.jsonl", "w") as f:
        for item in train_data:
            f.write(json.dumps(item) + "\n")
    
    with open(output_dir / "test.jsonl", "w") as f:
        for item in test_data:
            f.write(json.dumps(item) + "\n")
    
    print(f"Train: {len(train_data)} examples")
    print(f"Test: {len(test_data)} examples")
    print(f"Saved to {output_file.parent}")


def main():
    parser = argparse.ArgumentParser(description="Prepare LoRA training dataset")
    parser.add_argument("--input", type=str, default="../combined_20k.jsonl", 
                       help="Input combined dataset")
    parser.add_argument("--output", type=str, default="data", 
                       help="Output directory")
    parser.add_argument("--test-size", type=float, default=0.05, 
                       help="Test set percentage")
    args = parser.parse_args()
    
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    prepare_dataset(args.input, output_dir, args.test_size)


if __name__ == "__main__":
    main()