import asyncio
import os
import pty
import signal
import struct
import fcntl
import termios
from fastapi import WebSocket, WebSocketDisconnect


async def ws_endpoint(ws: WebSocket):
    # Block non-Tailscale IPs
    client_ip = ws.client.host if ws.client else ""
    if not client_ip.startswith("100."):
        await ws.close(code=1008, reason="Tailscale only")
        return

    await ws.accept()
    master_fd, slave_fd = pty.openpty()
    pid = os.fork()

    if pid == 0:
        os.close(master_fd)
        os.setsid()
        os.dup2(slave_fd, 0)
        os.dup2(slave_fd, 1)
        os.dup2(slave_fd, 2)
        os.close(slave_fd)
        os.execvp("/bin/bash", ["/bin/bash"])

    os.close(slave_fd)
    loop = asyncio.get_event_loop()

    async def read_pty():
        try:
            while True:
                await asyncio.sleep(0.01)
                if select_readable(master_fd):
                    data = os.read(master_fd, 4096)
                    if not data:
                        break
                    await ws.send_bytes(data)
        except (OSError, WebSocketDisconnect):
            pass

    def select_readable(fd):
        import select
        r, _, _ = select.select([fd], [], [], 0)
        return bool(r)

    read_task = asyncio.create_task(read_pty())

    try:
        while True:
            msg = await ws.receive()
            if "bytes" in msg and msg["bytes"]:
                data = msg["bytes"]
                # Handle resize: first byte 0x01 means resize message
                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]
                    fcntl.ioctl(master_fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
                else:
                    os.write(master_fd, data)
            elif "text" in msg and msg["text"]:
                os.write(master_fd, msg["text"].encode())
    except WebSocketDisconnect:
        pass
    finally:
        read_task.cancel()
        os.close(master_fd)
        os.kill(pid, signal.SIGTERM)
        os.waitpid(pid, 0)