import os
import subprocess
import threading
from flask import Flask, render_template, redirect, jsonify

app = Flask(__name__)

IMAGE_STACK = os.environ.get("IMAGE_STACK", "/docker/ai-image-gen")
LLM_STACK   = os.environ.get("LLM_STACK",   "/docker/ai-llm")

# Track current active stack and swap status
state = {
    "active": None,       # "images" | "chat" | None
    "swapping": False,
    "message": "No stack active"
}

def run_compose(stack_dir, command):
    """Run a docker compose command in a stack directory."""
    subprocess.run(
        ["docker", "compose", "-f", f"{stack_dir}/docker-compose.yml"] + command,
        capture_output=True
    )

def is_stack_running(stack_dir):
    """Check if the GPU container for this stack is running."""
    container = "ai-image-gen-invokeai-1" if "image" in stack_dir else "ai-llm-ollama-1"
    result = subprocess.run(
        ["docker", "inspect", "--format", "{{.State.Running}}", container],
        capture_output=True,
        text=True
    )
    return result.stdout.strip() == "true"

def do_swap(target):
    """Stop the GPU-heavy container of inactive stack, start target."""
    state["swapping"] = True

    if target == "images":
        state["message"] = "Stopping Ollama to free VRAM..."
        subprocess.run(["docker", "stop", "ai-llm-ollama-1"], capture_output=True)
        state["message"] = "Starting InvokeAI..."
        subprocess.run(
            ["docker", "compose", "-f", f"{IMAGE_STACK}/docker-compose.yml", "up", "-d"],
            capture_output=True
        )
    else:
        state["message"] = "Stopping InvokeAI to free VRAM..."
        subprocess.run(["docker", "stop", "ai-image-gen-invokeai-1"], capture_output=True)
        state["message"] = "Starting Ollama..."
        subprocess.run(["docker", "start", "ai-llm-ollama-1"], capture_output=True)

    state["active"] = target
    state["swapping"] = False
    state["message"] = f"{target.capitalize()} is ready."

@app.route("/")
@app.route("/portal")
def portal():
    image_running = is_stack_running(IMAGE_STACK)
    llm_running   = is_stack_running(LLM_STACK)
    return render_template(
        "portal.html",
        active=state["active"],
        swapping=state["swapping"],
        message=state["message"],
        image_running=image_running,
        llm_running=llm_running
    )

@app.route("/switch/<target>")
def switch(target):
    if target not in ("images", "chat"):
        return "Invalid target", 400

    if state["swapping"]:
        return redirect("/portal")

    # Warn if same stack requested
    if state["active"] == target:
        urls = {"chat": "https://chat.mydomain.com", "images": "https://images.mydomain.com"}
        return redirect(urls[target])

    # Run swap in background so browser gets immediate response
    thread = threading.Thread(target=do_swap, args=(target,))
    thread.daemon = True
    thread.start()

    return render_template(
        "switching.html",
        target=target,
        message=f"Switching to {target}... this takes about 10 seconds."
    )

@app.route("/chat")
def chat():
    return redirect(URLS["chat"])

@app.route("/images")
def images():
    return redirect(URLS["images"])

@app.route("/status")
def status():
    """Polled by the switching page to know when swap is done."""
    return jsonify(state)

if __name__ == "__main__":
    # Detect what's already running on startup
    if is_stack_running(IMAGE_STACK):
        state["active"] = "images"
        state["message"] = "Images stack is active."
    elif is_stack_running(LLM_STACK):
        state["active"] = "chat"
        state["message"] = "Chat stack is active."

    app.run(host="0.0.0.0", port=8080)