#!/usr/bin/env python3
"""Strict Docker profile switcher for the local AI stack.

Only status and activation of a fixed set of labelled llama.cpp containers are
exposed. Callers cannot provide images, commands, mounts or Docker API paths.
"""

from __future__ import annotations

import http.client
import json
import logging
import os
import socket
import threading
import urllib.parse
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

HOST = os.environ.get("CONTROLLER_HOST", "0.0.0.0")
PORT = int(os.environ.get("CONTROLLER_PORT", "8090"))
SOCKET_PATH = os.environ.get("DOCKER_SOCKET", "/var/run/docker.sock")
TOKEN_FILE = os.environ.get("CONTROLLER_TOKEN_FILE", "/run/secrets/controller-token")
ALLOWED = tuple(x.strip() for x in os.environ.get(
    "ALLOWED_PROFILES", "fast,medium,large,ultra,uncensored,experimental").split(",") if x.strip())
LABEL_KEY = "com.mike-ai.llama-profile"
IMAGE_LABEL_KEY = "com.mike-ai.image-worker"
IMAGE_WORKER = os.environ.get("IMAGE_WORKER", "image")
RESTORE_WORKER = os.environ.get("RESTORE_WORKER", "restore")
TTS_LABEL_KEY = "com.mike-ai.tts-worker"
TTS_WORKER = os.environ.get("TTS_WORKER", "qwen3")
MUSIC_LABEL_KEY = "com.mike-ai.music-worker"
MUSIC_WORKER = os.environ.get("MUSIC_WORKER", "").strip()
LOCK = threading.Lock()
log = logging.getLogger("profile-controller")


class UnixConnection(http.client.HTTPConnection):
    def connect(self) -> None:
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(SOCKET_PATH)


def docker_request(method: str, path: str) -> tuple[int, bytes]:
    conn = UnixConnection("localhost", timeout=30)
    try:
        conn.request(method, path, headers={"Content-Type": "application/json"})
        response = conn.getresponse()
        return response.status, response.read()
    finally:
        conn.close()


def containers() -> dict[str, dict]:
    filters = urllib.parse.quote(json.dumps({"label": [LABEL_KEY]}))
    status, body = docker_request("GET", f"/containers/json?all=1&filters={filters}")
    if status != 200:
        raise RuntimeError(f"Docker list failed with HTTP {status}")
    result: dict[str, dict] = {}
    for item in json.loads(body):
        profile = item.get("Labels", {}).get(LABEL_KEY)
        if profile in ALLOWED:
            if profile in result:
                raise RuntimeError(f"duplicate container for profile {profile}")
            result[profile] = item
    return result


def labelled_containers(label: str) -> list[dict]:
    filters = urllib.parse.quote(json.dumps({"label": [label]}))
    status, body = docker_request("GET", f"/containers/json?all=1&filters={filters}")
    if status != 200:
        raise RuntimeError(f"Docker list failed with HTTP {status}")
    return json.loads(body)


def image_container(kind: str = IMAGE_WORKER) -> dict:
    matches = [item for item in labelled_containers(IMAGE_LABEL_KEY)
               if item.get("Labels", {}).get(IMAGE_LABEL_KEY) == kind]
    if len(matches) != 1:
        raise RuntimeError(
            f"expected exactly one image worker {kind!r}, found {len(matches)}")
    return matches[0]


def image_containers() -> list[dict]:
    """All allowlisted GPU workers that must never overlap an LLM."""
    allowed = {IMAGE_WORKER, RESTORE_WORKER}
    return [item for item in labelled_containers(IMAGE_LABEL_KEY)
            if item.get("Labels", {}).get(IMAGE_LABEL_KEY) in allowed]


def tts_container() -> dict:
    matches = [item for item in labelled_containers(TTS_LABEL_KEY)
               if item.get("Labels", {}).get(TTS_LABEL_KEY) == TTS_WORKER]
    if len(matches) != 1:
        raise RuntimeError(
            f"expected exactly one TTS worker {TTS_WORKER!r}, found {len(matches)}")
    return matches[0]


