"""Atlus — X11 Display Manager.

Manages a virtual X11 display (Xvfb) with window manager (openbox),
VNC server (x11vnc), WebSocket proxy (websockify), and system tray
(stalonetray) for running native GUI apps on headless SBCs.

One shared display per user. Apps launch into the display and are
accessible via noVNC in the browser.
"""

import asyncio
import logging
import os
import shutil
import signal
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

from backend.config import DATA_DIR

log = logging.getLogger("atlus.display")

DISPLAY_DATA_DIR = DATA_DIR / "display"

# Default virtual display settings
DEFAULT_RESOLUTION = "1280x800"
DEFAULT_DEPTH = 24
XVFB_DISPLAY_BASE = 50  # Start numbering from :50 to avoid conflicts


# ---------------------------------------------------------------------------
# Dependency checks
# ---------------------------------------------------------------------------

def _find_bin(name: str) -> Optional[str]:
    """Find a binary, even with systemd's minimal PATH."""
    found = shutil.which(name)
    if found:
        return found
    for p in ("/usr/bin", "/usr/local/bin", "/usr/sbin", "/bin"):
        full = os.path.join(p, name)
        if os.path.isfile(full) and os.access(full, os.X_OK):
            return full
    return None


BINS = {
    "Xvfb": _find_bin("Xvfb"),
    "openbox": _find_bin("openbox"),
    "x11vnc": _find_bin("x11vnc"),
    "websockify": _find_bin("websockify"),
    "stalonetray": _find_bin("stalonetray"),
    "wmctrl": _find_bin("wmctrl"),
    "xdotool": _find_bin("xdotool"),
    "xdpyinfo": _find_bin("xdpyinfo"),
}


def check_dependencies() -> dict:
    """Return dependency status dict."""
    return {name: (path is not None) for name, path in BINS.items()}


def _require_display_deps():
    """Check that minimum required binaries are available."""
    missing = []
    for name in ("Xvfb", "openbox", "x11vnc", "websockify"):
        if not BINS[name]:
            missing.append(name)
    return missing


# ---------------------------------------------------------------------------
# Managed Process — async subprocess wrapper
# ---------------------------------------------------------------------------

@dataclass
class ManagedProcess:
    """A long-lived subprocess with monitoring."""
    name: str
    proc: Optional[asyncio.subprocess.Process] = None
    pid: Optional[int] = None
    started_at: float = 0
    _monitor_task: Optional[asyncio.Task] = field(default=None, repr=False)

    @property
    def alive(self) -> bool:
        return self.proc is not None and self.proc.returncode is None

    async def start(self, *args, env=None, **kwargs):
        """Launch the subprocess."""
        log.info("Starting %s: %s", self.name, " ".join(str(a) for a in args))
        self.proc = await asyncio.create_subprocess_exec(
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
            **kwargs,
        )
        self.pid = self.proc.pid
        self.started_at = time.time()
        log.info("%s started (pid %d)", self.name, self.pid)

    async def stop(self):
        """Terminate the process gracefully, then force-kill."""
        if not self.alive:
            return
        if self._monitor_task and not self._monitor_task.done():
            self._monitor_task.cancel()
        try:
            self.proc.terminate()
            try:
                await asyncio.wait_for(self.proc.wait(), timeout=5)
            except asyncio.TimeoutError:
                self.proc.kill()
                await self.proc.wait()
        except ProcessLookupError:
            pass
        log.info("%s stopped (pid %d)", self.name, self.pid)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "pid": self.pid,
            "alive": self.alive,
            "started_at": self.started_at,
        }


# ---------------------------------------------------------------------------
# X11 App — a launched GUI application
# ---------------------------------------------------------------------------

