#!/usr/bin/env python

"""
Gradio app for R1-Distill-LLama-8b training interface.
"""

import os
import sys
import json
import gradio as gr
import subprocess
from pathlib import Path
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

def load_config(config_path):
    """Load a JSON configuration file."""
    try:
        with open(config_path, 'r') as f:
            return json.load(f)
    except Exception as e:
        logger.error(f"Error loading config from {config_path}: {e}")
        return None

def start_training():
    """Start the training process."""
    try:
        # Check if training is already in progress
        lock_file = Path("TRAINING_IN_PROGRESS.lock")
        if lock_file.exists():
            return "Training is already in progress. Please wait for it to complete."
        
        # Start training script
        cmd = [sys.executable, "run_transformers_training.py"]
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True
        )
        
        # Return immediate confirmation
        return "Training started successfully! Check the logs for progress."
    except Exception as e:
        logger.error(f"Error starting training: {e}")
        return f"Error starting training: {str(e)}"

def check_training_status():
    """Check the current status of training."""
    try:
        # Check lock file
        lock_file = Path("TRAINING_IN_PROGRESS.lock")
        if not lock_file.exists():
            return "No training in progress"
        
        # Read lock file content
        with open(lock_file, 'r') as f:
            status = f.read()
        return f"Training in progress:\n{status}"
    except Exception as e:
        logger.error(f"Error checking training status: {e}")
        return f"Error checking status: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="R1-Distill-LLama-8b Training") as demo:
    gr.Markdown("# R1-Distill-LLama-8b Training Interface")
    gr.Markdown("This interface allows you to control and monitor the training process.")
    
    with gr.Row():
        start_btn = gr.Button("Start Training")
        status_btn = gr.Button("Check Status")
    
    output = gr.Textbox(label="Output", lines=5)
    
    start_btn.click(start_training, outputs=output)
    status_btn.click(check_training_status, outputs=output)

if __name__ == "__main__":
    demo.launch()