import { serve } from "bun";
import type { ServerWebSocket, TLSOptions } from "bun";
import { readFileSync } from "fs";

/**
 * Bun-powered WebSocket proxy that mirrors the Athena HTTP gateway.
 * Every WebSocket message is treated like an HTTP call to the Rust API,
 * allowing the full gateway surface area to ride on a single connection.
 * Supports both WS and WSS traffic. For WSS, either terminate TLS here,
 * or use a TLS-terminating reverse proxy in front (see below). This code
 * explicitly checks for TLS config if you want to serve WSS directly.
 */

const DEFAULT_API_BASE = process.env.ATHENA_API_BASE ?? "http://localhost:4052";
const DEFAULT_PORT = Number(
  process.env.PORT ?? process.env.ATHENA_WSS_PORT ?? 7070,
);
const DEFAULT_API_TIMEOUT_MS = Number(
  process.env.ATHENA_API_TIMEOUT_MS ?? 15_000,
);

// Optionally read TLS key/cert for WSS directly
const TLS_KEY_PATH = process.env.ATHENA_TLS_KEY;
const TLS_CERT_PATH = process.env.ATHENA_TLS_CERT;
let tls: { key: string; cert: string } | undefined = undefined;
if (TLS_KEY_PATH && TLS_CERT_PATH) {
  tls = {
    key: readFileSync(TLS_KEY_PATH, "utf8"),
    cert: readFileSync(TLS_CERT_PATH, "utf8"),
  };
}

const actionDefaults: Record<string, { path: string; method: string }> = {
  "gateway.fetch": { path: "/gateway/fetch", method: "POST" },
  "gateway.insert": { path: "/gateway/insert", method: "POST" },
  "gateway.update": { path: "/gateway/update", method: "POST" },
  "gateway.delete": { path: "/gateway/delete", method: "POST" },
  "gateway.query": { path: "/gateway/query", method: "POST" },
};

type JSONObject = Record<string, unknown>;

type GatewayWsRequest = {
  action: string;
  request_id?: string;
  client_name: string;
  path: string;
  method?: string;
  headers?: JSONObject;
  query?: JSONObject;
  body?: unknown;
};

type GatewayWsResponse = {
  request_id?: string;
  status: "success" | "error";
  message: string;
  data?: unknown;
};

function normalizePath(path: string): string {
  if (path.startsWith("http://") || path.startsWith("https://")) return path;
  if (!path.startsWith("/")) return `/${path}`;
  return path;
}

function buildUrl(base: string, path: string, query?: JSONObject): URL {
  const normalized = normalizePath(path);
  const url = normalized.startsWith("http")
    ? new URL(normalized)
    : new URL(normalized, base.endsWith("/") ? base : `${base}/`);

  if (query && typeof query === "object") {
    for (const [key, value] of Object.entries(query)) {
      if (value === undefined || value === null) continue;
      if (Array.isArray(value)) {
        for (const item of value) url.searchParams.append(key, String(item));
      } else if (typeof value === "object") {
        url.searchParams.append(key, JSON.stringify(value));
      } else {
        url.searchParams.append(key, String(value));
      }
    }
  }

  return url;
}

function serializeHeaders(
  headers: JSONObject | undefined,
  clientName: string,
): Headers {
  const result = new Headers();
  result.set("content-type", "application/json");
  result.set("x-athena-client", clientName);

  if (headers && typeof headers === "object") {
    for (const [key, value] of Object.entries(headers)) {
      if (value === undefined || value === null) continue;
      result.set(key.toLowerCase(), String(value));
    }
  }

  return result;
}

function tryParseMessage(raw: unknown): GatewayWsRequest {
  const text =
    typeof raw === "string"
      ? raw
      : raw instanceof ArrayBuffer
        ? new TextDecoder().decode(raw)
        : raw instanceof Uint8Array
          ? new TextDecoder().decode(raw)
          : (() => {
              throw new Error(
                "Unsupported payload type; expected string or ArrayBuffer",
              );
            })();

  const parsed = JSON.parse(text);
  if (!parsed || typeof parsed !== "object") {
    throw new Error("Invalid message: expected JSON object");
  }

  const { action, client_name, path } = parsed as Record<string, unknown>;
  if (!action || typeof action !== "string") throw new Error("Missing action");
  if (!client_name || typeof client_name !== "string")
    throw new Error("Missing client_name");
  if (!path || typeof path !== "string") throw new Error("Missing path");

  return parsed as GatewayWsRequest;
}