@dataclass
class X11App:
    """A GUI application running on the virtual display."""
    app_id: str
    command: str
    args: list[str] = field(default_factory=list)
    name: str = ""
    proc: Optional[asyncio.subprocess.Process] = None
    pid: Optional[int] = None
    started_at: float = 0
    has_tray_icon: bool = False

    @property
    def alive(self) -> bool:
        return self.proc is not None and self.proc.returncode is None

    async def start(self, env: dict):
        """Launch the application."""
        cmd = [self.command] + self.args
        log.info("Launching X11 app %s: %s", self.app_id, " ".join(cmd))
        try:
            self.proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env,
            )
            self.pid = self.proc.pid
            self.started_at = time.time()
            log.info("X11 app %s started (pid %d)", self.app_id, self.pid)
        except Exception:
            log.exception("Failed to launch X11 app %s", self.app_id)
            raise

    async def stop(self):
        """Terminate the application."""
        if not self.alive:
            return
        try:
            self.proc.terminate()
            try:
                await asyncio.wait_for(self.proc.wait(), timeout=5)
            except asyncio.TimeoutError:
                self.proc.kill()
                await self.proc.wait()
        except ProcessLookupError:
            pass
        log.info("X11 app %s stopped", self.app_id)

    def to_dict(self) -> dict:
        return {
            "app_id": self.app_id,
            "command": self.command,
            "name": self.name or self.command,
            "alive": self.alive,
            "pid": self.pid,
            "started_at": self.started_at,
            "has_tray_icon": self.has_tray_icon,
        }


# ---------------------------------------------------------------------------
# Display Session — one virtual display per user
# ---------------------------------------------------------------------------