def music_container() -> dict:
    if not MUSIC_WORKER:
        raise RuntimeError("music worker is not configured")
    matches = [item for item in labelled_containers(MUSIC_LABEL_KEY)
               if item.get("Labels", {}).get(MUSIC_LABEL_KEY) == MUSIC_WORKER]
    if len(matches) != 1:
        raise RuntimeError(
            f"expected exactly one music worker {MUSIC_WORKER!r}, found {len(matches)}")
    return matches[0]


def stop_music_if_configured() -> None:
    if MUSIC_WORKER:
        stop_container(music_container(), timeout=30)


def stop_container(item: dict, timeout: int = 120) -> None:
    if item.get("State") != "running":
        return
    status, _ = docker_request("POST", f"/containers/{item['Id']}/stop?t={timeout}")
    if status not in (204, 304):
        raise RuntimeError(f"failed to stop container: HTTP {status}")


def start_container(item: dict) -> None:
    if item.get("State") == "running":
        return
    status, _ = docker_request("POST", f"/containers/{item['Id']}/start")
    if status not in (204, 304):
        raise RuntimeError(f"failed to start container: HTTP {status}")


def stop_inference() -> dict:
    with LOCK:
        items = containers()
        previous = active_profile(items)
        for item in items.values():
            stop_container(item)
        return {"active_profile": None, "previous_profile": previous}


def set_image_worker(running: bool, kind: str = IMAGE_WORKER) -> dict:
    if kind not in {IMAGE_WORKER, RESTORE_WORKER}:
        raise ValueError("worker is not allowlisted")
    with LOCK:
        item = image_container(kind)
        if running:
            # The image worker may never overlap a llama profile on the 5080.
            for profile_item in containers().values():
                stop_container(profile_item)
            # The 9B beta text encoder temporarily borrows the RTX 3060 from
            # Qwen3-TTS. The gateway retains Piper as a fallback meanwhile.
            stop_container(tts_container(), timeout=30)
            stop_music_if_configured()
            for other in image_containers():
                if other["Id"] != item["Id"]:
                    stop_container(other, timeout=20)
            start_container(item)
        else:
            # CUDA/PyTorch may not react promptly to SIGTERM after an OOM.
            # Bound recovery time and let Docker issue SIGKILL afterwards.
            stop_container(item, timeout=20)
            # TTS is restored by the following profile activation. Keeping it
            # stopped here lets the router verify that both GPUs really
            # released the image model before Qwen and TTS are reloaded.
        return {"image_worker": kind,
                "state": "running" if running else "stopped"}


def set_music_worker(running: bool) -> dict:
    """Start ACE-Step exclusively, or stop it before LLM restoration."""
    with LOCK:
        item = music_container()
        if running:
            for profile_item in containers().values():
                stop_container(profile_item)
            for worker in image_containers():
                stop_container(worker, timeout=20)
            stop_container(tts_container(), timeout=30)
            start_container(item)
        else:
            stop_container(item, timeout=30)
        return {"music_worker": MUSIC_WORKER,
                "state": "running" if running else "stopped"}


def active_profile(items: dict[str, dict] | None = None) -> str | None:
    items = items or containers()
    active = [name for name, item in items.items() if item.get("State") == "running"]
    if len(active) > 1:
        raise RuntimeError(f"multiple llama profiles active: {', '.join(active)}")
    return active[0] if active else None


def activate(profile: str) -> dict:
    if profile not in ALLOWED:
        raise ValueError("profile is not allowlisted")
    with LOCK:
        # Defensive mutual exclusion even if a caller bypasses the router.
        for worker in image_containers():
            stop_container(worker)
        stop_music_if_configured()
        start_container(tts_container())
        items = containers()
        missing = [name for name in ALLOWED if name not in items]
        if missing:
            raise RuntimeError("profile containers missing: " + ", ".join(missing))
        current = active_profile(items)
        if current == profile:
            return {"active_profile": current, "changed": False}
        for name, item in items.items():
            if name == profile or item.get("State") != "running":
                continue
            status, _ = docker_request("POST", f"/containers/{item['Id']}/stop?t=120")
            if status not in (204, 304):
                raise RuntimeError(f"failed to stop profile {name}: HTTP {status}")
        target = items[profile]
        status, _ = docker_request("POST", f"/containers/{target['Id']}/start")
        if status not in (204, 304):
            raise RuntimeError(f"failed to start profile {profile}: HTTP {status}")
        log.info("activated profile %s (previous=%s)", profile, current)
        return {"active_profile": profile, "changed": True}


