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:
        chan, session = await conn.create_session(
            asyncssh.SSHClientProcess, term_type="xterm-256color",
            term_size=(80, 24),
        )

        async def read_ssh():
            try:
                async for data in chan.stdout:
                    await websocket.send_bytes(data.encode() if isinstance(data, str) else data)
            except (asyncssh.Error, WebSocketDisconnect):
                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]
                        chan.change_terminal_size(cols, rows)
                    else:
                        chan.write(data.decode("utf-8", errors="replace"))
                elif "text" in msg and msg["text"]:
                    chan.write(msg["text"])
        except WebSocketDisconnect:
            pass
        finally:
            read_task.cancel()
    finally:
        conn.close()