const { expect } = require("chai");
const { ethers } = require("hardhat");
const tss = require("./helpers/tss");

// ---------------------------------------------------------------------------
// AereThresholdRegistry -- on-chain registry of t-of-n threshold-ECDSA committees.
//
// The signatures in these tests are produced by an INDEPENDENT CJS reimplementation
// of the same Feldman-VSS DKG + threshold combine that ships in sdk-js/src/mpc. A
// threshold signature is an ordinary (r,s,v) recovering to the DKG group address, so
// the contract verifies it with plain ecrecover. Two facts these tests establish:
//   * a valid t-of-n signature verifies; a below-threshold (t-1) combine does NOT.
//   * committee governance (reshare / rotate / slash) is authorized by a threshold
//     signature of the CURRENT committee over a domain-separated intent hash.
//
// HONEST caveat mirrored from the contract: on-chain, a threshold sig is
// indistinguishable from a single-key sig; "t participated" is an off-chain property.
// ---------------------------------------------------------------------------

const ACTION_RESHARE = ethers.keccak256(ethers.toUtf8Bytes("AereThresholdRegistry.reshare.v1"));
const ACTION_ROTATE = ethers.keccak256(ethers.toUtf8Bytes("AereThresholdRegistry.rotateKey.v1"));
const ACTION_SLASH = ethers.keccak256(ethers.toUtf8Bytes("AereThresholdRegistry.slashMember.v1"));

const abi = ethers.AbiCoder.defaultAbiCoder();

// Random member identity addresses (not the group key; just committee roster).
function makeMembers(n) {
  const a = [];
  for (let i = 0; i < n; i++) a.push(ethers.Wallet.createRandom().address);
  return a;
}

// Recompute the contract's intent hashes in JS so we can threshold-sign them.
function reshareIntent(chainId, addr, id, epoch, nonce, newMembers, newT) {
  return ethers.keccak256(
    abi.encode(
      ["bytes32", "uint256", "address", "uint256", "uint64", "uint64", "address[]", "uint8"],
      [ACTION_RESHARE, chainId, addr, id, epoch, nonce, newMembers, newT],
    ),
  );
}
function rotateIntent(chainId, addr, id, epoch, nonce, newPubKey, newMembers, newT) {
  return ethers.keccak256(
    abi.encode(
      ["bytes32", "uint256", "address", "uint256", "uint64", "uint64", "bytes32", "address[]", "uint8"],
      [ACTION_ROTATE, chainId, addr, id, epoch, nonce, ethers.keccak256(newPubKey), newMembers, newT],
    ),
  );
}
function slashIntent(chainId, addr, id, epoch, nonce, member, evidence) {
  return ethers.keccak256(
    abi.encode(
      ["bytes32", "uint256", "address", "uint256", "uint64", "uint64", "address", "bytes32"],
      [ACTION_SLASH, chainId, addr, id, epoch, nonce, member, evidence],
    ),
  );
}