class DisplaySession:
    """Manages a complete virtual X11 display with all supporting services."""

    def __init__(self, username: str, display_num: int):
        self.username = username
        self.display_num = display_num
        self.display = f":{display_num}"
        self.resolution = DEFAULT_RESOLUTION
        self.depth = DEFAULT_DEPTH

        # Managed processes
        self.xvfb = ManagedProcess("Xvfb")
        self.openbox = ManagedProcess("openbox")
        self.x11vnc = ManagedProcess("x11vnc")
        self.websockify = ManagedProcess("websockify")
        self.stalonetray = ManagedProcess("stalonetray")

        # Launched GUI apps
        self.apps: dict[str, X11App] = {}

        # VNC/WebSocket ports
        self.vnc_port = 5900 + display_num
        self.ws_port = 6080 + display_num

        self.started = False
        self.created_at = time.time()
        self.last_active = time.time()

        # Data directory for this session
        self.session_dir = DISPLAY_DATA_DIR / username
        self.session_dir.mkdir(parents=True, exist_ok=True)

    def _display_env(self) -> dict:
        """Build environment dict with DISPLAY set."""
        env = {**os.environ}
        env["DISPLAY"] = self.display
        env["HOME"] = os.path.expanduser(f"~{self.username}") or "/"
        env["USER"] = self.username
        env["XAUTHORITY"] = str(self.session_dir / ".Xauthority")
        # Ensure full PATH for systemd
        path = env.get("PATH", "")
        for p in ("/usr/local/bin", "/usr/bin", "/usr/sbin", "/bin", "/sbin"):
            if p not in path:
                path = p + ":" + path
        env["PATH"] = path
        env["LC_ALL"] = "C"
        return env

    async def start(self):
        """Start the complete display stack: Xvfb → openbox → stalonetray → x11vnc → websockify."""
        if self.started:
            return

        env = self._display_env()

        # 1. Start Xvfb
        xvfb_bin = BINS["Xvfb"]
        await self.xvfb.start(
            xvfb_bin,
            self.display,
            "-screen", "0", f"{self.resolution}x{self.depth}",
            "-ac",  # disable access control
            "-nolisten", "tcp",
            env=env,
        )

        # Wait for display to be ready
        await self._wait_for_display(env, timeout=10)

        # 2. Start openbox (window manager)
        openbox_bin = BINS["openbox"]
        await self.openbox.start(
            openbox_bin,
            "--config-file", "/dev/null",
            env=env,
        )
        await asyncio.sleep(0.5)

        # 3. Start stalonetray (system tray)
        if BINS["stalonetray"]:
            stalonetray_bin = BINS["stalonetray"]
            # Configure stalonetray: small, at top-right, transparent
            await self.stalonetray.start(
                stalonetray_bin,
                "--geometry", "1x1+0+0",
                "--icon-size", "24",
                "--kludges", "force_icons_size",
                "--window-type", "dock",
                "--grow-gravity", "E",
                "--icon-gravity", "E",
                "--background", "#111318",
                "--no-shrink",
                env=env,
            )
            await asyncio.sleep(0.3)

        # 4. Start x11vnc
        x11vnc_bin = BINS["x11vnc"]
        await self.x11vnc.start(
            x11vnc_bin,
            "-display", self.display,
            "-rfbport", str(self.vnc_port),
            "-shared",
            "-forever",
            "-nopw",
            "-noxdamage",
            "-cursor", "most",
            "-ncache", "10",
            env=env,
        )
        await asyncio.sleep(0.5)

        # 5. Start websockify (WebSocket → VNC proxy)
        websockify_bin = BINS["websockify"]
        # Serve noVNC client files if available
        novnc_dir = self._find_novnc()
        ws_args = [
            websockify_bin,
            "--heartbeat", "30",
            str(self.ws_port),
            f"localhost:{self.vnc_port}",
        ]
        if novnc_dir:
            ws_args.insert(-2, "--web")
            ws_args.insert(-2, str(novnc_dir))

        await self.websockify.start(*ws_args, env=env)
        await asyncio.sleep(0.3)

        self.started = True
        self.last_active = time.time()
        log.info("Display session started for %s on %s (VNC:%d, WS:%d)",
                 self.username, self.display, self.vnc_port, self.ws_port)

    async def _wait_for_display(self, env: dict, timeout: float = 10):
        """Wait until the Xvfb display is accepting connections."""
        xdpyinfo = BINS.get("xdpyinfo")
        if not xdpyinfo:
            # No xdpyinfo available — just wait a bit
            await asyncio.sleep(1.5)
            return

        start = time.time()
        while time.time() - start < timeout:
            proc = await asyncio.create_subprocess_exec(
                xdpyinfo, "-display", self.display,
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
                env=env,
            )
            await proc.wait()
            if proc.returncode == 0:
                log.info("Display %s is ready", self.display)
                return
            await asyncio.sleep(0.3)
        log.warning("Display %s may not be ready after %.1fs", self.display, timeout)

    def _find_novnc(self) -> Optional[Path]:
        """Find noVNC installation directory."""
        candidates = [
            Path("/usr/share/novnc"),
            Path("/usr/share/noVNC"),
            Path("/usr/local/share/novnc"),
            Path("/usr/local/share/noVNC"),
            Path("/opt/novnc"),
            Path("/opt/noVNC"),
            Path("/snap/novnc/current/usr/share/novnc"),
        ]
        for p in candidates:
            if p.is_dir() and (p / "vnc.html").exists():
                return p
            if p.is_dir() and (p / "vnc_lite.html").exists():
                return p
        return None

    async def stop(self):
        """Stop everything in reverse order."""
        # Stop all apps first
        for app in list(self.apps.values()):
            await app.stop()
        self.apps.clear()

        # Stop services in reverse order
        for svc in (self.websockify, self.x11vnc, self.stalonetray, self.openbox, self.xvfb):
            await svc.stop()

        self.started = False
        log.info("Display session stopped for %s", self.username)

    async def launch_app(self, command: str, args: list[str] = None,
                         name: str = "", has_tray_icon: bool = False) -> X11App:
        """Launch a GUI application on this display."""
        if not self.started:
            await self.start()

        app_id = str(uuid.uuid4())[:8]
        app = X11App(
            app_id=app_id,
            command=command,
            args=args or [],
            name=name or os.path.basename(command),
            has_tray_icon=has_tray_icon,
        )

        env = self._display_env()
        await app.start(env)

        self.apps[app_id] = app
        self.last_active = time.time()

        # Start a monitor task that removes app when it exits
        asyncio.create_task(self._monitor_app(app_id))

        return app

    async def _monitor_app(self, app_id: str):
        """Wait for an app to exit, then clean up."""
        app = self.apps.get(app_id)
        if not app or not app.proc:
            return
        try:
            await app.proc.wait()
        except Exception:
            pass
        # Remove from apps dict (it's already dead)
        self.apps.pop(app_id, None)
        log.info("X11 app %s (%s) exited", app_id, app.name)

    async def stop_app(self, app_id: str) -> bool:
        """Stop a specific app."""
        app = self.apps.pop(app_id, None)
        if not app:
            return False
        await app.stop()
        return True

    async def list_windows(self) -> list[dict]:
        """List windows using wmctrl."""
        wmctrl = BINS.get("wmctrl")
        if not wmctrl or not self.started:
            return []

        env = self._display_env()
        try:
            proc = await asyncio.create_subprocess_exec(
                wmctrl, "-l", "-p",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env,
            )
            stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=5)
            if proc.returncode != 0:
                return []

            windows = []
            for line in stdout.decode().strip().split("\n"):
                if not line.strip():
                    continue
                parts = line.split(None, 4)
                if len(parts) >= 5:
                    windows.append({
                        "window_id": parts[0],
                        "desktop": parts[1],
                        "pid": int(parts[2]) if parts[2].isdigit() else 0,
                        "host": parts[3],
                        "title": parts[4],
                    })
            return windows
        except Exception:
            return []

    async def focus_window(self, window_id: str) -> bool:
        """Activate a window using wmctrl."""
        wmctrl = BINS.get("wmctrl")
        if not wmctrl or not self.started:
            return False

        env = self._display_env()
        try:
            proc = await asyncio.create_subprocess_exec(
                wmctrl, "-i", "-a", window_id,
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
                env=env,
            )
            await asyncio.wait_for(proc.wait(), timeout=5)
            return proc.returncode == 0
        except Exception:
            return False

    def touch(self):
        self.last_active = time.time()

    def to_dict(self) -> dict:
        return {
            "username": self.username,
            "display": self.display,
            "resolution": self.resolution,
            "vnc_port": self.vnc_port,
            "ws_port": self.ws_port,
            "started": self.started,
            "created_at": self.created_at,
            "last_active": self.last_active,
            "apps": {aid: a.to_dict() for aid, a in self.apps.items() if a.alive},
            "services": {
                "xvfb": self.xvfb.to_dict(),
                "openbox": self.openbox.to_dict(),
                "x11vnc": self.x11vnc.to_dict(),
                "websockify": self.websockify.to_dict(),
                "stalonetray": self.stalonetray.to_dict(),
            },
        }


