import { createCipheriv, createECDH, createPrivateKey, createPublicKey, diffieHellman, hkdfSync, randomBytes } from "node:crypto";
import type { Choice, CreateRequest, DevicePublicKeyRecord, PrivateRequestPrepareResponse, PrivateStatusUpdatePlaintext, PrivateStatusUpdatePrepareResponse } from "@self-deprecated/agent-tick-shared";
import type { CreateStatusUpdateInput, CreateToolActivityInput, PrivateStatusUpdatePayload, PrivateToolActivityPayload } from "./types";

const CONTENT_ALGORITHM = "aes-256-gcm" as const;
const ENVELOPE_ALGORITHM = "p256-ecdh-hkdf-sha256+aes-256-gcm" as const;

export function createPrivateStatusUpdateInput(
   base: CreateStatusUpdateInput,
   plaintext: PrivateStatusUpdatePlaintext,
   prepare: PrivateStatusUpdatePrepareResponse,
): CreateStatusUpdateInput {
   if (!prepare.deviceKeys.length) {
      throw new Error("Private Status Update could not be created: no recipient has a usable Private Request device key.");
   }
   return {
      ...base,
      contentMode: "private",
      encryptedPayload: encryptPrivateStatusUpdatePayload(plaintext, prepare.deviceKeys),
      privateRecipientVersion: prepare.recipientVersion,
   };
}

export function createPrivateToolActivityInput(
   base: CreateToolActivityInput,
   plaintext: Record<string, unknown>,
   prepare: PrivateStatusUpdatePrepareResponse,
): CreateToolActivityInput {
   if (!prepare.deviceKeys.length) {
      throw new Error("Private Tool Activity could not be created: no recipient has a usable Private Request device key.");
   }
   return {
      ...base,
      contentMode: "private",
      encryptedPayload: encryptPrivateStatusUpdatePayload(plaintext, prepare.deviceKeys) as PrivateToolActivityPayload,
      privateRecipientVersion: prepare.recipientVersion,
   };
}

export function createPrivateRequestInput(base: CreateRequest, prepare: PrivateRequestPrepareResponse): CreateRequest {
   if (!prepare.deviceKeys.length) {
      throw new Error("Private Request could not be created: no recipient has a usable Private Request device key.");
   }
   const choices = base.choices?.length ? base.choices.map(normalizeChoice) : defaultChoices(base.requestType ?? "sanction");
   return {
      ...base,
      title: "Private Request",
      body: undefined,
      command: undefined,
      questions: undefined,
      choices: publicChoicesForPrivateRequest(choices),
      allowFreeformReply: false,
      contentMode: "private",
      encryptedPayload: encryptPrivateStatusUpdatePayload({
         title: base.title,
         ...(base.body ? { body: base.body } : {}),
         ...(base.command ? { command: base.command } : {}),
         choices,
         ...(base.session ? { session: base.session } : {}),
         requestType: base.requestType ?? "sanction",
      }, prepare.deviceKeys),
      privateRecipientVersion: prepare.recipientVersion,
   };
}

function publicChoicesForPrivateRequest(choices: Choice[]): Choice[] {
   return choices.map((choice) => ({ id: choice.id, kind: choice.kind, label: choice.kind === "deny" ? "Deny" : "Approve" }));
}

function defaultChoices(requestType: string): Choice[] {
   if (requestType === "steering") return [{ id: "option_a", label: "Option A", kind: "approve" }, { id: "option_b", label: "Option B", kind: "approve" }, { id: "cancel", label: "Cancel", kind: "deny" }];
   return [{ id: "approve", label: "Approve", kind: "approve" }, { id: "deny", label: "Deny", kind: "deny" }];
}

function normalizeChoice(choice: NonNullable<CreateRequest["choices"]>[number]): Choice {
   return { ...choice, kind: choice.kind ?? "approve" };
}

export function encryptPrivateStatusUpdatePayload(plaintext: PrivateStatusUpdatePlaintext | Record<string, unknown>, deviceKeys: DevicePublicKeyRecord[]): PrivateStatusUpdatePayload {
   const contentKey = randomBytes(32);
   const nonce = randomBytes(12);
   const serialized = Buffer.from(JSON.stringify(plaintext), "utf8");
   const cipher = createCipheriv(CONTENT_ALGORITHM, contentKey, nonce);
   const ciphertext = Buffer.concat([cipher.update(serialized), cipher.final()]);
   const tag = cipher.getAuthTag();
   return {
      version: 1,
      algorithm: CONTENT_ALGORITHM,
      nonce: b64(nonce),
      ciphertext: b64(ciphertext),
      tag: b64(tag),
      keyEnvelopes: deviceKeys.map((key) => wrapContentKey(contentKey, key)),
   };
}

function wrapContentKey(contentKey: Buffer, deviceKey: DevicePublicKeyRecord): PrivateStatusUpdatePayload["keyEnvelopes"][number] {
   const recipientPublicKey = createPublicKey({ key: Buffer.from(deviceKey.publicKey, "base64url"), format: "der", type: "spki" });
   const ephemeral = createECDH("prime256v1");
   ephemeral.generateKeys();
   const ephemeralPrivateKey = ephemeralKeyPairPrivateKey(ephemeral);
   const sharedSecret = diffieHellman({ privateKey: ephemeralPrivateKey, publicKey: recipientPublicKey });
   // Use the existing native Private Request key agreement context so older mobile
   // crypto modules can decrypt private Status Updates without a native update.
   const wrappingKey = Buffer.from(hkdfSync("sha256", sharedSecret, Buffer.alloc(0), Buffer.from(`agent-tick-private-request:${deviceKey.deviceKeyId}`, "utf8"), 32));
   const nonce = randomBytes(12);
   const cipher = createCipheriv(CONTENT_ALGORITHM, wrappingKey, nonce);
   cipher.setAAD(Buffer.from(deviceKey.deviceKeyId, "utf8"));
   const ciphertext = Buffer.concat([cipher.update(contentKey), cipher.final()]);
   const tag = cipher.getAuthTag();
   return {
      deviceKeyId: deviceKey.deviceKeyId,
      algorithm: ENVELOPE_ALGORITHM,
      ephemeralPublicKey: b64(ephemeralKeyPairPublicSpki(ephemeral)),
      nonce: b64(nonce),
      ciphertext: b64(ciphertext),
      tag: b64(tag),
   };
}

function ephemeralKeyPairPrivateKey(ecdh: ReturnType<typeof createECDH>) {
   const publicKey = ecdh.getPublicKey(undefined, "uncompressed");
   const jwk = {
      kty: "EC",
      crv: "P-256",
      x: b64(publicKey.subarray(1, 33)),
      y: b64(publicKey.subarray(33, 65)),
      d: b64(ecdh.getPrivateKey()),
   };
   return createPrivateKey({ key: jwk, format: "jwk" });
}

function ephemeralKeyPairPublicSpki(ecdh: ReturnType<typeof createECDH>): Buffer {
   const key = createPublicKey(ephemeralKeyPairPrivateKey(ecdh));
   return key.export({ format: "der", type: "spki" }) as Buffer;
}

function b64(value: Buffer): string {
   return value.toString("base64url");
}