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


async def ws_endpoint(websocket: WebSocket, token: str = Query("")):
    try:
        payload = jwt.decode(token, config.SECRET_KEY, algorithms=[config.ALGORITHM])
        if payload.get("sub") != config.ADMIN_USER:
            await websocket.close(code=1008, reason="Unauthorized")
            return
    except Exception:
        await websocket.close(code=1008, reason="Unauthorized")
        return

    await websocket.accept()

    try:
        conn = await asyncssh.connect(
            config.SSH_HOST, username=config.SSH_USER,
            client_keys=[config.SSH_KEY_PATH], known_hosts=None,
        )
    except Exception:
        await websocket.close(code=1011, reason="SSH connection failed")
        return

    try:
        process = await conn.create_process(
            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())

        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"]:
                    process.stdin.write(msg["text"].encode())
        except WebSocketDisconnect:
            pass
        finally:
            read_task.cancel()
            process.close()
    finally:
        conn.close()