import os
import time
import zipfile
import gradio as gr
import subprocess
import threading

# Paths
OUTPUT_DIR = "train_output"
ZIP_FILE = "python_ai_trained_model.zip"
LOG_FILE = "train_log.txt"

# Zip function
def zip_trained_model():
    with zipfile.ZipFile(ZIP_FILE, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(OUTPUT_DIR):
            for file in files:
                filepath = os.path.join(root, file)
                arcname = os.path.relpath(filepath, OUTPUT_DIR)
                zipf.write(filepath, arcname)

# Tail logs function
def tail_logs(n=20):
    if not os.path.exists(LOG_FILE):
        return ""
    with open(LOG_FILE, 'r') as f:
        return ''.join(f.readlines()[-n:])

# Background training runner
def run_training(status_box, time_box, download_file, model_size_box, log_box):
    start_time = time.time()
    status_box.value = "🚀 Training started..."
    time_box.value = "Training in progress..."
    log_box.value = ""

    # Create log file
    with open(LOG_FILE, "w") as log:
        log.write("🚀 Launching train.py...\n")

    # Start training
    with open(LOG_FILE, "a") as log:
        process = subprocess.Popen(["python", "train.py"], stdout=log, stderr=subprocess.STDOUT)
        while process.poll() is None:
            log_box.update(value=tail_logs())
            time.sleep(5)

    # Check exit status
    if process.returncode != 0:
        status_box.value = f"❌ Training failed with exit code {process.returncode}"
        log_box.value = tail_logs()
        return

    # Training success
    elapsed = round(time.time() - start_time, 2)
    time_box.value = f"✅ Completed in {elapsed // 60:.0f} min {elapsed % 60:.0f} sec"
    status_box.value = "🔄 Compressing model..."

    # Zip it
    zip_trained_model()
    size_mb = round(os.path.getsize(ZIP_FILE) / (1024 * 1024), 2)
    model_size_box.value = f"📦 Model Size: {size_mb} MB"

    download_file.value = ZIP_FILE
    download_file.visible = True
    status_box.value = "✅ Training complete. Download below."

# Button trigger
def start_training(status_box, time_box, download_file, model_size_box, log_box):
    thread = threading.Thread(target=run_training, args=(status_box, time_box, download_file, model_size_box, log_box))
    thread.start()
    return "Training process started."

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Python AI Trainer (StarCoder 7B)")
    gr.Markdown("Train your Python AI with 1 click. Watch logs. Download model when done.")

    with gr.Row():
        train_btn = gr.Button("🚀 Start Training")
        status_box = gr.Textbox(label="Status", value="Ready", interactive=False)

    with gr.Row():
        time_box = gr.Textbox(label="Training Time", interactive=False)
        model_size_box = gr.Textbox(label="Final Model Size", interactive=False)

    log_box = gr.Textbox(label="Live Training Logs", lines=20, interactive=False, value="")
    download_file = gr.File(label="📥 Download Trained Model (.zip)", visible=False)

    train_btn.click(
        fn=start_training,
        inputs=[status_box, time_box, download_file, model_size_box, log_box],
        outputs=[status_box]
    )

    
demo.launch()