/**
 * restore.test.ts — coverage for ChunkBuffer, parseRawScan, reconstructAndVerify.
 *
 * Tests build synthetic envelopes via the same helpers the generator uses
 * (Subagent 1's `helpers.ts` + `erasure.ts`). We avoid the full
 * `generateCodex` pipeline so these tests don't depend on PDF rendering
 * or Aztec encoding — they only exercise the restoration math.
 */

import { beforeEach, describe, expect, it } from "vitest";
import { deflate } from "pako";

import {
  CodexError,
  deriveShardCounts,
  type ChunkEnvelope,
  type CodexPayload,
} from "../../shared/codex-types";
import {
  base64UrlToBytes,
  bytesToBase64Url,
  buildChunkEnvelope,
  sha256Hex,
} from "./helpers";
import { encodeWithErasure } from "./erasure";
import { _resetReedSolomonCacheForTest } from "./rs-loader";
import {
  ChunkBuffer,
  parseRawScan,
  reconstructAndVerify,
} from "./restore";

// Reset between tests: the WASM Reed-Solomon instance carries internal state
// across encode/reconstruct calls that breaks subsequent reconstructions in
// the same process. Until that's fixed in erasure.ts / rs-loader.ts, isolating
// each test keeps the suite deterministic.
beforeEach(() => {
  _resetReedSolomonCacheForTest();
});

const VERIFY_CODE = "PX-TEST-ABCD";

/** Build envelopes for an arbitrary CodexPayload, mirroring the real generator. */
async function buildEnvelopesForPayload(
  payload: CodexPayload,
): Promise<ChunkEnvelope[]> {
  const json = JSON.stringify(payload);
  const compressed = deflate(new TextEncoder().encode(json), { level: 9 });
  // encodeWithErasure shards at CHUNK_SIZE internally — no need to pass it.
  const { dataChunks, parityChunks, totalChunks } =
    await encodeWithErasure(compressed);
  const all = [...dataChunks, ...parityChunks];
  return Promise.all(
    all.map((bytes, i) =>
      buildChunkEnvelope(
        bytes,
        i,
        totalChunks,
        i >= dataChunks.length,
        payload.verifyCode,
      ),
    ),
  );
}

async function makePayload(seedLength = 12_000): Promise<CodexPayload> {
  // A vaultBlob big enough to span several chunks after deflate. We seed with
  // a deterministic but non-repeating pattern (xorshift-ish) so pako cannot
  // compress it to a single chunk — without that, dataShards collapses to 1
  // and the "lose 2 chunks" test has nothing to reconstruct.
  const blobBytes = new Uint8Array(seedLength);
  let s = 0x9e3779b9;
  for (let i = 0; i < seedLength; i++) {
    s = (s ^ (s << 13)) >>> 0;
    s = (s ^ (s >>> 17)) >>> 0;
    s = (s ^ (s << 5)) >>> 0;
    blobBytes[i] = s & 0xff;
  }
  const vaultBlob = bytesToBase64Url(blobBytes);
  const vaultFingerprint = (
    await sha256Hex(new TextEncoder().encode(vaultBlob))
  )
    .slice(0, 32)
    .toUpperCase();
  return {
    version: 2,
    verifyCode: VERIFY_CODE,
    vaultBlob,
    timestamp: 1_700_000_000_000,
    entryCount: 42,
    vaultFingerprint,
    generatedAt: 1_700_000_001_000,
  };
}

async function makeSingleEnvelope(opts: {
  index: number;
  total: number;
  isParity?: boolean;
  verifyCode?: string;
  bytes?: Uint8Array;
}): Promise<ChunkEnvelope> {
  return buildChunkEnvelope(
    opts.bytes ?? new Uint8Array([1, 2, 3, 4, 5]),
    opts.index,
    opts.total,
    opts.isParity ?? false,
    opts.verifyCode ?? VERIFY_CODE,
  );
}

