import type { Context, Session } from "context";
import { eq, or } from "drizzle-orm";
import schema from "schema";
import { consola } from "consola";
import { Devbox, RunloopSDK } from "@runloop/api-client";
import decrypt from "lib/decrypt";
import { createListener } from "pty/pty-tunnel";
import chalk from "chalk";
import { $ } from "zx";
import crypto from "crypto";
import fs from "fs/promises";

const TERM = "xterm-256color";
const PTY_SERVER_DOWNLOAD_URL =
  "https://github.com/tsirysndr/pty-tunnel-server/releases/download/v0.0.2/pty-server-linux-x86_64.tar.gz";
const SERVER_BIN_NAME = "pty-tunnel-server";
const PTY_PORT = 26661;

type SandboxEnvironmentOptions = {
  id: string;
  apiKey: string;
};

async function checkIfServerInstalled(sandbox: Devbox) {
  const exists = await sandbox.cmd.exec(`command -v ${SERVER_BIN_NAME}`);
  return exists.exitCode === 0;
}

async function setupSandboxEnvironment(
  options: SandboxEnvironmentOptions,
  stdout: (data: string) => void,
  stderr: (data: string) => void,
): Promise<{ sandbox: Devbox }> {
  const sdk = new RunloopSDK({
    bearerToken: options.apiKey,
  });
  consola.info("Runloop: fetching sandbox", chalk.greenBright(options.id));
  const sandbox = sdk.devbox.fromId(options.id);
  consola.info("Runloop: sandbox fetched", chalk.greenBright(options.id));

  consola.info(
    "Runloop: checking pty-tunnel-server",
    chalk.greenBright(options.id),
  );
  if (!(await checkIfServerInstalled(sandbox))) {
    await $`bash -c "type /tmp/${SERVER_BIN_NAME} || curl -L ${PTY_SERVER_DOWNLOAD_URL} | tar xz -C /tmp"`;

    consola.info(
      "Uploading pty-tunnel server binary to sandbox",
      chalk.greenBright(options.id),
    );

    const pathname = `pty-server-${crypto.randomUUID()}`;
    await sandbox.file.upload({
      path: pathname,
      file: new File(
        [await fs.readFile(`/tmp/${SERVER_BIN_NAME}`)],
        SERVER_BIN_NAME,
      ),
    });

    consola.info(
      "Setting up pty-tunnel server binary in sandbox",
      chalk.greenBright(options.id),
    );

    await sandbox.cmd.exec(
      `bash -c mv "${pathname}" /usr/bin/${SERVER_BIN_NAME} || sudo mv "${pathname}" /usr/bin/${SERVER_BIN_NAME}; chmod a+x /usr/bin/${SERVER_BIN_NAME} || sudo chmod a+x /usr/bin/${SERVER_BIN_NAME}`,
    );

    consola.info(
      "Pty-tunnel server binary set up in sandbox",
      chalk.greenBright(options.id),
    );
  }

  consola.info(
    "Starting pty-tunnel server in sandbox",
    chalk.greenBright(options.id),
  );

  await sandbox.cmd.execAsync(
    `bash -c "TERM==${TERM} ${SERVER_BIN_NAME} --port=${PTY_PORT} --mode=client --cols=${process.stdout.columns ?? 80} --rows=${process.stdout.rows ?? 24} bash"`,
    {
      stdout: stdout,
      stderr: stderr,
    },
  );

  consola.info(
    "Runloop: pty-tunnel-server process started",
    chalk.greenBright(options.id),
  );

  return { sandbox };
}

export async function createTerminalSession(
  ctx: Context,
  id: string,
  key = id,
) {
  const [record] = await ctx.db
    .select()
    .from(schema.sandboxes)
    .leftJoin(
      schema.runloopAuth,
      eq(schema.runloopAuth.sandboxId, schema.sandboxes.id),
    )
    .where(or(eq(schema.sandboxes.id, id), eq(schema.sandboxes.sandboxId, id)))
    .execute();

  if (!record?.runloop_auth) {
    consola.error("Runloop auth not found for sandbox", { id });
    throw new Error("Runloop auth not found for sandbox " + id);
  }

  if (!record.sandboxes.sandboxId) {
    consola.error("Sandbox ID not found for sandbox", { id });
    throw new Error("Sandbox ID not found for sandbox " + id);
  }

  const listener = createListener();

  // setup the sandbox environment for pty-tunnel server
  const { sandbox } = await setupSandboxEnvironment(
    {
      id: record.sandboxes.sandboxId,
      apiKey: decrypt(record.runloop_auth.apiKey),
    },
    // Pipe the pty-tunnel-server's stdout into the listener so
    // readConnectionInfo() can parse the JSON connection handshake.
    // We also accept stderr in case the binary writes there instead.
    (data) => {
      consola.debug(`pty-tunnel-server [stdout]:`, data.trimEnd());
      // jsonlines parser requires newline-terminated data
      const chunk = data.endsWith("\n") ? data : data + "\n";
      listener.stdoutStream.write(chunk);
    },
    (data) => consola.debug(`pty-tunnel-server [stderr]:`, data.trimEnd()),
  );

  consola.info("Runloop: fetching sandbox tunnels", chalk.greenBright(id));
  const tunnel = (await sandbox.getInfo()).tunnel;
  console.log(
    `Tunnel URL for port ${PTY_PORT}: https://${PTY_PORT}-${tunnel?.tunnel_key}.tunnel.runloop.ai`,
  );
  consola.info(
    "Runloop: awaiting pty-tunnel connection info",
    chalk.greenBright(id),
  );
  const details = await listener.connection;
  consola.info(
    "Runloop: pty-tunnel connection info received",
    chalk.greenBright(id),
  );

  const url =
    `wss://${PTY_PORT}-${tunnel?.tunnel_key}.tunnel.runloop.ai` as const;
  consola.info("Connecting to WebSocket URL:", url);

  const socket = details.createClient(url);

  const session: Session = {
    socket,
    clients: new Set(),
    wsClients: new Set(),
  };

  socket.addEventListener("message", async ({ data }) => {
    const text = data.toString("utf-8");
    for (const res of session.clients) {
      res.write(`event: output\n`);
      res.write(`data: ${JSON.stringify({ data: text })}\n\n`);
    }
    for (const ws of session.wsClients) {
      if (ws.readyState === ws.OPEN) ws.send(text);
    }
  });

  socket.addEventListener("close", () => {
    ctx.sessions.delete(key);
    for (const ws of session.wsClients) {
      if (ws.readyState === ws.OPEN) ws.close(1000, "exit");
    }
    session.clients.clear();
    session.wsClients.clear();
  });

  consola.info(
    "Runloop: waiting for pty-tunnel socket to open",
    chalk.greenBright(id),
  );
  await socket.waitForOpen();
  consola.info(
    "Runloop: pty-tunnel socket open, sending ready",
    chalk.greenBright(id),
  );
  socket.sendMessage({ type: "ready" });

  ctx.sessions.set(key, session);
  return session;
}