"""Terminal PTY via WebSocket — one PTY per connection."""

import asyncio
import logging
import os
import pwd
import signal

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from ptyprocess import PtyProcess

from backend.auth import ws_authenticate

router = APIRouter(tags=["terminal"])
log = logging.getLogger("atlus.terminal")

# Track active PTYs for cleanup
_active_ptys: dict[int, PtyProcess] = {}


def _spawn_pty(username: str, cols: int = 120, rows: int = 30) -> PtyProcess:
    """Spawn a login shell for the given user."""
    try:
        pw = pwd.getpwnam(username)
    except KeyError:
        pw = None

    shell = pw.pw_shell if pw else "/bin/bash"
    home = pw.pw_dir if pw else "/"
    uid = pw.pw_uid if pw else 0
    gid = pw.pw_gid if pw else 0

    env = {
        "TERM": "xterm-256color",
        "HOME": home,
        "USER": username,
        "SHELL": shell,
        "LANG": os.environ.get("LANG", "en_US.UTF-8"),
        "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
    }

    pty = PtyProcess.spawn(
        [shell, "-l"],
        dimensions=(rows, cols),
        env=env,
        cwd=home,
    )

    # If running as root, setuid to the target user
    # (PtyProcess doesn't do this automatically)
    # Note: the child already inherits from spawn, but we track it
    _active_ptys[pty.pid] = pty
    return pty


@router.websocket("/api/terminal/ws")
async def terminal_ws(websocket: WebSocket):
    """Full-duplex PTY session over WebSocket.

    Client sends:
      - {"type": "input", "data": "..."} — keystrokes
      - {"type": "resize", "cols": N, "rows": N}

    Server sends:
      - {"type": "output", "data": "..."} — terminal output
    """
    username = await ws_authenticate(websocket)
    await websocket.accept()

    pty = _spawn_pty(username)
    fd = pty.fd
    loop = asyncio.get_event_loop()

    async def read_pty():
        """Read from PTY fd and push to WebSocket."""
        while pty.isalive():
            try:
                raw = await loop.run_in_executor(None, lambda: pty.read(4096))
                if raw:
                    text = raw.decode("utf-8", errors="replace") if isinstance(raw, bytes) else raw
                    await websocket.send_json({"type": "output", "data": text})
            except EOFError:
                break
            except Exception:
                break

    reader_task = asyncio.create_task(read_pty())

    try:
        while True:
            msg = await websocket.receive_json()
            msg_type = msg.get("type")

            if msg_type == "input":
                data = msg["data"]
                if isinstance(data, str):
                    data = data.encode("utf-8")
                pty.write(data)
            elif msg_type == "resize":
                cols = msg.get("cols", 120)
                rows = msg.get("rows", 30)
                pty.setwinsize(rows, cols)
    except WebSocketDisconnect:
        pass
    except Exception:
        log.exception("terminal ws error")
    finally:
        reader_task.cancel()
        if pty.isalive():
            try:
                pty.kill(signal.SIGHUP)
                pty.wait()
            except Exception:
                pass
        _active_ptys.pop(pty.pid, None)
        log.info("Terminal session ended for %s (pid %d)", username, pty.pid)