describe("ChunkBuffer", () => {
  it("accepts a valid chunk and exposes context", async () => {
    const buf = new ChunkBuffer();
    const env = await makeSingleEnvelope({ index: 0, total: 5 });
    const res = await buf.add(env);
    expect(res.accepted).toBe(true);
    expect(buf.size()).toBe(1);
    expect(buf.getContext()).toEqual({
      verifyCode: VERIFY_CODE,
      stripeCount: 1,
      total: 5,
      dataShards: deriveShardCounts(5).dataShards,
      scanned: 1,
      allStripesObserved: true,
    });
  });

  it("rejects a chunk from a different verifyCode", async () => {
    const buf = new ChunkBuffer();
    await buf.add(await makeSingleEnvelope({ index: 0, total: 5 }));
    const foreign = await makeSingleEnvelope({
      index: 1,
      total: 5,
      verifyCode: "PX-OTHR-XXXX",
    });
    const res = await buf.add(foreign);
    expect(res).toEqual({ accepted: false, reason: "foreign_codex" });
    expect(buf.size()).toBe(1);
  });

  it("rejects a chunk with bad integrity hash", async () => {
    const buf = new ChunkBuffer();
    const env = await makeSingleEnvelope({ index: 0, total: 5 });
    // Corrupt the data field without updating the hash — verifyChunkIntegrity
    // must catch this.
    const corrupted: ChunkEnvelope = {
      ...env,
      d: bytesToBase64Url(new Uint8Array([9, 9, 9, 9])),
    };
    const res = await buf.add(corrupted);
    expect(res).toEqual({ accepted: false, reason: "bad_integrity" });
    expect(buf.size()).toBe(0);
  });

  it("rejects a chunk with wrong envelope version", async () => {
    const buf = new ChunkBuffer();
    const env = await makeSingleEnvelope({ index: 0, total: 5 });
    const wrongVersion = { ...env, v: 1 } as unknown as ChunkEnvelope;
    const res = await buf.add(wrongVersion);
    expect(res.accepted).toBe(false);
    expect(res.reason).toBe("parse_error");
  });

  it("deduplicates: adding the same index twice keeps only one", async () => {
    const buf = new ChunkBuffer();
    const env = await makeSingleEnvelope({ index: 2, total: 5 });
    expect((await buf.add(env)).accepted).toBe(true);
    const second = await buf.add(env);
    expect(second).toEqual({ accepted: false, reason: "duplicate" });
    expect(buf.size()).toBe(1);
  });

  it("isReady() flips true at dataShards count", async () => {
    const buf = new ChunkBuffer();
    const total = 5;
    const dataShards = deriveShardCounts(total).dataShards;
    for (let i = 0; i < dataShards; i++) {
      expect(buf.isReady()).toBe(false);
      const env = await makeSingleEnvelope({
        index: i,
        total,
        bytes: new Uint8Array([i, i, i]),
      });
      await buf.add(env);
    }
    expect(buf.isReady()).toBe(true);
  });

  it("toStripeChunks returns raw chunk bytes per stripe", async () => {
    const buf = new ChunkBuffer();
    const env = await makeSingleEnvelope({
      index: 0,
      total: 5,
      bytes: new Uint8Array([1, 2, 3]),
    });
    await buf.add(env);
    const stripes = buf.toStripeChunks();
    expect(stripes.length).toBe(1);
    const map = stripes[0]!.chunks;
    expect(map.size).toBe(1);
    expect(Array.from(map.get(0)!)).toEqual([1, 2, 3]);
  });

  it("stripeReadiness exposes per-stripe counts", async () => {
    const buf = new ChunkBuffer();
    const env = await makeSingleEnvelope({ index: 0, total: 5 });
    await buf.add(env);
    expect(buf.stripeReadiness()).toEqual([
      {
        stripe: 0,
        found: 1,
        required: deriveShardCounts(5).dataShards,
        isReady: false,
      },
    ]);
  });
});

