import { createHmac, timingSafeEqual } from "node:crypto";

export const MODEL_GATEWAY_AUDIENCE = "awp-model-gateway" as const;
export const MODEL_GATEWAY_AGENT = "codex" as const;

export interface ModelGatewayCapabilityClaims {
  readonly v: 1;
  readonly aud: typeof MODEL_GATEWAY_AUDIENCE;
  readonly agent: typeof MODEL_GATEWAY_AGENT;
  readonly attemptId: string;
  readonly accountId: string;
  readonly iat: number;
  readonly exp: number;
}

export interface ModelGatewayCapabilityInput {
  readonly attemptId: string;
  readonly accountId: string;
  readonly issuedAtSeconds?: number;
  readonly ttlSeconds?: number;
}

const DEFAULT_TTL_SECONDS = 4 * 60 * 60;
const MAX_TTL_SECONDS = 6 * 60 * 60;

function secretBuffer(secret: string): Buffer {
  const trimmed = secret.trim();
  if (trimmed.length < 32) throw new Error("Model gateway signing secret is too short");
  return Buffer.from(trimmed, "utf8");
}

function encode(value: Buffer | string): string {
  return Buffer.from(value).toString("base64url");
}

function signature(secret: string, encodedPayload: string): string {
  return createHmac("sha256", secretBuffer(secret)).update(encodedPayload).digest("base64url");
}

export function issueModelGatewayCapability(
  secret: string,
  input: ModelGatewayCapabilityInput,
): string {
  const attemptId = input.attemptId.trim();
  const accountId = input.accountId.trim();
  if (!attemptId || !accountId)
    throw new Error("Model gateway capability requires Attempt and account identity");
  const iat = input.issuedAtSeconds ?? Math.floor(Date.now() / 1000);
  const ttl = input.ttlSeconds ?? DEFAULT_TTL_SECONDS;
  if (!Number.isInteger(iat) || !Number.isInteger(ttl) || ttl < 60 || ttl > MAX_TTL_SECONDS) {
    throw new Error("Model gateway capability TTL must be between 60 seconds and 6 hours");
  }
  const claims: ModelGatewayCapabilityClaims = {
    v: 1,
    aud: MODEL_GATEWAY_AUDIENCE,
    agent: MODEL_GATEWAY_AGENT,
    attemptId,
    accountId,
    iat,
    exp: iat + ttl,
  };
  const payload = encode(JSON.stringify(claims));
  return `${payload}.${signature(secret, payload)}`;
}

export function verifyModelGatewayCapability(
  secret: string,
  token: string,
  nowSeconds = Math.floor(Date.now() / 1000),
): ModelGatewayCapabilityClaims {
  const [payload, suppliedSignature, extra] = token.trim().split(".");
  if (!payload || !suppliedSignature || extra !== undefined) {
    throw new Error("Invalid model gateway capability format");
  }
  const expected = signature(secret, payload);
  const expectedBuffer = Buffer.from(expected);
  const suppliedBuffer = Buffer.from(suppliedSignature);
  if (
    expectedBuffer.length !== suppliedBuffer.length ||
    !timingSafeEqual(expectedBuffer, suppliedBuffer)
  ) {
    throw new Error("Invalid model gateway capability signature");
  }
  let parsed: unknown;
  try {
    parsed = JSON.parse(Buffer.from(payload, "base64url").toString("utf8"));
  } catch {
    throw new Error("Invalid model gateway capability payload");
  }
  if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
    throw new Error("Invalid model gateway capability claims");
  }
  const claims = parsed as Partial<ModelGatewayCapabilityClaims>;
  if (
    claims.v !== 1 ||
    claims.aud !== MODEL_GATEWAY_AUDIENCE ||
    claims.agent !== MODEL_GATEWAY_AGENT ||
    typeof claims.attemptId !== "string" ||
    !claims.attemptId.trim() ||
    typeof claims.accountId !== "string" ||
    !claims.accountId.trim() ||
    typeof claims.iat !== "number" ||
    !Number.isInteger(claims.iat) ||
    typeof claims.exp !== "number" ||
    !Number.isInteger(claims.exp) ||
    claims.exp <= claims.iat ||
    claims.exp - claims.iat > MAX_TTL_SECONDS ||
    nowSeconds < claims.iat - 60 ||
    nowSeconds >= claims.exp
  ) {
    throw new Error("Expired or invalid model gateway capability claims");
  }
  return claims as ModelGatewayCapabilityClaims;
}