"""
HF Terminal — Real Bash shell in the browser via WebSocket + PTY.

Architecture: Tornado serves static files and a /ws endpoint.
Each WebSocket connection spawns an isolated PTY process so
nano, vim, htop and colour output all work correctly.

Why Tornado instead of Gradio:
  1. Gradio's WebSocket layer is built for component RPC, not raw byte streams.
  2. Tornado gives native HTTP + WebSocket in one process with zero overhead.
  3. HF Spaces supports any Python app listening on port 7860 — no Gradio needed.
"""

import base64
import fcntl
import json
import os
import pty
import select
import signal
import struct
import termios

import tornado.ioloop
import tornado.web
import tornado.websocket

# ── Configuration ────────────────────────────────────────────────
SHELL = os.environ.get("SHELL", "/bin/bash")
PORT = int(os.environ.get("PORT", 7860))
TERM = os.environ.get("TERM", "xterm-256color")
COLUMNS = int(os.environ.get("COLUMNS", 80))
ROWS = int(os.environ.get("ROWS", 24))
PERSIST_DIR = os.environ.get("PERSIST_DIR", "/data")

# ── Ensure persistent storage symlink ────────────────────────────
if os.path.isdir(PERSIST_DIR) and not os.path.ismount(PERSIST_DIR):
    os.makedirs(PERSIST_DIR, exist_ok=True)


# ── PTY Process Manager ─────────────────────────────────────────
class PTYProcess:
    """Wraps a single PTY fork and provides async read/write."""

    def __init__(self, cols=COLUMNS, rows=ROWS):
        self.cols = cols
        self.rows = rows
        self.pid = None
        self.fd = None

    def spawn(self):
        pid, fd = pty.fork()
        if pid == 0:
            # ── Child process ──
            env = os.environ.copy()
            env["TERM"] = TERM
            env["COLUMNS"] = str(self.cols)
            env["LINES"] = str(self.rows)
            # Persist bash history to /data if available
            if os.path.isdir(PERSIST_DIR):
                histfile = os.path.join(PERSIST_DIR, ".bash_history")
                env["HISTFILE"] = histfile
                env["PROMPT_COMMAND"] = "history -a; history -c; history -r"
            args = [SHELL, "--login"] if "bash" in SHELL else [SHELL]
            os.execvpe(SHELL, args, env)
        # ── Parent process ──
        self.pid = pid
        self.fd = fd
        self._set_nonblocking(self.fd)
        self.resize(self.rows, self.cols)

    @staticmethod
    def _set_nonblocking(fd):
        flags = fcntl.fcntl(fd, fcntl.F_GETFL)
        fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)

    def resize(self, rows, cols):
        if self.fd is None:
            return
        self.rows = rows
        self.cols = cols
        winsize = struct.pack("HHHH", rows, cols, 0, 0)
        fcntl.ioctl(self.fd, termios.TIOCSWINSZ, winsize)

    def write(self, data):
        if self.fd is None:
            return
        os.write(self.fd, data)

    def read(self):
        """Non-blocking read. Returns bytes or None if nothing available."""
        if self.fd is None:
            return None
        try:
            ready, _, _ = select.select([self.fd], [], [], 0)
            if ready:
                return os.read(self.fd, 65536)
        except (OSError, IOError):
            return None
        return None

    def is_alive(self):
        if self.pid is None:
            return False
        try:
            pid, status = os.waitpid(self.pid, os.WNOHANG)
            return pid == 0
        except ChildProcessError:
            return False

    def kill(self):
        if self.pid is not None:
            try:
                os.kill(self.pid, signal.SIGHUP)
            except ProcessLookupError:
                pass
        if self.fd is not None:
            try:
                os.close(self.fd)
            except OSError:
                pass
            self.fd = None


# ── WebSocket Handler ────────────────────────────────────────────
class TerminalWebSocket(tornado.websocket.WebSocketHandler):

    def initialize(self):
        self.pty = None
        self._read_task = None

    def open(self):
        self.pty = PTYProcess()
        self.pty.spawn()
        self._read_task = tornado.ioloop.PeriodicCallback(self._poll_pty, 16)
        self._read_task.start()

    def on_message(self, message):
        if self.pty is None:
            return
        try:
            msg = json.loads(message)
        except json.JSONDecodeError:
            return

        kind = msg.get("type")

        if kind == "input":
            data = msg.get("data", "")
            try:
                raw = base64.b64decode(data)
                self.pty.write(raw)
            except Exception:
                self.pty.write(data.encode("utf-8"))

        elif kind == "resize":
            rows = int(msg.get("rows", 24))
            cols = int(msg.get("cols", 80))
            self.pty.resize(rows, cols)

    def on_close(self):
        if self._read_task:
            self._read_task.stop()
            self._read_task = None
        if self.pty:
            self.pty.kill()
            self.pty = None

    def _poll_pty(self):
        if self.pty is None:
            return
        data = self.pty.read()
        if data:
            payload = base64.b64encode(data).decode("ascii")
            self.write_message(json.dumps({"type": "output", "data": payload}))
        if not self.pty.is_alive():
            self.write_message(json.dumps({"type": "exit", "data": "Process exited."}))
            if self._read_task:
                self._read_task.stop()
                self._read_task = None

    def check_origin(self, origin):
        return True


# ── Application Factory ──────────────────────────────────────────
BASE_DIR = os.path.dirname(os.path.abspath(__file__))


class IndexHandler(tornado.web.RequestHandler):
    """Serves terminal.html from the project root."""
    def get(self):
        path = os.path.join(BASE_DIR, "terminal.html")
        with open(path, "r", encoding="utf-8") as f:
            self.set_header("Content-Type", "text/html; charset=utf-8")
            self.write(f.read())


def make_app():
    static_path = os.path.join(BASE_DIR, "static")
    return tornado.web.Application(
        [
            (r"/", IndexHandler),
            (r"/ws", TerminalWebSocket),
            (r"/static/(.*)", tornado.web.StaticFileHandler, {"path": static_path}),
        ],
        debug=False,
    )


# ── Entry Point ──────────────────────────────────────────────────
if __name__ == "__main__":
    app = make_app()
    app.listen(PORT, address="0.0.0.0")
    print(f"hf-terminal listening on 0.0.0.0:{PORT}", flush=True)
    tornado.ioloop.IOLoop.current().start()