from fastapi import FastAPI, WebSocket
from fastapi.responses import PlainTextResponse
import asyncio
import os
import pty
import subprocess

app = FastAPI()


@app.get("/", response_class=PlainTextResponse)
async def root():
    return "Bash WebSocket Terminal OK"


@app.get("/health", response_class=PlainTextResponse)
async def health():
    return "OK"


@app.websocket("/terminals/websocket/{terminal_id}")
async def terminal(websocket: WebSocket, terminal_id: str):

    await websocket.accept()

    print("CONNECTED:", terminal_id, flush=True)

    master, slave = pty.openpty()

    process = subprocess.Popen(
        ["/bin/bash", "--noprofile", "--norc", "-i"],
        stdin=slave,
        stdout=slave,
        stderr=slave,
        start_new_session=True
    )

    os.close(slave)

    loop = asyncio.get_running_loop()

    async def bash_to_websocket():

        while process.poll() is None:

            try:

                data = await loop.run_in_executor(
                    None,
                    os.read,
                    master,
                    8192
                )

                if not data:
                    break

                output = data.decode(
                    "utf-8",
                    errors="replace"
                )

                print(
                    "BASH OUTPUT:",
                    repr(output),
                    flush=True
                )

                await websocket.send_text(output)

            except Exception as e:

                print(
                    "OUTPUT ERROR:",
                    repr(e),
                    flush=True
                )

                break

    output_task = asyncio.create_task(
        bash_to_websocket()
    )

    try:

        while True:

            data = await websocket.receive_text()

            print(
                "TERMUX INPUT:",
                repr(data),
                flush=True
            )

            os.write(
                master,
                data.encode("utf-8")
            )

    except Exception as e:

        print(
            "WEBSOCKET:",
            repr(e),
            flush=True
        )

    finally:

        output_task.cancel()

        try:
            process.terminate()
        except:
            pass

        try:
            os.close(master)
        except:
            pass

        print(
            "DISCONNECTED:",
            terminal_id,
            flush=True
        )


if __name__ == "__main__":

    import uvicorn

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860
    )