# ---------------------------------------------------------------------------
# Display Manager — singleton
# ---------------------------------------------------------------------------

class DisplayManager:
    """Manages all user display sessions."""

    def __init__(self):
        self._sessions: dict[str, DisplaySession] = {}
        self._next_display = XVFB_DISPLAY_BASE
        DISPLAY_DATA_DIR.mkdir(parents=True, exist_ok=True)

    def _allocate_display(self) -> int:
        """Allocate the next available display number."""
        num = self._next_display
        self._next_display += 1
        return num

    async def get_or_create(self, username: str) -> DisplaySession:
        """Get existing display session or create (but don't start) a new one."""
        if username not in self._sessions:
            display_num = self._allocate_display()
            self._sessions[username] = DisplaySession(username, display_num)
            log.info("Created display session for %s on :%d", username, display_num)
        session = self._sessions[username]
        session.touch()
        return session

    async def ensure_started(self, username: str) -> DisplaySession:
        """Get or create a display session and ensure it's started."""
        session = await self.get_or_create(username)
        if not session.started:
            await session.start()
        return session

    def get_session(self, username: str) -> Optional[DisplaySession]:
        """Get display session if it exists."""
        return self._sessions.get(username)

    async def stop_session(self, username: str) -> bool:
        """Stop and remove a display session."""
        session = self._sessions.pop(username, None)
        if not session:
            return False
        await session.stop()
        return True

    async def shutdown_all(self):
        """Stop all display sessions — called on server shutdown."""
        for username in list(self._sessions.keys()):
            await self.stop_session(username)
        log.info("All display sessions shut down")


# ---------------------------------------------------------------------------
# Singleton
# ---------------------------------------------------------------------------

display_manager = DisplayManager()