import { readFile } from "node:fs/promises";
import {
  createServer,
  request as httpRequest,
  type IncomingMessage,
  type ServerResponse,
} from "node:http";
import { connect as netConnect } from "node:net";
import { verifyModelGatewayCapability, type ModelGatewayCapabilityClaims } from "@awp/contracts";

export interface ModelGatewayOptions {
  readonly subrouterBaseUrl: string;
  readonly subrouterProxyToken: string;
  readonly signingSecret: string;
  readonly host?: string;
  readonly port?: number;
}

const BLOCKED_HEADERS = new Set([
  "authorization",
  "proxy-authorization",
  "x-api-key",
  "x-subrouter-account-id",
  "x-subrouter-account",
  "x-subrouter-agent",
  "x-subrouter-session",
  "host",
]);

function bearer(request: IncomingMessage): string {
  const value = request.headers.authorization;
  if (typeof value !== "string" || !value.startsWith("Bearer ")) {
    throw Object.assign(new Error("Missing model gateway capability"), { statusCode: 401 });
  }
  return value.slice("Bearer ".length).trim();
}

function claimsFor(request: IncomingMessage, secret: string): ModelGatewayCapabilityClaims {
  try {
    return verifyModelGatewayCapability(secret, bearer(request));
  } catch {
    throw Object.assign(new Error("Invalid or expired model gateway capability"), {
      statusCode: 401,
    });
  }
}

function allowedPath(request: IncomingMessage): string {
  const path = request.url ?? "/";
  const pathname = new URL(path, "http://awp.invalid").pathname;
  if (pathname !== "/v1" && !pathname.startsWith("/v1/")) {
    throw Object.assign(new Error("Model gateway only proxies the OpenAI-compatible /v1 surface"), {
      statusCode: 404,
    });
  }
  return path;
}

function upstreamHeaders(
  request: IncomingMessage,
  target: URL,
  claims: ModelGatewayCapabilityClaims,
  proxyToken: string,
): Record<string, string | string[]> {
  const headers: Record<string, string | string[]> = {};
  for (const [name, value] of Object.entries(request.headers)) {
    if (value === undefined || BLOCKED_HEADERS.has(name.toLowerCase())) continue;
    headers[name] = value;
  }
  headers.host = target.host;
  headers.authorization = `Bearer ${proxyToken}`;
  headers["x-subrouter-agent"] = claims.agent;
  headers["x-subrouter-account-id"] = claims.accountId;
  headers["x-subrouter-session"] = claims.attemptId;
  return headers;
}

function statusCode(error: unknown): number {
  if (typeof error === "object" && error !== null && "statusCode" in error) {
    const value = Number((error as { statusCode?: unknown }).statusCode);
    if (Number.isInteger(value) && value >= 400 && value <= 599) return value;
  }
  return 502;
}

function safeError(response: ServerResponse, error: unknown): void {
  const status = statusCode(error);
  if (response.headersSent) {
    response.destroy();
    return;
  }
  response.writeHead(status, { "content-type": "application/json; charset=utf-8" });
  response.end(
    JSON.stringify({
      error: status === 401 ? "unauthorized" : status === 404 ? "not-found" : "gateway-failure",
    }),
  );
}

export function createModelGateway(options: ModelGatewayOptions) {
  const upstream = new URL(options.subrouterBaseUrl);
  if (upstream.protocol !== "http:") {
    throw new Error("The I1 K3s model gateway requires an internal HTTP Subrouter target");
  }
  const proxyToken = options.subrouterProxyToken.trim();
  const signingSecret = options.signingSecret.trim();
  if (!proxyToken) throw new Error("Subrouter proxy token is empty");
  if (!signingSecret) throw new Error("Model gateway signing secret is empty");

  const server = createServer((request, response) => {
    try {
      const claims = claimsFor(request, signingSecret);
      const path = allowedPath(request);
      const target = new URL(path, upstream);
      const proxy = httpRequest(
        target,
        {
          method: request.method,
          headers: upstreamHeaders(request, target, claims, proxyToken),
        },
        (upstreamResponse) => {
          response.writeHead(upstreamResponse.statusCode ?? 502, upstreamResponse.headers);
          upstreamResponse.pipe(response);
        },
      );
      proxy.on("error", (error) => safeError(response, error));
      request.pipe(proxy);
    } catch (error) {
      safeError(response, error);
    }
  });

  server.on("upgrade", (request, socket, head) => {
    try {
      const claims = claimsFor(request, signingSecret);
      const path = allowedPath(request);
      const target = new URL(path, upstream);
      const port = Number(target.port || 80);
      const headers = upstreamHeaders(request, target, claims, proxyToken);
      const upstreamSocket = netConnect(port, target.hostname);
      upstreamSocket.once("connect", () => {
        const lines = [
          `${request.method ?? "GET"} ${target.pathname}${target.search} HTTP/${request.httpVersion}`,
        ];
        for (const [name, value] of Object.entries(headers)) {
          if (Array.isArray(value)) {
            for (const item of value) lines.push(`${name}: ${item}`);
          } else {
            lines.push(`${name}: ${value}`);
          }
        }
        upstreamSocket.write(`${lines.join("\r\n")}\r\n\r\n`);
        if (head.length > 0) upstreamSocket.write(head);
        socket.pipe(upstreamSocket).pipe(socket);
      });
      upstreamSocket.on("error", () => socket.destroy());
    } catch (error) {
      const status = statusCode(error);
      socket.end(
        `HTTP/1.1 ${status} ${status === 401 ? "Unauthorized" : "Not Found"}\r\nConnection: close\r\n\r\n`,
      );
    }
  });

  return server;
}

async function main(): Promise<void> {
  const subrouterBaseUrl = process.env.AWP_SUBROUTER_URL;
  const proxyTokenFile = process.env.AWP_SUBROUTER_PROXY_TOKEN_FILE;
  const signingSecretFile = process.env.AWP_MODEL_GATEWAY_SIGNING_SECRET_FILE;
  if (!subrouterBaseUrl || !proxyTokenFile || !signingSecretFile) {
    throw new Error(
      "AWP_SUBROUTER_URL, AWP_SUBROUTER_PROXY_TOKEN_FILE and AWP_MODEL_GATEWAY_SIGNING_SECRET_FILE are required",
    );
  }
  const server = createModelGateway({
    subrouterBaseUrl,
    subrouterProxyToken: await readFile(proxyTokenFile, "utf8"),
    signingSecret: await readFile(signingSecretFile, "utf8"),
  });
  const host = process.env.AWP_GATEWAY_HOST ?? "0.0.0.0";
  const port = Number(process.env.AWP_GATEWAY_PORT ?? "32180");
  await new Promise<void>((resolve, reject) => {
    server.once("error", reject);
    server.listen(port, host, resolve);
  });
  process.stdout.write(`AWP model gateway listening on ${host}:${port}\n`);
}

if (process.argv[1] && import.meta.url === new URL(`file://${process.argv[1]}`).href) {
  void main().catch((error) => {
    process.stderr.write(
      `${error instanceof Error ? (error.stack ?? error.message) : String(error)}\n`,
    );
    process.exitCode = 1;
  });
}