import gradio as gr
import spaces
import subprocess
import threading
import os

STATUS = "Idle"


# ── Training ──────────────────────────────────────────────────────────────────

@spaces.GPU
def run_training():
    global STATUS
    STATUS = "Training Running..."

    try:
        subprocess.run(
            ["python", "train.py"],
            check=True
        )
        STATUS = "Training Complete ✅"
    except Exception as e:
        STATUS = f"Training Failed ❌\n{e}"

    return STATUS


def start_training():
    thread = threading.Thread(target=run_training, daemon=False)
    thread.start()
    return "GPU requested. Training will start when GPU is allocated."


# ── Merge ─────────────────────────────────────────────────────────────────────


def run_merge(upload_repo: str):
    global STATUS
    STATUS = "Merging..."

    cmd = ["python", "train.py", "--merge"]
    if upload_repo.strip():
        cmd += ["--upload-merged-to", upload_repo.strip()]

    try:
        subprocess.run(cmd, check=True)
        STATUS = "Merge Complete ✅"
    except Exception as e:
        STATUS = f"Merge Failed ❌\n{e}"

    return STATUS


def start_merge(upload_repo: str):
    thread = threading.Thread(
        target=run_merge, args=(upload_repo,), daemon=False
    )
    thread.start()
    return "GPU requested. Merge will start when GPU is allocated."


# ── Status ────────────────────────────────────────────────────────────────────

def get_status():
    global STATUS
    return STATUS


# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="MamunAI Trainer") as demo:

    gr.Markdown(
        """
        # 🤖 MamunAI Trainer
        **Owner:** Al Mamun Khan  |  **Space:** almamunkhan/MamunAI2

        Space is online.
        Press **Start Training** once — the server handles the rest.
        Training always resumes from the latest adapter in `almamunkhan/MamunAI`.
        Dataset is downloaded from GitHub automatically before every run.
        """
    )

    # ── Training section ──────────────────────────────────────────────────────
    gr.Markdown("---")
    gr.Markdown("## ▶ Training")

    start_btn = gr.Button("🚀 Start Training", variant="primary")

    status_box = gr.Textbox(label="Status", value="Idle")

    refresh_btn = gr.Button("🔄 Refresh Status")

    # ── Merge section ─────────────────────────────────────────────────────────
    gr.Markdown("---")
    gr.Markdown(
        "## 🔀 Merge Model\n"
        "_Merges the saved LoRA adapter into the base model. "
        "Only run this when you want the final merged weights — "
        "training never merges automatically._"
    )

    upload_repo_box = gr.Textbox(
        label="Upload merged model to HF repo (optional)",
        placeholder="e.g. almamunkhan/MamunAI-merged  — leave blank to skip",
    )

    merge_btn = gr.Button("🔀 Merge Model", variant="secondary")

    merge_status_box = gr.Textbox(label="Merge Status", value="Idle")

    merge_refresh_btn = gr.Button("🔄 Refresh Merge Status")

    # ── Event wiring ──────────────────────────────────────────────────────────
    start_btn.click(fn=start_training, outputs=status_box)
    refresh_btn.click(fn=get_status, outputs=status_box)

    merge_btn.click(fn=start_merge, inputs=upload_repo_box, outputs=merge_status_box)
    merge_refresh_btn.click(fn=get_status, outputs=merge_status_box)


if __name__ == "__main__":
    demo.queue()
    demo.launch(server_name="0.0.0.0")