import os
import pty
import asyncio
import struct
import fcntl
import termios
import json
import subprocess
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()

    # Async task: read PTY output and stream to browser
    async def read_pty():
        while True:
            try:
                data = await loop.run_in_executor(None, lambda: os.read(master_fd, 4096))
                if data:
                    await websocket.send_text(data.decode('utf-8', errors='replace'))
            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