import asyncio
import struct
import logging
from collections import Counter
from fastapi import WebSocket, Query
import asyncssh
import config
from auth import get_current_user_ws

logger = logging.getLogger(__name__)

MAX_SESSIONS = 5
MAX_SESSIONS_PER_USER = 2
_session_lock = asyncio.Lock()
_active_sessions: dict[int, str] = {}

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},
    "mac": {"host": config.MAC_SSH_HOST, "user": config.MAC_SSH_USER, "key": config.MAC_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 _reserve_session(session_id: int, username: str) -> tuple[bool, str | None, int | None]:
    async with _session_lock:
        if len(_active_sessions) >= MAX_SESSIONS:
            return False, "Too many sessions", 1013
        user_sessions = Counter(_active_sessions.values())
        if user_sessions[username] >= MAX_SESSIONS_PER_USER:
            return False, "Too many sessions for user", 1013
        _active_sessions[session_id] = username
        return True, None, session_id


async def _release_session(session_id: int | None):
    if session_id is None:
        return
    async with _session_lock:
        _active_sessions.pop(session_id, None)


async def ws_endpoint(websocket: WebSocket, host: str = Query("nas"), token: str | None = Query(None)):
    if host not in HOSTS:
        await websocket.close(code=1008, reason="Unknown host")
        return

    if not token:
        await websocket.close(code=1008, reason="Unauthorized")
        return

    try:
        user = await get_current_user_ws(token)
    except Exception:
        await websocket.close(code=1008, reason="Unauthorized")
        return

    if user.pages != "*" and "terminal" not in user.pages:
        await websocket.close(code=1008, reason="Access denied")
        return

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

    session_id = id(websocket)
    reserved, reason, session_or_code = await _reserve_session(session_id, user.username)
    if not reserved:
        await websocket.close(code=session_or_code, reason=reason)
        return
    session_id = session_or_code

    await websocket.accept()

    h = HOSTS[host]
    try:
        conn = await asyncssh.connect(
            h["host"], username=h["user"],
            client_keys=[h["key"]], known_hosts=_known_hosts,
            keepalive_interval=30,
            keepalive_count_max=3,
        )
    except Exception:
        await _release_session(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, asyncio.CancelledError):
                pass

        read_task = asyncio.create_task(read_ssh())

        try:
            while True:
                msg = await websocket.receive()
                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"]:
                    if msg["text"] == "__ping__":
                        continue
                    process.stdin.write(msg["text"].encode())
        except Exception:
            pass
        finally:
            read_task.cancel()
            process.close()
    finally:
        conn.close()
        await _release_session(session_id)