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"

# Shared states
start_time = None
end_time = None

# Function: Zip the trained model
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)

# Function: Tail the logs from training
def tail_logs():
    if not os.path.exists(LOG_FILE):
        return ""
    with open(LOG_FILE, 'r') as f:
        lines = f.readlines()[-20:]  # last 20 lines
        return ''.join(lines)

# Background training function
def run_training(status_box, time_box, download_file, model_size_box, log_box):
    global start_time, end_time

    start_time = time.time()
    status_box.value = "🚀 Training started... Please wait. This can take 30–90+ minutes."
    time_box.value = "Training in progress..."
    log_box.value = ""

    # Run training and redirect stdout/stderr to log file
    with open(LOG_FILE, "w") as log:
        process = subprocess.Popen(["python", "train.py"], stdout=log, stderr=log)
        while process.poll() is None:
            log_box.value = tail_logs()
            time.sleep(5)

    end_time = time.time()
    elapsed = round(end_time - start_time, 2)
    time_box.value = f"✅ Training completed in {elapsed // 60:.0f} min {elapsed % 60:.0f} sec."

    status_box.value = "🔄 Compressing trained model..."
    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 = "✅ Done! You can now download your trained model."

# Trigger button click
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 has started..."

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Python AI Trainer (StarCoder 7B)")
    gr.Markdown("Train a custom Python AI that can write, fix, and explain Python code.")

    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)
    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()