import asyncio
import os
import signal
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent

try:
    import pty
    HAS_PTY = True
except ImportError:
    HAS_PTY = False

import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI(title="Kanjin")

STATIC_DIR = Path(__file__).resolve().parent / "static"
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/")
async def index():
    html = (STATIC_DIR / "index.html").read_text(encoding="utf-8")
    return HTMLResponse(html)


@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    if not HAS_PTY:
        await ws.accept()
        await ws.send_text("Server must run on Linux (or WSL) with PTY support.\n")
        await ws.close()
        return

    await ws.accept()

    master_fd, slave_fd = pty.openpty()
    loop = asyncio.get_event_loop()

    proc = await asyncio.create_subprocess_exec(
        sys.executable, str(PROJECT_ROOT / "main.py"),
        stdin=slave_fd,
        stdout=slave_fd,
        stderr=slave_fd,
        cwd=str(PROJECT_ROOT),
        close_fds=True,
        preexec_fn=os.setsid,
    )
    os.close(slave_fd)

    async def read_pty():
        try:
            while True:
                data = await loop.run_in_executor(None, os.read, master_fd, 4096)
                if not data:
                    break
                await ws.send_text(data.decode(errors="replace"))
        except (ConnectionError, WebSocketDisconnect):
            pass
        finally:
            try:
                os.close(master_fd)
            except OSError:
                pass
            if proc.returncode is None:
                try:
                    os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
                    await asyncio.wait_for(proc.wait(), timeout=5)
                except (ProcessLookupError, asyncio.TimeoutError):
                    proc.kill()
                    await proc.wait()

    async def write_pty():
        try:
            async for message in ws.iter_text():
                os.write(master_fd, message.encode())
        except (ConnectionError, WebSocketDisconnect):
            pass

    await asyncio.gather(read_pty(), write_pty())


if __name__ == "__main__":
    uvicorn.run("server:app", host="0.0.0.0", port=8080, reload=True)