import os
import pty
import asyncio
import struct
import fcntl
import termios
import json

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()

app.mount("/static", StaticFiles(directory="static"), name="static")


@app.get("/")
async def index():
    return FileResponse("static/index.html")


@app.websocket("/ws")
async def terminal(websocket: WebSocket):

    await websocket.accept()

    master_fd, slave_fd = pty.openpty()

    proc = await asyncio.create_subprocess_exec(
        "/bin/bash",
        stdin=slave_fd,
        stdout=slave_fd,
        stderr=slave_fd,
        env={
            **os.environ,
            "TERM": "xterm-256color",
            "HOME": "/home/ubuntu",
            "USER": "ubuntu"
        }
    )

    os.close(slave_fd)

    loop = asyncio.get_event_loop()

    # Read terminal output
    async def read_pty():

        while True:

            try:

                data = await loop.run_in_executor(
                    None,
                    lambda: os.read(master_fd, 4096)
                )

                if data:

                    try:
                        text = data.decode("utf-8")

                    except UnicodeDecodeError:
                        text = data.decode(
                            "utf-8",
                            errors="ignore"
                        )

                    await websocket.send_text(text)

            except OSError:
                break

    read_task = asyncio.create_task(read_pty())

    try:

        while True:

            message = await websocket.receive_text()

            try:

                msg = json.loads(message)

                if msg.get("type") == "input":

                    os.write(
                        master_fd,
                        msg["data"].encode("utf-8")
                    )

                elif msg.get("type") == "resize":

                    cols = msg.get("cols", 80)
                    rows = msg.get("rows", 24)

                    fcntl.ioctl(
                        master_fd,
                        termios.TIOCSWINSZ,
                        struct.pack(
                            "HHHH",
                            rows,
                            cols,
                            0,
                            0
                        )
                    )

            except (json.JSONDecodeError, OSError):
                break

    except (WebSocketDisconnect, Exception):
        pass

    finally:

        read_task.cancel()

        try:
            proc.terminate()

            await asyncio.wait_for(
                proc.wait(),
                timeout=2
            )

        except Exception:
            proc.kill()

        try:
            os.close(master_fd)

        except Exception:
            pass