import asyncio
import struct
import logging
import time
from fastapi import WebSocket, WebSocketDisconnect, Query
import asyncssh
import jwt
import config

logger = logging.getLogger(__name__)

MAX_SESSIONS = 5
IDLE_TIMEOUT = 30 * 60  # 30 minutes
_active_sessions: set = set()

HOSTS = {
    "nas": {"host": config.SSH_HOST, "user": config.SSH_USER, "key": config.SSH_KEY_PATH},
    "vps": {"host": config.VPS_SSH_HOST, "user": config.VPS_SSH_USER, "key": config.VPS_SSH_KEY_PATH},
    "caddy-vps": {"host": config.CADDY_VPS_SSH_HOST, "user": config.CADDY_VPS_SSH_USER, "key": config.VPS_SSH_KEY_PATH},
    "claude-dev": {"host": config.SSH_HOST, "user": config.SSH_USER, "key": config.SSH_KEY_PATH,
                   "command": "/volume1/@appstore/ContainerManager/usr/bin/docker exec -it claude-dev bash"},
}

# Load known_hosts once at import time — fail closed if missing
_known_hosts_available = False
try:
    _known_hosts = asyncssh.read_known_hosts(config.SSH_KNOWN_HOSTS)
    _known_hosts_available = True
except Exception:
    _known_hosts = None
    logger.error("SSH known_hosts file not found at %s — terminal connections will be refused", config.SSH_KNOWN_HOSTS)


async def ws_endpoint(websocket: WebSocket, host: str = Query("nas")):
    await websocket.accept()

    # Auth: expect token as first text message
    try:
        first_msg = await asyncio.wait_for(websocket.receive_text(), timeout=5)
        payload = jwt.decode(first_msg, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        if payload.get("type") != "access":
            await websocket.close(code=1008, reason="Unauthorized")
            return
        username = payload.get("sub")
        role = payload.get("role", "admin")
        if username is None:
            await websocket.close(code=1008, reason="Unauthorized")
            return
        # RBAC: check terminal page access
        from rbac import get_pages
        pages = get_pages(username, role)
        if pages != "*" and "terminal" not in pages:
            await websocket.close(code=1008, reason="Access denied")
            return
    except Exception:
        await websocket.close(code=1008, reason="Unauthorized")
        return

    if not _known_hosts_available:
        await websocket.close(code=1011, reason="SSH host verification unavailable")
        return

    if len(_active_sessions) >= MAX_SESSIONS:
        await websocket.close(code=1013, reason="Too many sessions")
        return

    session_id = id(websocket)
    _active_sessions.add(session_id)

    h = HOSTS.get(host, HOSTS["nas"])
    try:
        conn = await asyncssh.connect(
            h["host"], username=h["user"],
            client_keys=[h["key"]], known_hosts=_known_hosts,
        )
    except Exception:
        _active_sessions.discard(session_id)
        await websocket.close(code=1011, reason="SSH connection failed")
        return

    try:
        process = await conn.create_process(
            h.get("command"), term_type="xterm-256color", term_size=(80, 24),
            encoding=None,
        )

        async def read_ssh():
            try:
                while True:
                    data = await process.stdout.read(4096)
                    if not data:
                        break
                    await websocket.send_bytes(data)
            except (asyncssh.Error, WebSocketDisconnect, asyncio.CancelledError):
                pass

        read_task = asyncio.create_task(read_ssh())
        last_activity = time.monotonic()

        try:
            while True:
                try:
                    msg = await asyncio.wait_for(websocket.receive(), timeout=60)
                except asyncio.TimeoutError:
                    if time.monotonic() - last_activity > IDLE_TIMEOUT:
                        await websocket.close(code=1000, reason="Idle timeout")
                        break
                    continue
                last_activity = time.monotonic()
                if "bytes" in msg and msg["bytes"]:
                    data = msg["bytes"]
                    if data[0:1] == b"\x01" and len(data) == 5:
                        cols = struct.unpack("!H", data[1:3])[0]
                        rows = struct.unpack("!H", data[3:5])[0]
                        process.channel.change_terminal_size(cols, rows)
                    else:
                        process.stdin.write(data)
                elif "text" in msg and msg["text"]:
                    process.stdin.write(msg["text"].encode())
        except WebSocketDisconnect:
            pass
        finally:
            read_task.cancel()
            process.close()
    finally:
        conn.close()
        _active_sessions.discard(session_id)