const crypto = require("crypto");
const identities = new Map();

function nodeIdentity(identityPurpose, node) {
  const identityKey = `${identityPurpose}:${node}`;
  const existing = identities.get(identityKey);
  if (existing) return existing;
  const { privateKey: privateKeyObject, publicKey } =
    crypto.generateKeyPairSync("ed25519");
  const privateDer = privateKeyObject.export({ format: "der", type: "pkcs8" });
  const publicDer = publicKey.export({
    format: "der",
    type: "spki",
  });
  const privateSeed = Buffer.from(privateDer).subarray(-32);
  const identity = {
    privateKey: `ed25519:${privateSeed.toString("base64")}`,
    publicKey: `ed25519:${Buffer.from(publicDer).subarray(-32).toString("base64")}`,
    privateKeyObject,
  };
  identities.set(identityKey, identity);
  return identity;
}

function nodeIdentityFromPrivateKey(privateKey) {
  if (typeof privateKey !== "string" || !privateKey.startsWith("ed25519:")) {
    throw new Error("node private key must use ed25519:<base64> encoding");
  }
  const seed = Buffer.from(privateKey.slice("ed25519:".length), "base64");
  if (seed.length !== 32) throw new Error("node private key must contain 32 bytes");
  const privateKeyObject = crypto.createPrivateKey({
    key: Buffer.concat([
      Buffer.from("302e020100300506032b657004220420", "hex"),
      seed,
    ]),
    format: "der",
    type: "pkcs8",
  });
  const publicKeyObject = crypto.createPublicKey(privateKeyObject);
  const publicDer = publicKeyObject.export({ format: "der", type: "spki" });
  return {
    privateKey,
    publicKey: `ed25519:${Buffer.from(publicDer).subarray(-32).toString("base64")}`,
    privateKeyObject,
  };
}

function canonicalSignedRequest(value, topLevel = true) {
  if (Array.isArray(value)) {
    return value.map((entry) => canonicalSignedRequest(entry, false));
  }
  if (value && typeof value === "object") {
    return Object.fromEntries(
      Object.entries(value)
        .filter(
          ([key, entry]) =>
            entry !== null &&
            (!topLevel || !["agent_signature", "node_signature"].includes(key))
        )
        .sort(([left], [right]) => (left < right ? -1 : left > right ? 1 : 0))
        .map(([key, entry]) => [key, canonicalSignedRequest(entry, false)])
    );
  }
  return value;
}

function withWireDefaults(request) {
  const value = { ...request };
  if (value.type === "report_node_capabilities") {
    value.dependency_cache_digests ??= [];
  } else if (value.type === "launch_task" || value.type === "launch_child_task") {
    value.wait_for_node ??= false;
    if (value.task_spec && typeof value.task_spec === "object") {
      value.task_spec = { ...value.task_spec };
      value.task_spec.failure_policy ??= "fail_fast";
    }
  } else if (value.type === "start_process") {
    value.restart ??= false;
  } else if (value.type === "report_debug_state") {
    value.stack_frames ??= [];
    value.local_values ??= [];
    value.task_args ??= [];
    value.handles ??= [];
    value.recent_output ??= [];
  } else if (value.type === "report_task_log") {
    value.stdout_tail ??= "";
    value.stderr_tail ??= "";
  } else if (value.type === "task_completed") {
    value.stdout_tail ??= "";
    value.stderr_tail ??= "";
    value.stdout_truncated ??= false;
    value.stderr_truncated ??= false;
  }
  return value;
}

function signedRequestPayloadDigest(request) {
  return `sha256:${crypto
    .createHash("sha256")
    .update(JSON.stringify(canonicalSignedRequest(withWireDefaults(request))))
    .digest("hex")}`;
}

function nodeSignatureMessage(
  node,
  requestKind,
  payloadDigest,
  nonce,
  issuedAtEpochSeconds
) {
  const parts = [
    "clusterflux-node-request-signature:v2",
    node,
    requestKind,
    payloadDigest,
    nonce,
    String(issuedAtEpochSeconds),
  ];
  return Buffer.concat(
    parts.flatMap((part) => [
      Buffer.from(`${Buffer.byteLength(part)}:`),
      Buffer.from(part),
      Buffer.from("\n"),
    ])
  );
}

function signedNodeProof(node, identity, requestKind, request, options = {}) {
  const nonce =
    options.nonce ||
    `${requestKind}-${process.pid}-${Date.now()}-${crypto
      .randomBytes(8)
      .toString("hex")}`;
  const issuedAt =
    options.issuedAtEpochSeconds ?? Math.floor(Date.now() / 1000);
  const signature = crypto.sign(
    null,
    nodeSignatureMessage(
      node,
      requestKind,
      signedRequestPayloadDigest(request),
      nonce,
      issuedAt
    ),
    identity.privateKeyObject
  );
  return {
    nonce,
    issued_at_epoch_seconds: issuedAt,
    signature: `ed25519:${signature.toString("base64")}`,
  };
}

function signedNodeHeartbeat(node, identity, options = {}) {
  const request = { type: "node_heartbeat", node };
  return signedNodeProof(node, identity, "node_heartbeat", request, options);
}

function signedNodeRequest(node, identity, requestKind, request, options = {}) {
  return {
    type: "signed_node",
    node,
    node_signature: signedNodeProof(
      node,
      identity,
      requestKind,
      request,
      options
    ),
    request,
  };
}

module.exports = {
  nodeIdentity,
  nodeIdentityFromPrivateKey,
  signedNodeProof,
  signedRequestPayloadDigest,
  signedNodeHeartbeat,
  signedNodeRequest,
};