def load_token() -> str:
    token = os.environ.get("CONTROLLER_TOKEN", "").strip()
    if not token:
        with open(TOKEN_FILE, encoding="utf-8") as handle:
            token = handle.read().strip()
    if len(token) < 32:
        raise RuntimeError("controller token is missing or too short")
    return token


TOKEN = load_token()


class Handler(BaseHTTPRequestHandler):
    server_version = "mike-ai-profile-controller/1"

    def log_message(self, fmt: str, *args: object) -> None:
        log.info("%s - %s", self.client_address[0], fmt % args)

    def reply(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, separators=(",", ":")).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def authenticated(self) -> bool:
        return self.headers.get("Authorization", "") == f"Bearer {TOKEN}"

    def do_GET(self) -> None:  # noqa: N802
        if self.path == "/health":
            self.reply(200, {"status": "ok"})
            return
        if self.path != "/status":
            self.reply(404, {"error": "not found"})
            return
        if not self.authenticated():
            self.reply(401, {"error": "unauthorized"})
            return
        try:
            items = containers()
            music = music_container() if MUSIC_WORKER else {}
            music_status = music.get("Status", "")
            music_health = ("disabled" if not MUSIC_WORKER else
                            "healthy" if "(healthy)" in music_status else
                            "unhealthy" if "(unhealthy)" in music_status else
                            "starting" if music.get("State") == "running" else
                            "stopped")
            self.reply(200, {"active_profile": active_profile(items),
                             "music_worker": music.get("State", "disabled"),
                             "music_health": music_health,
                             "profiles": {name: items.get(name, {}).get(
                                 "State", "missing") for name in ALLOWED}})
        except Exception as exc:
            log.exception("status failed")
            self.reply(503, {"error": str(exc)})

    def do_POST(self) -> None:  # noqa: N802
        if not self.authenticated():
            self.reply(401, {"error": "unauthorized"})
            return
        if self.path == "/inference/stop":
            try:
                self.reply(200, stop_inference())
            except Exception as exc:
                log.exception("stopping inference failed")
                self.reply(503, {"error": str(exc)})
            return
        if self.path in {"/workers/music/start", "/workers/music/stop"}:
            try:
                self.reply(200, set_music_worker(self.path.endswith("/start")))
            except Exception as exc:
                log.exception("music worker transition failed")
                self.reply(503, {"error": str(exc)})
            return
        worker_paths = {
            "/workers/image/start": (IMAGE_WORKER, True),
            "/workers/image/stop": (IMAGE_WORKER, False),
            "/workers/restore/start": (RESTORE_WORKER, True),
            "/workers/restore/stop": (RESTORE_WORKER, False),
        }
        if self.path in worker_paths:
            try:
                kind, running = worker_paths[self.path]
                self.reply(200, set_image_worker(running, kind))
            except Exception as exc:
                log.exception("image worker transition failed")
                self.reply(503, {"error": str(exc)})
            return
        prefix, suffix = "/profiles/", "/activate"
        if not self.path.startswith(prefix) or not self.path.endswith(suffix):
            self.reply(404, {"error": "not found"})
            return
        profile = self.path[len(prefix):-len(suffix)]
        if profile not in ALLOWED:
            self.reply(400, {"error": "profile is not allowlisted"})
            return
        try:
            self.reply(200, activate(profile))
        except Exception as exc:
            log.exception("activation failed for %s", profile)
            self.reply(503, {"error": str(exc)})


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()