import type { FastifyRequest } from "fastify";
import { dirname } from "node:path";
import * as pty from "node-pty";
import type { WebSocket } from "ws";
import type { AppConfig } from "./config.js";
import { CSRF_COOKIE, safeEqual } from "./security.js";

type ClientMessage =
  | { type: "authenticate"; csrfToken: string }
  | { type: "input"; data: string }
  | { type: "resize"; cols: number; rows: number };

function send(socket: WebSocket, value: Record<string, unknown>): void {
  if (socket.readyState === 1) socket.send(JSON.stringify(value));
}

function trustedOrigin(request: FastifyRequest, config: AppConfig): boolean {
  const origin = request.headers.origin;
  if (!origin) return config.nodeEnv === "test";
  try {
    const parsed = new URL(origin);
    return (
      (["http:", "https:"].includes(parsed.protocol) && Boolean(request.headers.host) && parsed.host === request.headers.host) ||
      config.allowedOrigins.has(origin)
    );
  } catch {
    return false;
  }
}

export function openShell(socket: WebSocket, request: FastifyRequest, config: AppConfig): void {
  if (!request.authUser || !trustedOrigin(request, config)) {
    socket.close(1008, "Authentication required");
    return;
  }

  let terminal: pty.IPty | null = null;
  const authenticationTimer = setTimeout(() => socket.close(1008, "Authentication timed out"), 5_000);
  const lifetimeTimer = setTimeout(() => socket.close(1000, "Terminal session expired"), 30 * 60 * 1_000);

  const closeTerminal = () => {
    clearTimeout(authenticationTimer);
    clearTimeout(lifetimeTimer);
    if (terminal) {
      try {
        terminal.kill();
      } catch {
        // The shell may already have exited.
      }
      terminal = null;
    }
  };

  socket.on("close", closeTerminal);
  socket.on("error", closeTerminal);
  socket.on("message", (raw) => {
    const serialized = Array.isArray(raw)
      ? Buffer.concat(raw).toString("utf8")
      : raw instanceof ArrayBuffer
        ? Buffer.from(new Uint8Array(raw)).toString("utf8")
        : raw.toString("utf8");
    if (Buffer.byteLength(serialized) > 16_384) {
      socket.close(1009, "Message too large");
      return;
    }

    let message: ClientMessage;
    try {
      message = JSON.parse(serialized) as ClientMessage;
    } catch {
      socket.close(1003, "Invalid terminal message");
      return;
    }

    if (!terminal) {
      const cookieToken = request.cookies[CSRF_COOKIE];
      if (
        message.type !== "authenticate" ||
        typeof message.csrfToken !== "string" ||
        !cookieToken ||
        !safeEqual(cookieToken, message.csrfToken)
      ) {
        socket.close(1008, "Authentication failed");
        return;
      }
      clearTimeout(authenticationTimer);
      const shellHome = dirname(config.databasePath);
      terminal = pty.spawn("/bin/bash", ["--noprofile", "--norc"], {
        name: "xterm-256color",
        cols: 100,
        rows: 28,
        cwd: shellHome,
        env: {
          HOME: shellHome,
          LANG: process.env.LANG ?? "C.UTF-8",
          PATH: "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
          SHELL: "/bin/bash",
          TERM: "xterm-256color",
          USER: "pi-companion"
        }
      });
      terminal.onData((data) => send(socket, { type: "output", data }));
      terminal.onExit(({ exitCode }) => {
        terminal = null;
        send(socket, { type: "exit", exitCode });
        socket.close(1000, "Shell exited");
      });
      send(socket, { type: "ready" });
      return;
    }

    if (message.type === "input" && typeof message.data === "string" && message.data.length <= 8_192) {
      terminal.write(message.data);
    } else if (
      message.type === "resize" &&
      Number.isInteger(message.cols) &&
      Number.isInteger(message.rows) &&
      message.cols >= 20 && message.cols <= 240 &&
      message.rows >= 5 && message.rows <= 100
    ) {
      terminal.resize(message.cols, message.rows);
    }
  });
}
