import subprocess
import sys

# Force-install torchvision (requirements.txt sometimes fails)
subprocess.run([sys.executable, "-m", "pip", "install", "torchvision", "-q"], check=False)

import gradio as gr
import threading
import os

LOG_FILE = "/tmp/training.log"
training_status = {"running": False, "log": "Not started yet."}

def run_training():
    training_status["running"] = True
    training_status["log"] = "Starting training...\n"
    try:
        proc = subprocess.Popen(
            [sys.executable, "train_canon_ocr.py"],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
        )
        for line in proc.stdout:
            training_status["log"] += line
            if len(training_status["log"]) > 8000:
                training_status["log"] = training_status["log"][-8000:]
        proc.wait()
        training_status["log"] += f"\nProcess exited with code {proc.returncode}"
    except Exception as e:
        training_status["log"] += f"\nERROR: {e}"
    training_status["running"] = False

def start_training():
    if training_status["running"]:
        return "Training already running!"
    t = threading.Thread(target=run_training, daemon=True)
    t.start()
    return "Training started! Check logs below."

def get_logs():
    status = "RUNNING" if training_status["running"] else "STOPPED"
    return f"[{status}]\n\n{training_status["log"]}"

with gr.Blocks() as demo:
    gr.Markdown("# Canon OCR — LoRA Fine-Tuning")
    gr.Markdown("Fine-tunes QARI-OCR on Muharaf + Wellcome MS samples.")
    start_btn = gr.Button("Start Training", variant="primary")
    status_text = gr.Textbox(label="Status", value="Not started")
    log_box = gr.Textbox(label="Training Log", lines=25, value="Click Start to begin.")
    refresh_btn = gr.Button("Refresh Logs")
    start_btn.click(fn=start_training, outputs=status_text)
    refresh_btn.click(fn=get_logs, outputs=log_box)

demo.launch()