#
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 MindTenet LLC. All rights reserved.
# SPDX-License-Identifier: AGPL-3.0-or-later
#
"""WebSocket terminal routes for the hijack hub.

Registers:
- ``/ws/worker/{worker_id}/term``  — worker → hub (terminal output, snapshots)
- ``/ws/browser/{worker_id}/term`` — browser → hub (dashboard viewer + hijack control)
"""

from __future__ import annotations

import asyncio
import json
import secrets
import time
from contextlib import suppress
from typing import TYPE_CHECKING, Annotated, Any

from undef.telemetry import get_logger

try:
    from fastapi import APIRouter, Path, WebSocket, WebSocketDisconnect, WebSocketException
except ImportError as _e:  # pragma: no cover
    raise ImportError("fastapi is required for hijack routes: pip install 'undef-terminal[websocket]'") from _e


from undef.terminal.hijack.hub.connections import _background_tasks
from undef.terminal.hijack.models import VALID_ROLES, _safe_float, _safe_int, extract_prompt_id
from undef.terminal.hijack.ratelimit import TokenBucket
from undef.terminal.hijack.routes.browser_handlers import handle_browser_message

if TYPE_CHECKING:
    from undef.terminal.hijack.hub import BrowserRoleResolutionError, TermHub
else:
    from undef.terminal.hijack.hub import BrowserRoleResolutionError

logger = get_logger(__name__)
_WORKER_HIJACK_CLEANUP_INTERVAL_S = 1.0
_BROWSER_HIJACK_CLEANUP_INTERVAL_S = 1.0


async def _periodic_hijack_cleanup(hub: TermHub, worker_id: str, interval_s: float) -> None:
    """Run lease cleanup on a fixed cadence while a WS handler is active."""
    while True:
        await asyncio.sleep(interval_s)
        await hub.cleanup_expired_hijack(worker_id)


