import type { Server } from "bun";
import type {
  CoreHostSshCloseMessage,
  CoreHostSshEofMessage,
  CoreHostSshInputMessage,
  CoreHostSshOpenMessage,
  CoreHostSshResizeMessage,
  ProviderHostSshDataMessage,
  ProviderHostSshErrorMessage,
  ProviderHostSshExitMessage,
  ProviderHostSshOpenedMessage,
} from "../../shared/src/index";
import { ctx, sql, config, logger } from "./context";
import type { ProviderSocket, WsData } from "./types";
import { jsonResponse, wsSendJson } from "./http";

function safeSessionId(): string {
  return `ssh_${Date.now()}_${Math.random().toString(16).slice(2)}`;
}

function sshClientFor(sessionId: string): ProviderSocket | null {
  return ctx.activeSshClients.get(sessionId) ?? null;
}

function providerForSsh(providerId: string): ProviderSocket | null {
  return ctx.activeProviders.get(providerId) ?? null;
}

export function closeSshClient(sessionId: string, code = 1000, reason = "ssh session closed"): void {
  const client = sshClientFor(sessionId);
  ctx.activeSshClients.delete(sessionId);
  if (client !== null && client.readyState === WebSocket.OPEN) client.close(code, reason);
}

export function forwardSshProviderMessage(
  message: ProviderHostSshOpenedMessage | ProviderHostSshDataMessage | ProviderHostSshExitMessage | ProviderHostSshErrorMessage,
): void {
  const client = sshClientFor(message.sessionId);
  if (client === null) {
    logger("warn", "ssh_client_missing", { providerId: message.providerId, sessionId: message.sessionId, type: message.type });
    return;
  }
  if (message.type === "host_ssh_opened") {
    wsSendJson(client, { type: "ssh.opened", providerId: message.providerId, sessionId: message.sessionId });
    return;
  }
  if (message.type === "host_ssh_data") {
    wsSendJson(client, { type: "ssh.data", stream: message.stream, data: message.data, encoding: message.encoding });
    return;
  }
  if (message.type === "host_ssh_exit") {
    wsSendJson(client, { type: "ssh.exit", exitCode: message.exitCode, signal: message.signal });
    setTimeout(() => closeSshClient(message.sessionId), 50);
    return;
  }
  wsSendJson(client, { type: "ssh.error", message: message.message });
  setTimeout(() => closeSshClient(message.sessionId, 1011, "ssh session error"), 50);
}

function numberFromUnknown(value: unknown, fallback: number, min: number, max: number): number {
  const parsed = typeof value === "number" ? value : typeof value === "string" ? Number(value) : fallback;
  if (!Number.isFinite(parsed)) return fallback;
  return Math.max(min, Math.min(max, Math.floor(parsed)));
}

async function providerSupports(providerId: string, capability: string): Promise<boolean> {
  if (!ctx.dbReady) return false;
  const rows = await sql()<Array<{ labels: unknown }>>`
    SELECT labels
    FROM unidesk_nodes
    WHERE provider_id = ${providerId}
    LIMIT 1
  `;
  const labels = rows[0]?.labels;
  if (typeof labels !== "object" || labels === null || Array.isArray(labels)) return false;
  const capabilities = (labels as Record<string, unknown>).unideskCapabilities;
  return Array.isArray(capabilities) ? capabilities.filter((item): item is string => typeof item === "string").includes(capability) : false;
}