describe("AereThresholdRegistry", function () {
  let reg, chainId, deployer, guardian, outsider;

  beforeEach(async function () {
    [deployer, guardian, outsider] = await ethers.getSigners();
    const F = await ethers.getContractFactory("AereThresholdRegistry");
    reg = await F.deploy();
    await reg.waitForDeployment();
    chainId = (await ethers.provider.getNetwork()).chainId;
  });

  async function register(n, t, guardianAddr = ethers.ZeroAddress) {
    const dkg = tss.runDkg(n, t);
    const members = makeMembers(n);
    const rc = await reg.registerCommittee(members, t, dkg.groupPubKey64, guardianAddr);
    const rcpt = await rc.wait();
    // committeeId is committeeCount after the call
    const id = await reg.committeeCount();
    return { dkg, members, id: Number(id) };
  }

  it("registers a committee and stores the DKG group key + roster", async function () {
    const { dkg, members, id } = await register(3, 2);
    const c = await reg.getCommittee(id);
    expect(c.threshold).to.equal(2);
    expect(c.size).to.equal(3);
    expect(c.epoch).to.equal(1n);
    expect(c.groupAddress.toLowerCase()).to.equal(dkg.groupAddress.toLowerCase());
    expect(c.groupPubKey.toLowerCase()).to.equal(dkg.groupPubKey64.toLowerCase());
    const onchainMembers = await reg.getMembers(id);
    expect(onchainMembers.map((x) => x.toLowerCase())).to.deep.equal(members.map((x) => x.toLowerCase()));
    for (const m of members) expect(await reg.isMember(id, m)).to.equal(true);
  });

  it("verifies a valid t-of-n threshold signature (2 of 3)", async function () {
    const { dkg, id } = await register(3, 2);
    const digest = ethers.keccak256(ethers.toUtf8Bytes("withdraw 5 AERE"));
    const sig = tss.thresholdSign(ethers.getBytes(digest), [dkg.shares[0], dkg.shares[1]]);
    expect(await reg.verifyThresholdSignature(id, digest, sig)).to.equal(true);
    // a different valid quorum {1,3} verifies too (same group key)
    const sig2 = tss.thresholdSign(ethers.getBytes(digest), [dkg.shares[0], dkg.shares[2]]);
    expect(await reg.verifyThresholdSignature(id, digest, sig2)).to.equal(true);
    // EIP-1271-style helper returns the magic value
    expect(await reg.isValidGroupSignature(id, digest, sig)).to.equal("0x1626ba7e");
  });

  it("REJECTS a below-threshold signature (1 of 3 -> wrong key)", async function () {
    const { dkg, id } = await register(3, 2);
    const digest = ethers.keccak256(ethers.toUtf8Bytes("withdraw 5 AERE"));
    // combine with only ONE share -> interpolates to the wrong secret
    const badSig = tss.thresholdSign(ethers.getBytes(digest), [dkg.shares[0]]);
    expect(await reg.verifyThresholdSignature(id, digest, badSig)).to.equal(false);
  });

  it("REJECTS a signature over a different message, a wrong committee, malleable high-s, and bad v", async function () {
    const { dkg, id } = await register(3, 2);
    const digest = ethers.keccak256(ethers.toUtf8Bytes("msg A"));
    const other = ethers.keccak256(ethers.toUtf8Bytes("msg B"));
    const sig = tss.thresholdSign(ethers.getBytes(digest), [dkg.shares[0], dkg.shares[1]]);
    // wrong message
    expect(await reg.verifyThresholdSignature(id, other, sig)).to.equal(false);
    // unknown committee id
    expect(await reg.verifyThresholdSignature(999, digest, sig)).to.equal(false);
    // high-s malleated copy must be rejected (s -> N-s, flip v)
    const r = ethers.dataSlice(sig, 0, 32);
    const s = BigInt(ethers.dataSlice(sig, 32, 64));
    const v = parseInt(ethers.dataSlice(sig, 64, 65), 16);
    const highS = tss.N - s;
    const flippedV = v === 27 ? 28 : 27;
    const malleable = ethers.concat([r, ethers.toBeHex(highS, 32), ethers.toBeHex(flippedV, 1)]);
    expect(await reg.verifyThresholdSignature(id, digest, malleable)).to.equal(false);
    // bad v (=29)
    const badV = ethers.concat([ethers.dataSlice(sig, 0, 64), "0x1d"]);
    expect(await reg.verifyThresholdSignature(id, digest, badV)).to.equal(false);
    // wrong length
    expect(await reg.verifyThresholdSignature(id, digest, ethers.dataSlice(sig, 0, 64))).to.equal(false);
  });

  it("reshares membership (2of3 -> 3of4) keeping the SAME group key, authorized by the current committee", async function () {
    const { dkg, id } = await register(3, 2);
    const newMembers = makeMembers(4);
    const newT = 3;
    const c0 = await reg.getCommittee(id);
    const intent = reshareIntent(chainId, await reg.getAddress(), id, c0.epoch, c0.actionNonce, newMembers, newT);
    const authSig = tss.thresholdSign(ethers.getBytes(intent), [dkg.shares[0], dkg.shares[1]]);

    await expect(reg.reshareCommittee(id, newMembers, newT, authSig))
      .to.emit(reg, "CommitteeReshared")
      .withArgs(id, 2, newT, 4, ethers.keccak256(abi.encode(["address[]"], [newMembers])));

    const c = await reg.getCommittee(id);
    expect(c.threshold).to.equal(3);
    expect(c.size).to.equal(4);
    expect(c.epoch).to.equal(2n);
    // SAME group key preserved
    expect(c.groupAddress.toLowerCase()).to.equal(dkg.groupAddress.toLowerCase());
    // new roster installed, old roster cleared
    const onchain = await reg.getMembers(id);
    expect(onchain.map((x) => x.toLowerCase())).to.deep.equal(newMembers.map((x) => x.toLowerCase()));

    // The SAME group key still signs (shares unchanged in this reshare model).
    const digest = ethers.keccak256(ethers.toUtf8Bytes("post-reshare"));
    const sig = tss.thresholdSign(ethers.getBytes(digest), [dkg.shares[0], dkg.shares[2]]);
    expect(await reg.verifyThresholdSignature(id, digest, sig)).to.equal(true);
  });

  it("REJECTS a reshare not authorized by the current committee (forged auth)", async function () {
    const { id } = await register(3, 2);
    const foreign = tss.runDkg(3, 2); // a DIFFERENT group key
    const newMembers = makeMembers(3);
    const c0 = await reg.getCommittee(id);
    const intent = reshareIntent(chainId, await reg.getAddress(), id, c0.epoch, c0.actionNonce, newMembers, 2);
    const badAuth = tss.thresholdSign(ethers.getBytes(intent), [foreign.shares[0], foreign.shares[1]]);
    await expect(reg.reshareCommittee(id, newMembers, 2, badAuth)).to.be.revertedWith("bad committee auth");
  });

  it("prevents replaying a reshare auth after the epoch/nonce advances", async function () {
    const { dkg, id } = await register(3, 2);
    const newMembers1 = makeMembers(3);
    const c0 = await reg.getCommittee(id);
    const intent1 = reshareIntent(chainId, await reg.getAddress(), id, c0.epoch, c0.actionNonce, newMembers1, 2);
    const auth1 = tss.thresholdSign(ethers.getBytes(intent1), [dkg.shares[0], dkg.shares[1]]);
    await reg.reshareCommittee(id, newMembers1, 2, auth1);
    // replay the exact same call+sig -> epoch/nonce changed so intent no longer matches
    await expect(reg.reshareCommittee(id, newMembers1, 2, auth1)).to.be.revertedWith("bad committee auth");
  });

  it("rotates to a NEW group key (post-compromise), authorized by the OLD committee", async function () {
    const { dkg, id } = await register(3, 2);
    const fresh = tss.runDkg(5, 3); // brand new committee + key
    const newMembers = makeMembers(5);
    const c0 = await reg.getCommittee(id);
    const intent = rotateIntent(chainId, await reg.getAddress(), id, c0.epoch, c0.actionNonce, fresh.groupPubKey64, newMembers, 3);
    // authorized by the CURRENT (old) committee
    const authSig = tss.thresholdSign(ethers.getBytes(intent), [dkg.shares[1], dkg.shares[2]]);
    await expect(reg.rotateKey(id, fresh.groupPubKey64, newMembers, 3, authSig)).to.emit(reg, "KeyRotated");

    const c = await reg.getCommittee(id);
    expect(c.groupAddress.toLowerCase()).to.equal(fresh.groupAddress.toLowerCase());
    expect(c.threshold).to.equal(3);
    expect(c.size).to.equal(5);
    expect(c.epoch).to.equal(2n);

    // OLD key no longer verifies; NEW key does.
    const digest = ethers.keccak256(ethers.toUtf8Bytes("after rotation"));
    const oldSig = tss.thresholdSign(ethers.getBytes(digest), [dkg.shares[0], dkg.shares[1]]);
    const newSig = tss.thresholdSign(ethers.getBytes(digest), [fresh.shares[0], fresh.shares[1], fresh.shares[2]]);
    expect(await reg.verifyThresholdSignature(id, digest, oldSig)).to.equal(false);
    expect(await reg.verifyThresholdSignature(id, digest, newSig)).to.equal(true);
  });

  it("rejects rotate to the SAME key (must actually change)", async function () {
    const { dkg, id } = await register(3, 2);
    const newMembers = makeMembers(3);
    const c0 = await reg.getCommittee(id);
    const intent = rotateIntent(chainId, await reg.getAddress(), id, c0.epoch, c0.actionNonce, dkg.groupPubKey64, newMembers, 2);
    const authSig = tss.thresholdSign(ethers.getBytes(intent), [dkg.shares[0], dkg.shares[1]]);
    await expect(reg.rotateKey(id, dkg.groupPubKey64, newMembers, 2, authSig)).to.be.revertedWith("key unchanged");
  });

  it("slashes a member via committee threshold signature, and marks the member", async function () {
    const { dkg, members, id } = await register(3, 2);
    const target = members[2];
    const evidence = ethers.keccak256(ethers.toUtf8Bytes("equivocation proof blob"));
    const c0 = await reg.getCommittee(id);
    const intent = slashIntent(chainId, await reg.getAddress(), id, c0.epoch, c0.actionNonce, target, evidence);
    const authSig = tss.thresholdSign(ethers.getBytes(intent), [dkg.shares[0], dkg.shares[1]]);

    expect(await reg.slashed(id, target)).to.equal(false);
    await expect(reg.slashMember(id, target, evidence, authSig))
      .to.emit(reg, "MemberSlashed")
      .withArgs(id, target, evidence, deployer.address);
    expect(await reg.slashed(id, target)).to.equal(true);
    // double slash blocked
    await expect(reg.slashMember(id, target, evidence, authSig)).to.be.revertedWith("already slashed");
  });

  it("lets the guardian slash and pause/unpause; non-guardian cannot", async function () {
    const { members, id } = await register(3, 2, guardian.address);
    // guardian slash (no threshold sig needed)
    const target = members[1];
    await expect(reg.connect(guardian).slashMember(id, target, ethers.ZeroHash, "0x"))
      .to.emit(reg, "MemberSlashed");
    expect(await reg.slashed(id, target)).to.equal(true);

    // guardian pause -> verification returns false even for a valid sig
    await expect(reg.connect(guardian).pauseCommittee(id)).to.emit(reg, "CommitteePaused");
    // outsider cannot unpause
    await expect(reg.connect(outsider).unpauseCommittee(id)).to.be.revertedWith("not guardian");
    await expect(reg.connect(guardian).unpauseCommittee(id)).to.emit(reg, "CommitteeUnpaused");
  });

  it("a PAUSED committee fails verification; ACTIVE again after unpause", async function () {
    const { dkg, id } = await register(3, 2, guardian.address);
    const digest = ethers.keccak256(ethers.toUtf8Bytes("x"));
    const sig = tss.thresholdSign(ethers.getBytes(digest), [dkg.shares[0], dkg.shares[1]]);
    expect(await reg.verifyThresholdSignature(id, digest, sig)).to.equal(true);
    await reg.connect(guardian).pauseCommittee(id);
    expect(await reg.verifyThresholdSignature(id, digest, sig)).to.equal(false);
    await reg.connect(guardian).unpauseCommittee(id);
    expect(await reg.verifyThresholdSignature(id, digest, sig)).to.equal(true);
  });

  it("rejects bad registration params (bad t, empty members, dup member, bad pubkey length)", async function () {
    const dkg = tss.runDkg(3, 2);
    const m = makeMembers(3);
    await expect(reg.registerCommittee(m, 0, dkg.groupPubKey64, ethers.ZeroAddress)).to.be.revertedWith("bad t");
    await expect(reg.registerCommittee(m, 4, dkg.groupPubKey64, ethers.ZeroAddress)).to.be.revertedWith("bad t");
    await expect(reg.registerCommittee([], 1, dkg.groupPubKey64, ethers.ZeroAddress)).to.be.revertedWith("bad n");
    const dup = [m[0], m[0], m[1]];
    await expect(reg.registerCommittee(dup, 2, dkg.groupPubKey64, ethers.ZeroAddress)).to.be.revertedWith("dup member");
    // 65-byte (with 0x04 prefix) pubkey is wrong -- contract wants raw 64-byte x||y
    const bad65 = ethers.concat(["0x04", dkg.groupPubKey64]);
    await expect(reg.registerCommittee(m, 2, bad65, ethers.ZeroAddress)).to.be.revertedWith("bad groupPubKey");
  });
});