async function proxyToAthena(
  req: GatewayWsRequest,
  apiBase: string,
  apiTimeoutMs: number,
) {
  const defaultShape = actionDefaults[req.action] ?? {
    path: req.path,
    method: req.method ?? "POST",
  };
  const method = (req.method ?? defaultShape.method ?? "POST").toUpperCase();
  const targetUrl = buildUrl(apiBase, req.path || defaultShape.path, req.query);
  const headers = serializeHeaders(req.headers, req.client_name);
  const controller = new AbortController();
  const timeout = setTimeout(() => controller.abort(), apiTimeoutMs);

  const body =
    method === "GET" || method === "HEAD"
      ? undefined
      : JSON.stringify(req.body ?? {});

  const response = await fetch(targetUrl, {
    method,
    headers,
    body,
    signal: controller.signal,
  }).finally(() => clearTimeout(timeout));

  const rawText = await response.text();
  let parsedBody: unknown = rawText;
  if (rawText.length) {
    try {
      parsedBody = JSON.parse(rawText);
    } catch (_) {
      parsedBody = rawText;
    }
  } else {
    parsedBody = null;
  }

  const payload: GatewayWsResponse = {
    request_id: req.request_id,
    status: response.ok ? "success" : "error",
    message: response.ok
      ? `proxied ${method} ${targetUrl.pathname}`
      : `upstream responded with ${response.status}`,
    data: {
      upstream: {
        status: response.status,
        status_text: response.statusText,
      },
      body: parsedBody,
    },
  };

  return payload;
}

function send(ws: ServerWebSocket<unknown>, payload: GatewayWsResponse) {
  try {
    ws.send(JSON.stringify(payload));
  } catch (error) {
    console.error("Failed to send WS response", error);
    ws.close(1011, "send_failed");
  }
}

export type GatewayServerOptions = {
  port?: number;
  apiBase?: string;
  apiTimeoutMs?: number;
  tls?: {
    key: string;
    cert: string;
  };
};

export function startGatewayServer(opts: GatewayServerOptions = {}) {
  const port = opts.port ?? DEFAULT_PORT;
  const apiBase = opts.apiBase ?? DEFAULT_API_BASE;
  const apiTimeoutMs = opts.apiTimeoutMs ?? DEFAULT_API_TIMEOUT_MS;

  // Brief WSS explanation for users
  if (!!opts.tls || tls) {
    console.log(
      "INFO: Starting Athena WSS Gateway with built-in TLS (WSS endpoint). " +
        "Ensure your cert/key is valid. If running behind a reverse proxy like nginx, do not set TLS here.",
    );
  } else {
    console.log(
      "INFO: Starting Athena WS Gateway without TLS. To serve WSS directly, provide ATHENA_TLS_CERT and ATHENA_TLS_KEY env vars, " +
        "or pass 'tls' in options. Otherwise, use a TLS reverse proxy in front of this server for WSS.",
    );
  }

  const finalTls = opts.tls || tls;
  const tlsOptions: TLSOptions | undefined = finalTls
    ? {
        key: finalTls.key,
        cert: finalTls.cert,
      }
    : undefined;

  const server = serve<{ startedAt: number }>({
    port,
    ...(tlsOptions ? { tls: tlsOptions } : {}),
    idleTimeout: 0, // keep connections alive; upstream timeout guards requests
    fetch(req, server) {
      const url = new URL(req.url);

      if (url.pathname === "/wss/info") {
        return Response.json({
          status: "success",
          message: "Athena WebSocket gateway",
          data: {
            transport: finalTls ? "wss" : "ws",
            path: "/wss/gateway",
            api_base: apiBase,
            actions: Object.keys(actionDefaults),
          },
        });
      }

      if (url.pathname === "/health" || url.pathname === "/ready") {
        return Response.json({ status: "ok", port, api_base: apiBase });
      }

      if (url.pathname === "/wss/gateway") {
        if (server.upgrade(req, { data: { startedAt: Date.now() } })) {
          return undefined;
        }
        return new Response("Expected WebSocket upgrade", { status: 426 });
      }

      return new Response("Not Found", { status: 404 });
    },
    websocket: {
      open(ws) {
        send(ws, {
          status: "success",
          message: "connected",
          data: {
            started_at: ws.data.startedAt,
            api_base: apiBase,
            protocol: finalTls ? "wss" : "ws",
          },
        });
      },
      async message(ws, raw) {
        let parsed: GatewayWsRequest;
        try {
          parsed = tryParseMessage(raw);
        } catch (error) {
          send(ws, {
            status: "error",
            message: error instanceof Error ? error.message : "Invalid payload",
          });
          return;
        }

        try {
          const payload = await proxyToAthena(parsed, apiBase, apiTimeoutMs);
          send(ws, payload);
        } catch (error) {
          const isAbort =
            error instanceof DOMException && error.name === "AbortError";
          send(ws, {
            request_id: parsed.request_id,
            status: "error",
            message: isAbort
              ? "Upstream request timed out"
              : "Upstream request failed",
            data: {
              error: error instanceof Error ? error.message : String(error),
            },
          });
        }
      },
      close(ws, code, reason) {
        if (code !== 1000) {
          console.warn(`WebSocket closed abnormally ${code}: ${reason}`);
        }
      },
    },
  });

  // Bun's Server exposes `tls` when configured; check dynamically.
  if ((server as any).tls) {
    console.log(
      `Athena WSS gateway listening on wss://0.0.0.0:${server.port}/wss/gateway -> ${apiBase}`,
    );
  } else {
    console.log(
      `Athena WS gateway listening on ws://0.0.0.0:${server.port}/wss/gateway -> ${apiBase}`,
    );
  }

  return server;
}

const mainServer = import.meta.main ? startGatewayServer() : null;
export default mainServer;
