import { assertEquals } from "jsr:@std/assert";

import { buildMerkleTree, leafHash, nodeHash, verifyInclusionProof } from "./audit_merkle.ts";
import { encodeHex } from "jsr:@std/encoding/hex";

Deno.test("Audit Merkle - leafHash", async () => {
  const entry = new Uint8Array([0x01, 0x02, 0x03]);
  const hash = await leafHash(entry);

  // SHA-256(0x00 || 0x01, 0x02, 0x03)
  const expectedData = new Uint8Array([0x00, 0x01, 0x02, 0x03]);
  const expectedHash = await crypto.subtle.digest("SHA-256", expectedData);

  assertEquals(encodeHex(hash), encodeHex(new Uint8Array(expectedHash)));
});

Deno.test("Audit Merkle - nodeHash", async () => {
  const left = new Uint8Array([0x11, 0x22]);
  const right = new Uint8Array([0x33, 0x44]);
  const hash = await nodeHash(left, right);

  // SHA-256(0x01 || 0x11, 0x22 || 0x33, 0x44)
  const expectedData = new Uint8Array([0x01, 0x11, 0x22, 0x33, 0x44]);
  const expectedHash = await crypto.subtle.digest("SHA-256", expectedData);

  assertEquals(encodeHex(hash), encodeHex(new Uint8Array(expectedHash)));
});

Deno.test("Audit Merkle - buildMerkleTree (empty)", async () => {
  const root = await buildMerkleTree([]);
  const expectedHash = await crypto.subtle.digest("SHA-256", new Uint8Array(0));
  assertEquals(encodeHex(root), encodeHex(new Uint8Array(expectedHash)));
});

Deno.test("Audit Merkle - buildMerkleTree (1 leaf)", async () => {
  const leaf = await leafHash(new Uint8Array([0x01]));
  const root = await buildMerkleTree([leaf]);
  assertEquals(encodeHex(root), encodeHex(leaf));
});

Deno.test("Audit Merkle - buildMerkleTree (2 leaves)", async () => {
  const leaf1 = await leafHash(new Uint8Array([0x01]));
  const leaf2 = await leafHash(new Uint8Array([0x02]));
  const root = await buildMerkleTree([leaf1, leaf2]);

  const expectedNode = await nodeHash(leaf1, leaf2);
  assertEquals(encodeHex(root), encodeHex(expectedNode));
});

Deno.test("Audit Merkle - buildMerkleTree (3 leaves)", async () => {
  const leaf1 = await leafHash(new Uint8Array([0x01]));
  const leaf2 = await leafHash(new Uint8Array([0x02]));
  const leaf3 = await leafHash(new Uint8Array([0x03]));

  const root = await buildMerkleTree([leaf1, leaf2, leaf3]);

  const leftChild = await nodeHash(leaf1, leaf2);
  const expectedNode = await nodeHash(leftChild, leaf3);
  assertEquals(encodeHex(root), encodeHex(expectedNode));
});

import { auditLog, flush } from "./audit.ts";
import { sqlWrapper } from "./db.ts";

Deno.test("Audit Logger - Micro-batcher Flush", async () => {
  const queries: any[] = [];

  const originalSql = sqlWrapper.sql;
  // Mock SQL
  const mockSql = ((strings: any, ...values: any[]) => {
    const queryStr = strings.join("?");
    queries.push({ queryStr, values });

    if (queryStr.includes("SELECT tree_size FROM audit_sths")) {
      return Promise.resolve([{ tree_size: "0" }]);
    }

    if (queryStr.includes("SELECT leaf_hash FROM audit_records")) {
      return Promise.resolve([
        { leaf_hash: "010203" }, // dummy hex
        { leaf_hash: "040506" },
      ]);
    }

    return Promise.resolve([]);
  }) as any;

  sqlWrapper.sql = mockSql;

  const { valkey } = await import("./valkey.ts");
  const pubCalls: any[] = [];
  const setCalls: any[] = [];

  const originalPub = valkey.publish;
  const originalSet = valkey.set;

  valkey.publish = (async (channel: string, message: string) => {
    pubCalls.push({ channel, message });
    return 1;
  }) as any;

  valkey.set = (async (key: string, value: string) => {
    setCalls.push({ key, value });
    return "OK";
  }) as any;

  await flush();

  // Clean up mocks
  sqlWrapper.sql = originalSql;
  valkey.publish = originalPub;
  valkey.set = originalSet;

  // The flush should have called Valkey pub/sub and set
  assertEquals(pubCalls.length, 1);
  assertEquals(pubCalls[0].channel, "auth:audit:sth");

  assertEquals(setCalls.length, 1);
  assertEquals(setCalls[0].key, "auth:audit:latest_sth");

  const payload = JSON.parse(pubCalls[0].message);
  assertEquals(payload.tree_size, 2);
  assertEquals(typeof payload.root_hash, "string");
  assertEquals(typeof payload.signature, "string");
  assertEquals(typeof payload.created_at, "string");

  const insertSTHQuery = queries.find((q) =>
    q.queryStr.includes("INSERT INTO audit_sths")
  );
  assertEquals(insertSTHQuery !== undefined, true);
});

Deno.test("Audit Logger - auditLog computes leaf hash synchronously", async () => {
  const queries: any[] = [];

  const originalSql = sqlWrapper.sql;
  const mockSql = ((strings: any, ...values: any[]) => {
    const queryStr = strings.join("?");
    queries.push({ queryStr, values });
    return Promise.resolve([]);
  }) as any;

  sqlWrapper.sql = mockSql;

  auditLog("user1", "test_action", "res1", { test: 123 }, "127.0.0.1");

  // Wait a small amount for the promise to resolve internally
  await new Promise((r) => setTimeout(r, 10));

  sqlWrapper.sql = originalSql;

  assertEquals(queries.length, 1);
  const insertQuery = queries[0];
  assertEquals(
    insertQuery.queryStr.includes("INSERT INTO audit_records"),
    true,
  );

  // Should have leaf_hash
  assertEquals(insertQuery.queryStr.includes("leaf_hash"), true);

  // The last value in the values array is the leaf_hash
  const lastValue = insertQuery.values[insertQuery.values.length - 1];
  assertEquals(typeof lastValue, "string");
  assertEquals(lastValue.length, 64); // SHA-256 hex is 64 chars
});

Deno.test("Audit Merkle - verifyInclusionProof", async () => {
  const leaf1 = await leafHash(new Uint8Array([0x01]));
  const leaf2 = await leafHash(new Uint8Array([0x02]));
  const leaf3 = await leafHash(new Uint8Array([0x03]));
  const leaf4 = await leafHash(new Uint8Array([0x04]));

  const root = await buildMerkleTree([leaf1, leaf2, leaf3, leaf4]);

  const node12 = await nodeHash(leaf1, leaf2);
  const node34 = await nodeHash(leaf3, leaf4);

  // Proof for leaf 1 (index 0): sibling is leaf2, then sibling is node34
  const proof1 = [leaf2, node34];
  const isValid1 = await verifyInclusionProof(leaf1, proof1, 0, 4, root);
  assertEquals(isValid1, true);
});