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

output_path = "train_output"
zipped_model = "python_ai_trained_model.zip"
status = gr.Textbox(label="Status", value="Ready", interactive=False)
download_link = gr.File(label="Download Trained Model", visible=False)

def run_training():
    global zipped_model
    status.value = "Training started... this may take a while (15–60+ mins)."

    # Run train.py
    subprocess.run(["python", "train.py"], check=True)

    # Compress trained model
    status.value = "Training complete. Compressing model..."
    with zipfile.ZipFile(zipped_model, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(output_path):
            for file in files:
                filepath = os.path.join(root, file)
                arcname = os.path.relpath(filepath, output_path)
                zipf.write(filepath, arcname)

    # Move zip to Gradio-visible path
    status.value = "Done. Model ready for download."
    download_link.visible = True

def start_training():
    training_thread = threading.Thread(target=run_training)
    training_thread.start()
    return "Training in progress...", None

with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Python AI Model Trainer (StarCoder 7B)")
    gr.Markdown("Click the button below to start fine-tuning your custom Python AI. After training, download the model and use it anywhere.")
    
    with gr.Row():
        train_btn = gr.Button("🚀 Start Training")
        train_btn.click(start_training, outputs=[status, download_link])

    status.render()
    download_link.render()

demo.launch()