describe("parseRawScan", () => {
  it("parses a well-formed envelope", async () => {
    const env = await makeSingleEnvelope({ index: 0, total: 5 });
    const parsed = parseRawScan(JSON.stringify(env));
    expect(parsed).toEqual(env);
  });

  it("returns parse_error for non-JSON input", () => {
    expect(parseRawScan("not json")).toEqual({ error: "parse_error" });
  });

  it("returns wrong_version for v=1 envelopes", () => {
    const v1 = {
      v: 1,
      c: VERIFY_CODE,
      i: 0,
      t: 5,
      p: false,
      d: "AAAA",
      h: "ABCDEF12",
    };
    expect(parseRawScan(JSON.stringify(v1))).toEqual({
      error: "wrong_version",
    });
  });

  it("returns wrong_version for a v2 object", () => {
    expect(parseRawScan(JSON.stringify({ v: 2, c: "PX" }))).toEqual({
      error: "wrong_version",
    });
  });
});

describe("reconstructAndVerify", () => {
  it("round-trips a full envelope set back to the original payload", async () => {
    const payload = await makePayload();
    const envelopes = await buildEnvelopesForPayload(payload);

    const buf = new ChunkBuffer();
    for (const env of envelopes) {
      const res = await buf.add(env);
      expect(res.accepted).toBe(true);
    }
    expect(buf.isReady()).toBe(true);

    const restored = await reconstructAndVerify(buf);
    expect(restored.verifyCode).toBe(payload.verifyCode);
    expect(restored.vaultBlob).toBe(payload.vaultBlob);
    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);
    expect(restored.entryCount).toBe(payload.entryCount);
  });

  it("reconstructs after losing 2 chunks", async () => {
    const payload = await makePayload();
    const envelopes = await buildEnvelopesForPayload(payload);
    // Need at least dataShards + 2 chunks total so dropping 2 still leaves
    // enough for reconstruction (small vaults pin to the 2-shard parity floor).
    expect(envelopes.length).toBeGreaterThan(2);

    // Drop two indices that exercise the parity-decode path: one data shard
    // and one parity shard. The parity floor of 2 means we can lose any two.
    const dropA = 1;
    const dropB = envelopes.length - 1;

    const buf = new ChunkBuffer();
    for (let i = 0; i < envelopes.length; i++) {
      if (i === dropA || i === dropB) continue;
      const res = await buf.add(envelopes[i]!);
      expect(res.accepted).toBe(true);
    }
    expect(buf.isReady()).toBe(true);

    const restored = await reconstructAndVerify(buf);
    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);
    expect(restored.vaultBlob).toBe(payload.vaultBlob);
  });

  it("throws RECONSTRUCTION_FINGERPRINT_MISMATCH if the reconstructed blob does not match", async () => {
    // Build a payload whose stored fingerprint deliberately does NOT match
    // its vaultBlob — simulates a successful chunk decode that nevertheless
    // pulled the wrong bytes into a parity-reconstructed slot.
    const payload = await makePayload();
    const tampered: CodexPayload = {
      ...payload,
      vaultFingerprint: "DEADBEEFDEADBEEFDEADBEEFDEADBEEF",
    };
    const envelopes = await buildEnvelopesForPayload(tampered);

    const buf = new ChunkBuffer();
    for (const env of envelopes) await buf.add(env);

    await expect(reconstructAndVerify(buf)).rejects.toMatchObject({
      name: "CodexError",
      code: "RECONSTRUCTION_FINGERPRINT_MISMATCH",
    });
  });

  it("throws CodexError when called on an empty buffer", async () => {
    const buf = new ChunkBuffer();
    await expect(reconstructAndVerify(buf)).rejects.toBeInstanceOf(CodexError);
  });
});

describe("ChunkBuffer + parseRawScan integration", () => {
  it("a full scan loop accepts every raw envelope string", async () => {
    const payload = await makePayload(2_500);
    const envelopes = await buildEnvelopesForPayload(payload);
    const buf = new ChunkBuffer();

    for (const env of envelopes) {
      const parsed = parseRawScan(JSON.stringify(env));
      expect("error" in parsed).toBe(false);
      const res = await buf.add(parsed as ChunkEnvelope);
      expect(res.accepted).toBe(true);
    }

    const restored = await reconstructAndVerify(buf);
    expect(restored.vaultFingerprint).toBe(payload.vaultFingerprint);
    // sanity-check that base64UrlToBytes is reachable from this file
    expect(base64UrlToBytes(restored.vaultBlob).length).toBeGreaterThan(0);
  });
});