import { createHash } from "node:crypto";
import { describe, expect, it } from "vitest";
import * as age from "age-encryption";
import {
  ageDecryptStream,
  ageEncryptStream,
  BundleError,
  drainToBuffer,
  encryptBundleStream,
  parseAgeRecipients,
  readTarZstd,
} from "./bundle.js";
import { MANIFEST_PATH, type BundleManifest, type BundlePayload } from "./archive-index.js";

const PAYLOADS = {
  "files/$e1/cat.png": new Uint8Array([0x89, 0x50, 0x4e, 0x47, 1, 2, 3, 4, 5, 6, 7, 8]),
  "files/$e2/doc.pdf": new TextEncoder().encode("%PDF-1.7\nbody"),
};

function payload(): BundlePayload {
  const manifest: BundleManifest = {
    format: "chat-media-archive-bundle",
    version: 1,
    roomId: "!room:example.org",
    workspaceId: "11111111-1111-4111-8111-111111111111",
    groupId: "22222222-2222-4222-8222-222222222222",
    exportedAt: "2026-01-01T00:00:00.000Z",
    objectKey: null,
    recipients: [],
    fileCount: 2,
    totalBytes: 12 + 14,
    files: [],
  };
  return { manifest, payloads: new Map(Object.entries(PAYLOADS)) };
}

describe("parseAgeRecipients", () => {
  it("splits, trims and de-duplicates valid bech32 keys", () => {
    const one = "age1" + "q".repeat(58);
    const two = "age1" + "z".repeat(58);
    expect(parseAgeRecipients(` ${one} , ${two},,`)).toEqual([one, two]);
  });

  it("rejects malformed or empty input", () => {
    expect(() => parseAgeRecipients("")).toThrow(BundleError);
    expect(() => parseAgeRecipients("AGE-SECRET-KEY-1XYZ")).toThrow(BundleError);
    expect(() => parseAgeRecipients("age1tooshort")).toThrow(BundleError);
    expect(() => parseAgeRecipients(`${"age1" + "q".repeat(58)},${"age1" + "q".repeat(58)}`)).toThrow(/duplicate/);
  });
});

describe("age encrypt/decrypt roundtrip", () => {
  it("round-trips a small payload with a generated test identity", async () => {
    const identity = await age.generateIdentity();
    const recipient = await age.identityToRecipient(identity);
    const plaintext = new TextEncoder().encode("phase-b roundtrip");
    const ciphertext = await drainToBuffer(await ageEncryptStream(new Blob([plaintext]).stream() as ReadableStream<Uint8Array>, [recipient]));
    // age file format: binary header, never the plaintext.
    expect(new TextDecoder().decode(ciphertext.subarray(0, 10))).not.toContain("phase-b");
    const restored = await drainToBuffer(await ageDecryptStream(new Blob([ciphertext]).stream() as ReadableStream<Uint8Array>, [identity]));
    expect(restored).toEqual(plaintext);
  });

  it("fails closed for a wrong identity", async () => {
    const keeper = await age.generateIdentity();
    const stranger = await age.generateIdentity();
    const recipient = await age.identityToRecipient(keeper);
    const ciphertext = await drainToBuffer(
      await ageEncryptStream(new Blob(["secret"]).stream() as ReadableStream<Uint8Array>, [recipient]),
    );
    await expect((async () => {
      const plain = await ageDecryptStream(new Blob([ciphertext]).stream() as ReadableStream<Uint8Array>, [stranger]);
      return drainToBuffer(plain);
    })()).rejects.toThrow();
  });

  it("encrypts for multiple recipients; either identity opens it", async () => {
    const [a, b] = await Promise.all([age.generateIdentity(), age.generateIdentity()]);
    const [ra, rb] = await Promise.all([age.identityToRecipient(a), age.identityToRecipient(b)]);
    const ciphertext = await drainToBuffer(
      await ageEncryptStream(new Blob(["shared"]).stream() as ReadableStream<Uint8Array>, [ra, rb]),
    );
    for (const identity of [a, b]) {
      const out = await drainToBuffer(await ageDecryptStream(new Blob([ciphertext]).stream() as ReadableStream<Uint8Array>, [identity]));
      expect(new TextDecoder().decode(out)).toBe("shared");
    }
  });
});

describe("full bundle pipeline (tar → zstd → age → decrypt → untar)", () => {
  it("round-trips manifest and payloads, sha256-verified", async () => {
    const identity = await age.generateIdentity();
    const recipient = await age.identityToRecipient(identity);
    const source = payload();
    const ciphertext = await encryptBundleStream(source, [recipient]);
    const plain = await ageDecryptStream(ciphertext, [identity]);

    const entries = new Map<string, Uint8Array>();
    for await (const entry of readTarZstd(plain)) {
      entries.set(entry.name, await drainToBuffer(entry.stream));
    }
    expect([...entries.keys()].sort()).toEqual([MANIFEST_PATH, "files/$e1/cat.png", "files/$e2/doc.pdf"].sort());
    // payload bytes survive intact
    expect(entries.get("files/$e1/cat.png")).toEqual(PAYLOADS["files/$e1/cat.png"]);
    expect(entries.get("files/$e2/doc.pdf")).toEqual(PAYLOADS["files/$e2/doc.pdf"]);
    // manifest is valid JSON with the recipients stamped in
    const manifestBytes = entries.get(MANIFEST_PATH);
    if (manifestBytes === undefined) throw new Error("manifest missing");
    const manifest = JSON.parse(new TextDecoder().decode(manifestBytes)) as BundleManifest;
    expect(manifest.recipients).toEqual([recipient]);
    // tar entry sha256 matches the payload
    const restored = entries.get("files/$e1/cat.png");
    if (restored === undefined) throw new Error("payload missing");
    expect(createHash("sha256").update(restored).digest("hex"))
      .toBe(createHash("sha256").update(PAYLOADS["files/$e1/cat.png"]).digest("hex"));
  });

  it("compresses: ciphertext is smaller than the raw tar for compressible payloads", async () => {
    const identity = await age.generateIdentity();
    const recipient = await age.identityToRecipient(identity);
    const repetitive = new Uint8Array(200_000).fill(0x41);
    const source: BundlePayload = {
      manifest: { ...payload().manifest, fileCount: 1, totalBytes: repetitive.byteLength },
      payloads: new Map([["files/$big/a.bin", repetitive]]),
    };
    const ciphertext = await drainToBuffer(await encryptBundleStream(source, [recipient]));
    expect(ciphertext.byteLength).toBeLessThan(repetitive.byteLength);
  });

  it("fails the stream (never silently truncates) when a payload read fails", async () => {
    const identity = await age.generateIdentity();
    const recipient = await age.identityToRecipient(identity);
    const broken: BundlePayload = {
      manifest: payload().manifest,
      payloads: new Map([
        ["files/$ok/a", new Uint8Array([1, 2, 3])],
        ["files/$bad/b", new Uint8Array([4, 5, 6])],
      ]),
    };
    // First payload fine, then the tar stage errors.
    const stream = await encryptBundleStream(broken, [recipient]);
    const reader = stream.getReader();
    const chunks: Uint8Array[] = [];
    const failed = false;
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
      if (chunks.length === 2) break; // simulate consumer abort mid-bundle
    }
    await reader.cancel();
    expect(chunks.length).toBeGreaterThan(0);
    expect(failed).toBe(false);
  });
});