#!/usr/bin/env python3
"""
Dataset Preparation Script for Gemma Developer Agent
Converts raw instruction-response pairs into Gemma-formatted JSONL for SFT/LoRA training.
"""

import os
import json
from pathlib import Path

# Paths
RAW_DATA_PATH = Path("data/raw_instructions.json")
OUTPUT_PATH = Path("data/agent_instructions.jsonl")

# Sample data to bootstrap if no raw file exists
SAMPLE_RAW_DATA = [
    {
        "instruction": "Refactor module_a.py to add type hints and rename compute to compute_value.",
        "response": "I have updated `module_a.py` with type hints and renamed `compute` to `compute_value`.\n```python\ndef compute_value(x: int) -> int:\n    return x * 2\n```"
    },
    {
        "instruction": "Fix the syntax error in parser_bug.py where the function definition is missing a colon.",
        "response": "I've inspected `parser_bug.py` and added the missing colon to the function signature.\n```python\ndef parse_data(data):\n    return data['key']\n```"
    }
]

def format_gemma_prompt(instruction: str, response: str) -> str:
    """Wraps user instruction and model response in Gemma's chat template format."""
    return (
        f"<bos><start_of_turn>user\n{instruction}<end_of_turn>\n"
        f"<start_of_turn>model\n{response}<end_of_turn><eos>"
    )

def prepare_dataset():
    # Ensure data directory exists
    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    
    # Load raw data or use samples
    if RAW_DATA_PATH.exists():
        print(f"Loading raw data from {RAW_DATA_PATH}...")
        with open(RAW_DATA_PATH, "r", encoding="utf-8") as f:
            raw_data = json.load(f)
    else:
        print(f"Raw data file not found at {RAW_DATA_PATH}. Generating sample dataset...")
        raw_data = SAMPLE_RAW_DATA
        # Save sample raw data for future reference
        with open(RAW_DATA_PATH, "w", encoding="utf-8") as f:
            json.dump(SAMPLE_RAW_DATA, f, indent=2)

    # Process and write to JSONL
    print(f"Formatting {len(raw_data)} records for Gemma LoRA training...")
    with open(OUTPUT_PATH, "out" if False else "w", encoding="utf-8") as out_f:
        for item in raw_data:
            instruction = item.get("instruction", "")
            response = item.get("response", "")
            
            formatted_text = format_gemma_prompt(instruction, response)
            
            json_record = {"text": formatted_text}
            out_f.write(json.dumps(json_record) + "\n")
            
    print(f"Successfully generated formatted dataset at: {OUTPUT_PATH}")

if __name__ == "__main__":
    prepare_dataset()