import os
import gradio as gr
import subprocess
import zipfile

model_dir = "trained_model"
zip_path = "trained_model.zip"

def model_exists():
    return os.path.exists(zip_path)

def train_model():
    result = subprocess.run(["python", "train.py"], capture_output=True, text=True)
    if result.returncode == 0 and os.path.exists(zip_path):
        return "✅ Model trained successfully! Ready for download.", zip_path
    else:
        return f"❌ Training failed:\n\n{result.stderr}", None

with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Python AI Model Trainer")

    if model_exists():
        gr.Markdown("✅ Trained model found. Click below to download:")
        gr.File(value=zip_path, label="Download Trained Model")
    else:
        gr.Markdown("🚫 No trained model found yet.")
        output = gr.Textbox(label="Training Log")
        download = gr.File(visible=False)
        train_button = gr.Button("🚀 Train Model")

        def on_click_train():
            message, path = train_model()
            return message, gr.update(value=path, visible=True) if path else gr.update(visible=False)

        train_button.click(fn=on_click_train, outputs=[output, download])

demo.launch()