export async function handleSshClientMessage(ws: ProviderSocket, raw: string | Buffer): Promise<void> {
  const text = typeof raw === "string" ? raw : raw.toString("utf8");
  const message = JSON.parse(text) as { type?: unknown; providerId?: unknown; cwd?: unknown; command?: unknown; tty?: unknown; data?: unknown; encoding?: unknown; cols?: unknown; rows?: unknown };
  if (message.type === "ssh.open") {
    const providerId = typeof message.providerId === "string" ? message.providerId : "";
    if (providerId.length === 0) {
      wsSendJson(ws, { type: "ssh.error", message: "providerId is required" });
      return;
    }
    const provider = providerForSsh(providerId);
    if (provider === null) {
      wsSendJson(ws, { type: "ssh.error", message: `provider is not online: ${providerId}` });
      return;
    }
    if (!(await providerSupports(providerId, "host.ssh"))) {
      wsSendJson(ws, { type: "ssh.error", message: `provider does not declare host.ssh capability: ${providerId}` });
      return;
    }
    const sessionId = safeSessionId();
    ws.data.channel = "ssh";
    ws.data.providerId = providerId;
    ws.data.sessionId = sessionId;
    ctx.activeSshClients.set(sessionId, ws);
    const openMessage: CoreHostSshOpenMessage = {
      type: "host_ssh_open",
      sessionId,
      cols: numberFromUnknown(message.cols, 100, 20, 300),
      rows: numberFromUnknown(message.rows, 30, 8, 120),
    };
    if (typeof message.cwd === "string" && message.cwd.length > 0) openMessage.cwd = message.cwd;
    if (typeof message.command === "string" && message.command.length > 0) openMessage.command = message.command;
    if (typeof message.tty === "boolean") openMessage.tty = message.tty;
    provider.send(JSON.stringify(openMessage));
    wsSendJson(ws, { type: "ssh.dispatched", providerId, sessionId });
    logger("info", "ssh_session_dispatched", { providerId, sessionId, hasCommand: typeof openMessage.command === "string" });
    return;
  }

  const sessionId = ws.data.sessionId;
  const providerId = ws.data.providerId;
  if (sessionId === undefined || providerId === undefined) {
    if (message.type === "ssh.eof" || message.type === "ssh.close") return;
    wsSendJson(ws, { type: "ssh.error", message: "ssh session is not open" });
    return;
  }
  const provider = providerForSsh(providerId);
  if (provider === null) {
    wsSendJson(ws, { type: "ssh.error", message: `provider went offline: ${providerId}` });
    return;
  }

  if (message.type === "ssh.input") {
    if (typeof message.data !== "string" || message.encoding !== "base64") {
      wsSendJson(ws, { type: "ssh.error", message: "ssh.input requires base64 data" });
      return;
    }
    const inputMessage: CoreHostSshInputMessage = { type: "host_ssh_input", sessionId, data: message.data, encoding: "base64" };
    provider.send(JSON.stringify(inputMessage));
    return;
  }

  if (message.type === "ssh.resize") {
    const resizeMessage: CoreHostSshResizeMessage = {
      type: "host_ssh_resize",
      sessionId,
      cols: numberFromUnknown(message.cols, 100, 20, 300),
      rows: numberFromUnknown(message.rows, 30, 8, 120),
    };
    provider.send(JSON.stringify(resizeMessage));
    return;
  }

  if (message.type === "ssh.eof") {
    const eofMessage: CoreHostSshEofMessage = { type: "host_ssh_eof", sessionId };
    provider.send(JSON.stringify(eofMessage));
    return;
  }

  if (message.type === "ssh.close") {
    const closeMessage: CoreHostSshCloseMessage = { type: "host_ssh_close", sessionId };
    provider.send(JSON.stringify(closeMessage));
    closeSshClient(sessionId);
    return;
  }

  wsSendJson(ws, { type: "ssh.error", message: `unsupported ssh client message: ${String(message.type)}` });
}

export async function sshRoute(req: Request, server: Server<WsData>): Promise<Response | undefined> {
  const url = new URL(req.url);
  const token = url.searchParams.get("token") ?? req.headers.get("x-provider-token");
  if (token !== config().providerToken) return jsonResponse({ ok: false, error: "invalid ssh bridge token" }, 401);
  const upgraded = server.upgrade(req, { data: { channel: "ssh" } satisfies WsData });
  return upgraded ? undefined : jsonResponse({ ok: false, error: "websocket upgrade failed" }, 400);
}
