import os
import subprocess
import sys
import threading
from pathlib import Path

import gradio as gr

LOG_FILE = Path("/tmp/apochat_training.log")
STATUS_FILE = Path("/tmp/apochat_training_status.txt")

DATASET = os.environ.get("DATASET", "apoapps/apochat-gemma4-e2b-chat-v1")
OUTPUT_REPO = os.environ.get("OUTPUT_REPO", "apoapps/apochat-gemma4-e2b-apochat-tuned-v2")


def _status() -> str:
    if STATUS_FILE.exists():
        return STATUS_FILE.read_text().strip()
    return "idle"


def _log_tail(n: int = 80) -> str:
    if not LOG_FILE.exists():
        return "No logs yet."
    lines = LOG_FILE.read_text().splitlines()
    return "\n".join(lines[-n:])


def run_training(use_qlora: bool, epochs: float, learning_rate: float) -> None:
    STATUS_FILE.write_text("running")
    LOG_FILE.write_text("")
    cmd = [
        sys.executable,
        "finetune_apochat_peft.py",
        "--dataset", DATASET,
        "--output-dir", "/tmp/apochat-peft-output",
        "--push-to-hub", OUTPUT_REPO,
        "--epochs", str(epochs),
        "--learning-rate", str(learning_rate),
    ]
    if use_qlora:
        cmd.append("--use-qlora")
    with open(LOG_FILE, "a") as log_f:
        log_f.write(f"Running: {' '.join(cmd)}\n")
        proc = subprocess.Popen(
            cmd,
            stdout=log_f,
            stderr=subprocess.STDOUT,
            cwd=str(Path(__file__).parent),
        )
        proc.wait()
    STATUS_FILE.write_text("done" if proc.returncode == 0 else f"failed:{proc.returncode}")


def start_training(use_qlora: bool, epochs: float, learning_rate: float):
    if _status() == "running":
        return "already running", _log_tail()
    threading.Thread(target=run_training, args=(use_qlora, epochs, learning_rate), daemon=True).start()
    return "started", _log_tail()


def refresh() -> tuple[str, str]:
    return _status(), _log_tail()


with gr.Blocks(title="Apochat Gemma 4 E2B Trainer") as demo:
    gr.Markdown("""
    # Apochat Gemma 4 E2B Trainer

    Fine-tune the base `google/gemma-4-E2B-it` model on the Apochat chat dataset using
    PEFT/LoRA (optionally QLoRA). Training is started manually after you upgrade the
    Space hardware to a GPU.

    **Before clicking Start:** go to the Space settings and set Hardware to a GPU
    (e.g. `t4-small`, `a10g-small`, `l4x1`). Free `cpu-basic` will not finish training.
    """)
    with gr.Row():
        use_qlora = gr.Checkbox(label="Use QLoRA (4-bit, saves VRAM)", value=True)
        epochs = gr.Number(label="Epochs", value=1.0, minimum=0.1, maximum=5.0)
        learning_rate = gr.Number(label="Learning rate", value=2e-4, minimum=1e-5, maximum=1e-3)
    start_btn = gr.Button("Start training", variant="primary")
    status_box = gr.Textbox(label="Status", value=_status())
    log_box = gr.Textbox(label="Log", lines=25, value=_log_tail(), max_lines=25)
    start_btn.click(start_training, inputs=[use_qlora, epochs, learning_rate], outputs=[status_box, log_box])
    refresh_btn = gr.Button("Refresh")
    refresh_btn.click(refresh, outputs=[status_box, log_box])
    demo.load(refresh, outputs=[status_box, log_box], every=10)

if __name__ == "__main__":
    demo.launch()