import sys
import json
from transformers import AutoTokenizer
from tqdm import tqdm

# IMPORTANT: The tokenizer MUST match the student model you are fine-tuning.
STUDENT_MODEL_ID = "google/t5gemma-s-s-ul2-it"

def reformat_with_template(input_filepath: str, output_filepath: str):
    """
    Reads a JSONL file with 'text' (prompt) and 'target' (completion) fields
    and reformats it into a single 'text' field using the model's chat template.
    """
    print(f"Loading tokenizer for {STUDENT_MODEL_ID}...")
    tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL_ID)
    
    print(f"Reformatting {input_filepath}...")
    
    with open(input_filepath, 'r', encoding='utf-8') as infile, \
         open(output_filepath, 'w', encoding='utf-8') as outfile:
        
        # Count lines for tqdm progress bar
        num_lines = sum(1 for line in open(input_filepath, 'r', encoding='utf-8'))
        infile.seek(0) # Reset file pointer
        
        for line in tqdm(infile, total=num_lines, desc="Reformatting"):
            original_entry = json.loads(line)
            
            # The original prompt is in the 'text' field
            prompt_text = original_entry.get("text", "")
            # The verified completion is in the 'target' field
            completion_text = original_entry.get("target", "")
            
            # The prompt needs to end with the cue for the model to respond to
            if "Timestamp:" not in prompt_text:
                 prompt_text += "\n\nTimestamp:\n"

            # Create the message structure that the chat template expects
            messages = [
                {"role": "user", "content": prompt_text},
                {"role": "assistant", "content": completion_text}
            ]
            
            # Apply the template to create the new, single text field
            # add_generation_prompt=False is important for training data
            formatted_text = tokenizer.apply_chat_template(
                messages, 
                tokenize=False, 
                add_generation_prompt=False
            )
            
            new_entry = {
                "id": original_entry.get("id"),
                "text": formatted_text,
                "target": completion_text # Keep target for reference if needed
            }
            outfile.write(json.dumps(new_entry) + "\n")

    print(f"\nSuccessfully reformatted {num_lines} entries. Output saved to {output_filepath}.")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python reformat_with_template.py <input_dataset.jsonl> <output_dataset_templated.jsonl>")
        sys.exit(1)
        
    reformat_with_template(sys.argv[1], sys.argv[2])