def register_ws_routes(hub: TermHub, router: APIRouter) -> None:
    """Attach WebSocket terminal routes to *router*."""

    @router.websocket("/ws/worker/{worker_id}/term")
    async def ws_worker_term(websocket: WebSocket, worker_id: Annotated[str, Path(pattern=r"^[\w\-]+$")]) -> None:
        worker_token = hub.worker_token()
        if worker_token is not None:
            auth_header = websocket.headers.get("authorization", "")
            provided = auth_header.removeprefix("Bearer ").strip() if auth_header.startswith("Bearer ") else ""
            if not secrets.compare_digest(provided, worker_token):
                # Accept first so the close code is transmitted to the client.
                # Calling close() before accept() silently drops the connection
                # without sending the 1008 policy-violation code.
                await websocket.accept()
                await websocket.close(code=1008, reason="authentication required")
                return
        await websocket.accept()
        # Register worker, atomically clearing any stale hijack state from a
        # crashed previous connection.  A crashed worker may reconnect before its
        # old finally block clears state; the identity check `worker_ws is old_ws`
        # in deregister_worker skips cleanup when a new connection has already
        # overwritten worker_ws, so stale REST clients cannot send keystrokes under
        # a dead session.
        prev_was_hijacked = await hub.register_worker(worker_id, websocket)
        logger.info("term_worker_connected worker_id=%s", worker_id)
        if prev_was_hijacked:
            hub.notify_hijack_changed(worker_id, enabled=False, owner=None)
            await hub.broadcast_hijack_state(worker_id)
        await hub.broadcast(
            worker_id,
            {"type": "worker_connected", "worker_id": worker_id, "ts": time.time()},
        )
        await hub.request_snapshot(worker_id)

        cleanup_task = asyncio.create_task(_periodic_hijack_cleanup(hub, worker_id, _WORKER_HIJACK_CLEANUP_INTERVAL_S))
        try:
            while True:
                raw = await websocket.receive_text()
                if len(raw) > hub.max_ws_message_bytes:
                    logger.warning("ws_worker_oversized worker_id=%s size=%d", worker_id, len(raw))
                    continue
                if not await hub.is_active_worker(worker_id, websocket):
                    with suppress(Exception):
                        await websocket.close()
                    break
                try:
                    msg = json.loads(raw)
                except Exception as exc:
                    logger.debug("ws_worker_bad_json worker_id=%s: %s", worker_id, exc)
                    continue
                mtype = msg.get("type")
                if mtype not in {"worker_hello", "term", "snapshot", "analysis", "status"}:
                    logger.debug("ws_worker_ignored worker_id=%s type=%r", worker_id, mtype)
                    continue
                if mtype == "worker_hello":
                    _hello_mode = msg.get("input_mode")
                    if _hello_mode in ("hijack", "open"):
                        mode_applied = await hub.set_worker_hello_mode(worker_id, _hello_mode)
                        if mode_applied:
                            await hub.broadcast_hijack_state(worker_id)
                        logger.info(
                            "worker_hello worker_id=%s input_mode=%s applied=%s", worker_id, _hello_mode, mode_applied
                        )
                    elif _hello_mode is not None:
                        logger.warning(
                            "worker_hello_invalid_mode worker_id=%s input_mode=%r — expected 'hijack' or 'open', ignoring",
                            worker_id,
                            _hello_mode,
                        )
                    continue
                if mtype == "term":
                    data = msg.get("data", "")
                    if data:
                        await hub.broadcast(worker_id, {"type": "term", "data": data, "ts": msg.get("ts", time.time())})
                elif mtype == "snapshot":
                    snapshot: dict[str, Any] = {
                        "type": "snapshot",
                        "screen": msg.get("screen", ""),
                        "cursor": msg.get("cursor", {"x": 0, "y": 0}),
                        "cols": _safe_int(msg.get("cols"), 80, min_val=1),
                        "rows": _safe_int(msg.get("rows"), 25, min_val=1),
                        "screen_hash": msg.get("screen_hash", ""),
                        "cursor_at_end": bool(msg.get("cursor_at_end", True)),
                        "has_trailing_space": bool(msg.get("has_trailing_space", False)),
                        "prompt_detected": msg.get("prompt_detected"),
                        "ts": _safe_float(msg.get("ts"), time.time()),
                    }
                    await hub.update_last_snapshot(worker_id, snapshot)
                    await hub.broadcast(worker_id, snapshot)
                    await hub.append_event(
                        worker_id,
                        "snapshot",
                        {"prompt_id": extract_prompt_id(snapshot), "screen_hash": snapshot.get("screen_hash")},
                    )
                elif mtype == "analysis":
                    await hub.broadcast(
                        worker_id,
                        {
                            "type": "analysis",
                            "formatted": msg.get("formatted", ""),
                            "raw": msg.get("raw"),
                            "ts": msg.get("ts", time.time()),
                        },
                    )
                elif mtype == "status":  # pragma: no branch
                    await hub.broadcast(worker_id, msg)
                    await hub.append_event(worker_id, "worker_status", {"status": msg})
        except WebSocketDisconnect:
            pass
        except Exception as exc:  # pragma: no cover
            logger.warning("term_worker_ws_error worker_id=%s error=%s", worker_id, exc)
        finally:
            cleanup_task.cancel()
            with suppress(asyncio.CancelledError):
                await cleanup_task
            should_broadcast, was_hijacked = await hub.deregister_worker(worker_id, websocket)
            if should_broadcast:
                hub.metric("ws_disconnect_total")
                hub.metric("ws_disconnect_worker_total")
                logger.info("term_worker_disconnected worker_id=%s", worker_id)
                _broadcast_task = asyncio.create_task(
                    hub.broadcast(
                        worker_id,
                        {"type": "worker_disconnected", "worker_id": worker_id, "ts": time.time()},
                    )
                )
                _background_tasks.add(_broadcast_task)
                _broadcast_task.add_done_callback(_background_tasks.discard)
                _broadcast_task.add_done_callback(
                    lambda t: (
                        logger.warning(
                            "worker_disconnected_broadcast_failed worker_id=%s error=%s", worker_id, t.exception()
                        )
                        if not t.cancelled() and t.exception() is not None
                        else None
                    )
                )
                if was_hijacked:
                    hub.notify_hijack_changed(worker_id, enabled=False, owner=None)
                    _hijack_state_task = asyncio.create_task(hub.broadcast_hijack_state(worker_id))
                    _background_tasks.add(_hijack_state_task)
                    _hijack_state_task.add_done_callback(_background_tasks.discard)
                    _hijack_state_task.add_done_callback(
                        lambda t: (
                            logger.warning(
                                "broadcast_hijack_state_failed worker_id=%s error=%s", worker_id, t.exception()
                            )
                            if not t.cancelled() and t.exception() is not None
                            else None
                        )
                    )
            await hub.prune_if_idle(worker_id)

    @router.websocket("/ws/browser/{worker_id}/term")
    async def ws_browser_term(
        websocket: WebSocket,
        worker_id: Annotated[str, Path(pattern=r"^[\w\-]+$")],
    ) -> None:
        await websocket.accept()
        try:
            role = await hub.resolve_role_for_browser(websocket, worker_id)
        except BrowserRoleResolutionError:
            await websocket.close(code=1008, reason="browser role resolution failed")
            return
        except WebSocketException:
            raise  # re-raise so FastAPI closes the already-accepted socket with the exception's code
        if role not in VALID_ROLES:  # pragma: no cover
            role = "viewer"
        can_hijack = role == "admin"
        # True once this browser has owned a dashboard hijack this session.
        # Retained even after the hijack is released so the finally block can
        # send a resume if the worker is still paused.  Does NOT reflect current
        # ownership — check hub state for that.
        owned_hijack = False
        # Capture all startup state atomically while registering the browser.
        browser_state = await hub.register_browser(worker_id, websocket, role)
        is_hijacked = browser_state["is_hijacked"]
        hijacked_by_me = browser_state["hijacked_by_me"]
        worker_online = browser_state["worker_online"]
        input_mode = browser_state["input_mode"]
        initial_snapshot = browser_state["initial_snapshot"]

        await websocket.send_text(
            json.dumps(
                {
                    "type": "hello",
                    "worker_id": worker_id,
                    "can_hijack": can_hijack,
                    "hijacked": is_hijacked,
                    "hijacked_by_me": hijacked_by_me,
                    "worker_online": worker_online,
                    "input_mode": input_mode,
                    "role": role,
                    "hijack_control": "ws",
                    "hijack_step_supported": True,
                    "capabilities": {
                        "hijack_control": "ws",
                        "hijack_step_supported": True,
                    },
                },
                ensure_ascii=True,
            )
        )
        await websocket.send_text(json.dumps(await hub.hijack_state_msg_for(worker_id, websocket), ensure_ascii=True))

        if initial_snapshot is not None:
            await websocket.send_text(json.dumps(initial_snapshot, ensure_ascii=True))
        else:
            await hub.request_snapshot(worker_id)

        cleanup_task = asyncio.create_task(_periodic_hijack_cleanup(hub, worker_id, _BROWSER_HIJACK_CLEANUP_INTERVAL_S))
        try:
            _browser_bucket = TokenBucket(hub.browser_rate_limit_per_sec)
            while True:
                raw = await websocket.receive_text()
                if len(raw) > hub.max_ws_message_bytes:
                    logger.warning("ws_browser_oversized worker_id=%s size=%d", worker_id, len(raw))
                    continue
                try:
                    msg_b: dict[str, Any] = json.loads(raw)
                except Exception as exc:
                    logger.debug("ws_browser_bad_json worker_id=%s: %s", worker_id, exc)
                    continue
                mtype = msg_b.get("type")
                if mtype == "input" and not _browser_bucket.allow():
                    logger.warning("ws_browser_rate_limited worker_id=%s", worker_id)
                    with suppress(Exception):
                        await websocket.send_text(
                            json.dumps({"type": "error", "message": "rate_limited"}, ensure_ascii=True)
                        )
                    continue

                owned_hijack = await handle_browser_message(hub, websocket, worker_id, role, msg_b, owned_hijack)

        except WebSocketDisconnect:
            pass
        except Exception as exc:  # pragma: no cover
            logger.warning("term_browser_ws_error worker_id=%s error=%s", worker_id, exc)
        finally:
            hub.metric("ws_disconnect_total")
            hub.metric("ws_disconnect_browser_total")
            cleanup_task.cancel()
            with suppress(asyncio.CancelledError):
                await cleanup_task
            # Atomically detect ownership, capture REST-session liveness, and clear
            # the owner — all in one lock block to avoid the TOCTOU window where
            # _is_owner() returns True but another coroutine steals hijack_owner
            # (or vice-versa), and to avoid a second lock round-trip for
            # has_valid_rest_lease after the owner has already been cleared.
            disconnect_result = await hub.cleanup_browser_disconnect(worker_id, websocket, owned_hijack)
            was_owner = disconnect_result["was_owner"]
            rest_still_active = disconnect_result["rest_still_active"]
            resume_without_owner = disconnect_result["resume_without_owner"]
            if was_owner:
                _do_resume = not rest_still_active
                # Re-check: a concurrent hijack_acquire may have written a new
                # session between the lock release above and _send_worker below.
                if _do_resume and await hub.check_still_hijacked(worker_id):
                    _do_resume = False
                if _do_resume:
                    _resume_task = asyncio.create_task(
                        hub.send_worker(
                            worker_id,
                            {
                                "type": "control",
                                "action": "resume",
                                "owner": "dashboard",
                                "lease_s": 0,
                                "ts": time.time(),
                            },
                        )
                    )
                    _background_tasks.add(_resume_task)
                    _resume_task.add_done_callback(_background_tasks.discard)
                    _resume_task.add_done_callback(
                        lambda t: (
                            logger.warning(
                                "ws_disconnect_resume_failed worker_id=%s error=%s", worker_id, t.exception()
                            )
                            if not t.cancelled() and t.exception() is not None
                            else None
                        )
                    )
                await hub.broadcast_hijack_state(worker_id)
                if _do_resume:
                    hub.notify_hijack_changed(worker_id, enabled=False, owner=None)
                await hub.append_event(worker_id, "hijack_released", {"owner": "dashboard_ws_disconnect"})
            elif resume_without_owner:
                if await hub.check_still_hijacked(worker_id):
                    resume_without_owner = False
                if resume_without_owner:
                    _resume_task = asyncio.create_task(
                        hub.send_worker(
                            worker_id,
                            {
                                "type": "control",
                                "action": "resume",
                                "owner": "dashboard",
                                "lease_s": 0,
                                "ts": time.time(),
                            },
                        )
                    )
                    _background_tasks.add(_resume_task)
                    _resume_task.add_done_callback(_background_tasks.discard)
                    _resume_task.add_done_callback(
                        lambda t: (
                            logger.warning(
                                "ws_disconnect_resume_failed worker_id=%s error=%s", worker_id, t.exception()
                            )
                            if not t.cancelled() and t.exception() is not None
                            else None
                        )
                    )
            await hub.prune_if_idle(worker_id)
