# app.py - HF Space launcher (Option B, HF tree aware, clean logs + Gradio UI)
import os
import subprocess
import threading
import shlex
import time
import requests
import gradio as gr
import socket

# -----------------------------
# Configuration
# -----------------------------
DOWNLOADS = {
    "lora": {
        "url": "https://huggingface.co/latent-consistency/lcm-lora-sdv1-5/resolve/main/pytorch_lora_weights.safetensors",
        "dest": "stable-diffusion-webui/models/Lora/pytorch_lora_weights.safetensors"
    },
    "controlnet_tile": {
        "url": "https://huggingface.co/lllyasviel/ControlNet-v1-1/resolve/main/control_v11f1e_sd15_tile.pth",
        "dest": "stable-diffusion-webui/extensions/ControlNet/models/control_v11f1e_sd15_tile.pth"
    },
    "temporalnet": {
        "url": "https://huggingface.co/CiaraRowles/TemporalNet/resolve/main/diff_control_sd15_temporalnet_fp16.safetensors",
        "dest": "stable-diffusion-webui/extensions/ControlNet/models/diff_control_sd15_temporalnet_fp16.safetensors"
    },
    "civitai_model": {
        "url": "https://civitai.com/api/download/models/143906?type=Model&format=SafeTensor&size=pruned&fp=fp16",
        "dest": "stable-diffusion-webui/models/Stable-diffusion/civitai_model.safetensors"
    }
}

WEBUI_ARGS = "--listen --xformers --enable-insecure-extension-access --medvram"
LAUNCH_PY = "launch.py"  # root of repo

# -----------------------------
# Utilities
# -----------------------------
def ensure_folders():
    """Ensure HF tree persistent folders exist."""
    folders = [
        "stable-diffusion-webui/deforum/input",
        "stable-diffusion-webui/deforum/output_committed/frames",
        "stable-diffusion-webui/deforum/output_committed/video",
        "stable-diffusion-webui/models/Stable-diffusion",
        "stable-diffusion-webui/models/Lora",
        "stable-diffusion-webui/extensions/ControlNet/models",
    ]
    for f in folders:
        os.makedirs(f, exist_ok=True)
    print("✅ Persistent folders ensured.")

def download_file(url, dest, retries=3, backoff=5):
    """Download a file if missing, with retries."""
    if os.path.exists(dest):
        return True
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    for attempt in range(1, retries + 1):
        try:
            with requests.get(url, stream=True, timeout=60) as r:
                r.raise_for_status()
                with open(dest + ".part", "wb") as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
            os.replace(dest + ".part", dest)
            return True
        except Exception:
            time.sleep(backoff * attempt)
    return False

def fetch_models():
    """Download runtime models if missing."""
    for key, info in DOWNLOADS.items():
        download_file(info["url"], info["dest"])

# -----------------------------
# Wait for WebUI port
# -----------------------------
def wait_for_port(port, host="127.0.0.1", timeout=180):
    start = time.time()
    while time.time() - start < timeout:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.settimeout(1)
                s.connect((host, port))
                return True
            except:
                time.sleep(1)
    return False

# -----------------------------
# Start WebUI in background
# -----------------------------
def start_webui():
    ensure_folders()
    fetch_models()
    port = int(os.environ.get("PORT", 7860))
    cmd = ["python", LAUNCH_PY] + shlex.split(WEBUI_ARGS) + [f"--port={port}"]

    # Redirect output to avoid HF log spam
    with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as proc:
        if wait_for_port(port, timeout=180):
            print(f"✅ WebUI is ready on port {port}")
        else:
            print("⚠️ WebUI did not start within timeout.")

threading.Thread(target=start_webui, daemon=True).start()
print("🧵 WebUI background thread started.")

# -----------------------------
# Gradio UI
# -----------------------------
def show_status():
    lines = ["✅ HF Space running (Option B launcher)."]

    # Show public shareable URL for the Space WebUI
    try:
        # share=True generates a public URL automatically
        port = int(os.environ.get("PORT", 7860))
        url = demo.share_url  # updated after launch
        lines.append(f"WebUI URL (open in browser when ready): {url}")
    except Exception:
        lines.append("WebUI URL not yet available. Refresh after a few seconds.")

    # Model status
    for key, info in DOWNLOADS.items():
        present = "yes" if os.path.exists(info["dest"]) else "no"
        lines.append(f"{key}: {present}")
    return "\n".join(lines)

with gr.Blocks() as demo:
    gr.Markdown("## Automatic1111 WebUI — HF Space launcher (Option B)")
    gr.Markdown("This Space launches Automatic1111 in a background thread and keeps a small Gradio UI alive.")
    status_btn = gr.Button("Check status")
    status_out = gr.Textbox(lines=10)
    status_btn.click(fn=show_status, inputs=None, outputs=status_out)

# Launch Gradio with share=True to get a public clickable URL
demo.launch(
    server_name="0.0.0.0",
    server_port=int(os.environ.get("PORT", 7860)),
    ssr_mode=False,
    share=True  # <--- makes the WebUI URL publicly accessible
)