import { describe, expect, test } from 'bun:test';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import {
  KTLogManager,
  LightWitness,
  MemoryKTLogStore,
  computeBundleHash,
  computeLogId,
  signSth,
  sthToWire,
} from '../src/index.js';
import {
  KTLogIdMismatchError,
  KTSplitViewError,
  KTStaleSTHError,
  KTVerificationError,
} from '../src/errors.js';
import { toBase64, fromBase64 } from '../src/util.js';

const crypto = new SubtleCryptoProvider();

async function setup() {
  const kp = await crypto.generateEd25519KeyPair();
  const store = new MemoryKTLogStore();
  const mgr = await KTLogManager.create({
    crypto,
    store,
    signingPrivateKey: kp.privateKey,
    signingPublicKey: kp.publicKey,
  });
  return { kp, mgr };
}

function fakeBundle(seed: number) {
  return {
    identitySigningKey: new Uint8Array(32).fill(seed),
    identityDHKey: new Uint8Array(32).fill(seed + 1),
    signedPreKey: {
      keyId: 1,
      publicKey: new Uint8Array(32).fill(seed + 2),
      signature: new Uint8Array(64).fill(seed + 3),
    },
  };
}

describe('LightWitness', () => {
  test('observes valid STH and stores it', async () => {
    const { kp, mgr } = await setup();
    await mgr.recordRegister('alice', computeBundleHash(fakeBundle(0x10)));
    const sth = await mgr.publishSTH();

    const witness = new LightWitness({
      crypto,
      logPublicKey: kp.publicKey,
      fetcher: {
        async fetchLatestSTH() {
          return sthToWire(sth, toBase64);
        },
        async fetchConsistencyProof() {
          return { proof: [] };
        },
      },
    });
    const polled = await witness.pollOnce();
    expect(polled.treeSize).toBe(1);
    expect(witness.compare(sth)).toBe('agree');
  });

  test('rejects STH whose log_id does not match pinned key', async () => {
    const { mgr } = await setup();
    const wrong = await crypto.generateEd25519KeyPair();
    await mgr.recordRegister('alice', computeBundleHash(fakeBundle(0x10)));
    const sth = await mgr.publishSTH();

    const witness = new LightWitness({
      crypto,
      logPublicKey: wrong.publicKey, // pinned to wrong key
      fetcher: {
        async fetchLatestSTH() {
          return sthToWire(sth, toBase64);
        },
        async fetchConsistencyProof() {
          return { proof: [] };
        },
      },
    });
    await expect(witness.pollOnce()).rejects.toBeInstanceOf(KTLogIdMismatchError);
  });

  test('rejects STH older than maxStaleMs', async () => {
    const { kp, mgr } = await setup();
    await mgr.recordRegister('alice', computeBundleHash(fakeBundle(0x10)));
    const sth = await mgr.publishSTH(1000); // far in the past
    const witness = new LightWitness({
      crypto,
      logPublicKey: kp.publicKey,
      fetcher: {
        async fetchLatestSTH() {
          return sthToWire(sth, toBase64);
        },
        async fetchConsistencyProof() {
          return { proof: [] };
        },
      },
      maxStaleMs: 1000,
      now: () => 10_000_000,
    });
    await expect(witness.pollOnce()).rejects.toBeInstanceOf(KTStaleSTHError);
  });

  test('detects split-view at same tree_size', async () => {
    const { kp, mgr } = await setup();
    await mgr.recordRegister('alice', computeBundleHash(fakeBundle(0x10)));
    const sth1 = await mgr.publishSTH();

    // Forge another signed STH with same tree_size but different rootHash
    const tamperedRoot = new Uint8Array(sth1.rootHash);
    tamperedRoot[0] ^= 0xff;
    const sth2 = await signSth(crypto, kp.privateKey, {
      treeSize: sth1.treeSize,
      timestampMs: sth1.timestampMs,
      rootHash: tamperedRoot,
      indexRoot: sth1.indexRoot,
      logId: computeLogId(kp.publicKey),
    });

    const witness = new LightWitness({
      crypto,
      logPublicKey: kp.publicKey,
      fetcher: {
        async fetchLatestSTH() {
          return sthToWire(sth1, toBase64);
        },
        async fetchConsistencyProof() {
          return { proof: [] };
        },
      },
    });
    await witness.observe(sth1);
    await expect(witness.observe(sth2)).rejects.toBeInstanceOf(KTSplitViewError);
  });

  test('verifies consistency between two successive STHs', async () => {
    const { kp, mgr } = await setup();
    await mgr.recordRegister('alice', computeBundleHash(fakeBundle(0x10)));
    const sth1 = await mgr.publishSTH();
    await mgr.recordRegister('bob', computeBundleHash(fakeBundle(0x20)));
    const sth2 = await mgr.publishSTH();

    const consistency = await mgr.buildConsistencyProof(sth1.treeSize);
    const witness = new LightWitness({
      crypto,
      logPublicKey: kp.publicKey,
      fetcher: {
        async fetchLatestSTH() {
          return sthToWire(sth2, toBase64);
        },
        async fetchConsistencyProof() {
          return { proof: consistency.proof.map(toBase64) };
        },
      },
    });
    await witness.observe(sth1);
    await witness.observe(sth2);
    expect(witness.compare(sth2)).toBe('agree');
  });

  test('rejects STH where log re-wrote history (consistency proof fails)', async () => {
    const { kp, mgr } = await setup();
    await mgr.recordRegister('alice', computeBundleHash(fakeBundle(0x10)));
    const sth1 = await mgr.publishSTH();

    // Build a forked log where leaf 0 is different.
    const forkStore = new MemoryKTLogStore();
    const forkMgr = await KTLogManager.create({
      crypto,
      store: forkStore,
      signingPrivateKey: kp.privateKey,
      signingPublicKey: kp.publicKey,
    });
    await forkMgr.recordRegister('mallory', computeBundleHash(fakeBundle(0xee)));
    await forkMgr.recordRegister('bob', computeBundleHash(fakeBundle(0x20)));
    const forkedSth2 = await forkMgr.publishSTH();
    const forkedConsistency = await forkMgr.buildConsistencyProof(sth1.treeSize);

    const witness = new LightWitness({
      crypto,
      logPublicKey: kp.publicKey,
      fetcher: {
        async fetchLatestSTH() {
          return sthToWire(forkedSth2, toBase64);
        },
        async fetchConsistencyProof() {
          return { proof: forkedConsistency.proof.map(toBase64) };
        },
      },
    });
    await witness.observe(sth1);
    await expect(witness.observe(forkedSth2)).rejects.toBeInstanceOf(KTVerificationError);
  });

  test('compare returns "unknown" for tree_size we have not seen', async () => {
    const { kp, mgr } = await setup();
    await mgr.recordRegister('alice', computeBundleHash(fakeBundle(0x10)));
    const sth = await mgr.publishSTH();
    const witness = new LightWitness({
      crypto,
      logPublicKey: kp.publicKey,
      fetcher: {
        async fetchLatestSTH() {
          return sthToWire(sth, toBase64);
        },
        async fetchConsistencyProof() {
          return { proof: [] };
        },
      },
    });
    expect(witness.compare(sth)).toBe('unknown');
  });
});

// Make TS happy about unused fromBase64
void fromBase64;