"""Authenticated localhost transport for the K1 shadow perception worker."""

from __future__ import annotations

import asyncio
import hmac
import os
import secrets
from collections.abc import Callable
from contextlib import suppress
from pathlib import Path
from uuid import uuid4

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from k1link.compute.live_perception import LivePerceptionIngress

TOKEN_BYTES = 32
TOKEN_FILE_NAME = "shadow-worker.token"


def ensure_live_shadow_token(repository_root: Path) -> tuple[Path, str]:
    """Load or create the private bearer used only through the SSH tunnel."""

    token_root = repository_root.resolve() / ".runtime" / "live-perception"
    token_root.mkdir(mode=0o700, parents=True, exist_ok=True)
    with suppress(OSError):
        token_root.chmod(0o700)
    token_path = token_root / TOKEN_FILE_NAME
    try:
        token = token_path.read_text(encoding="ascii").strip()
    except FileNotFoundError:
        token = secrets.token_urlsafe(TOKEN_BYTES)
        descriptor = os.open(
            token_path,
            os.O_WRONLY | os.O_CREAT | os.O_EXCL,
            0o600,
        )
        try:
            os.write(descriptor, f"{token}\n".encode("ascii"))
            os.fsync(descriptor)
        finally:
            os.close(descriptor)
    if len(token) < 40 or len(token) > 128 or not token.isascii():
        raise RuntimeError("live perception shadow token is invalid")
    with suppress(OSError):
        token_path.chmod(0o600)
    return token_path, token


def build_live_perception_shadow_router(
    ingress: LivePerceptionIngress,
    plugin_id: str,
    *,
    bearer_token: str,
    result_receiver: Callable[[bytes], bool] | None = None,
) -> APIRouter:
    """Expose one exclusive sensor stream with bounded diagnostic results back."""

    router = APIRouter(include_in_schema=False)

    @router.websocket(
        f"/api/v1/device-plugins/{plugin_id}/live-perception-shadow"
    )
    async def live_perception_shadow(websocket: WebSocket) -> None:
        authorization = websocket.headers.get("authorization", "")
        supplied = authorization.removeprefix("Bearer ")
        if not supplied or not hmac.compare_digest(supplied, bearer_token):
            await websocket.close(code=1008, reason="Shadow worker authentication failed")
            return

        consumer_id = f"shadow-worker-{uuid4().hex}"
        try:
            ingress.open_consumer(consumer_id)
        except RuntimeError:
            await websocket.close(code=1008, reason="Shadow worker lease is unavailable")
            return

        await websocket.accept()
        client_event = asyncio.create_task(websocket.receive())
        ingress_event = asyncio.create_task(
            asyncio.to_thread(
                ingress.take_next,
                consumer_id,
                timeout=0.5,
            )
        )
        try:
            while True:
                completed, _ = await asyncio.wait(
                    {client_event, ingress_event},
                    return_when=asyncio.FIRST_COMPLETED,
                )
                if client_event in completed:
                    message = client_event.result()
                    if message.get("type") == "websocket.disconnect":
                        break
                    result = message.get("bytes")
                    if not isinstance(result, bytes) or result_receiver is None:
                        await websocket.close(
                            code=1008,
                            reason="Shadow result direction is unavailable",
                        )
                        return
                    try:
                        await asyncio.to_thread(result_receiver, result)
                    except (RuntimeError, ValueError):
                        await websocket.close(
                            code=1008,
                            reason="Shadow result contract is invalid",
                        )
                        return
                    client_event = asyncio.create_task(websocket.receive())
                if ingress_event in completed:
                    event = ingress_event.result()
                    if event is None:
                        if ingress.snapshot()["closed"]:
                            break
                    else:
                        await websocket.send_bytes(event.wire_bytes())
                    ingress_event = asyncio.create_task(
                        asyncio.to_thread(
                            ingress.take_next,
                            consumer_id,
                            timeout=0.5,
                        )
                    )
        except (WebSocketDisconnect, RuntimeError):
            return
        finally:
            client_event.cancel()
            with suppress(asyncio.CancelledError):
                await client_event
            ingress_event.cancel()
            with suppress(asyncio.CancelledError):
                await ingress_event
            ingress.close_consumer(consumer_id)
            with suppress(RuntimeError):
                await websocket.close()

    return router
