// AerePQCTxAccount — a post-quantum-authorized ERC-4337 v0.7 smart account (deploy-today reference
// path for AIP-PQ-TX). The account's sole owner is a NIST PQC public key, and every fund-moving
// authorization (validateUserOp / executeWithPQC) is decided by the LIVE native PQC precompiles on
// chain 2800: Falcon-512 0x0AE1, Falcon-1024 0x0AE2, ML-DSA-44 0x0AE3, SLH-DSA-128s 0x0AE4. There is
// no ECDSA fallback; a userOp may be relayed by any bundler/EOA and only the PQC key holder can
// authorize it.
//
// Two verification backends are exercised at 0x0AE1..0x0AE4 (the exact pattern used by
// AerePQCSocialRecoveryModule.test.js and AereThresholdAccount.test.js):
//   1. MockPQCPrecompile — the repo's faithful mock (accept iff commitment == keccak256(pk||msg)),
//      parsing the input at the precompiles' SPEC offsets, used for the account LOGIC across all four
//      schemes (accept + reject branches, domain separation, nonce/replay).
//   2. RealFalcon512PrecompileAdapter — routes the precompile staticcall into the NIST-KAT-proven
//      AereFalcon512Verifier, so a subset of tests drive GENUINE Falcon-512 signatures (produced by
//      @noble/post-quantum, byte-compatible with AERE's live precompile) through validateUserOp.
//
// Run: npx hardhat test test/AerePQCTxAccount.test.js
const { expect } = require("chai");
const { ethers, network } = require("hardhat");
const path = require("path");
const { pathToFileURL } = require("url");

const PRECOMPILE = {
  1: ethers.getAddress("0x0000000000000000000000000000000000000ae1"),
  2: ethers.getAddress("0x0000000000000000000000000000000000000ae2"),
  3: ethers.getAddress("0x0000000000000000000000000000000000000ae3"),
  4: ethers.getAddress("0x0000000000000000000000000000000000000ae4"),
};
const META = {
  1: { pkLen: 897, pkHeader: 0x09, esigHeader: 0x29, falcon: true, name: "Falcon-512" },
  2: { pkLen: 1793, pkHeader: 0x0a, esigHeader: 0x2a, falcon: true, name: "Falcon-1024" },
  3: { pkLen: 1312, sigLen: 2420, falcon: false, name: "ML-DSA-44" },
  4: { pkLen: 32, sigLen: 7856, falcon: false, name: "SLH-DSA-128s" },
};
const ZERO = ethers.ZeroHash;

// Domain separators — MUST match AerePQCTxAccount exactly.
const USEROP_DOMAIN = ethers.id("AerePQCTxAccount.validateUserOp.v1");
const EXEC_DOMAIN = ethers.id("AerePQCTxAccount.executeWithPQC.v1");
// Fund-moving paths verify over abi.encodePacked(domain, hash) = 64 bytes; EIP-1271 over raw 32 bytes.
const domained = (domain, hash) => ethers.concat([domain, hash]);

// ---- deterministic seeded RNG (keccak hash-chain) --------------------------
function makeRng(seedText) {
  let seed = ethers.keccak256(ethers.toUtf8Bytes(seedText));
  return {
    bytes(n) {
      const out = new Uint8Array(n);
      let i = 0;
      while (i < n) {
        seed = ethers.keccak256(seed);
        const chunk = ethers.getBytes(seed);
        for (let j = 0; j < chunk.length && i < n; j++, i++) out[i] = chunk[j];
      }
      return out;
    },
  };
}
function makePubKey(scheme, rng) {
  const m = META[scheme];
  const raw = rng.bytes(m.pkLen);
  if (m.falcon) raw[0] = m.pkHeader;
  return ethers.hexlify(raw);
}
function commitmentFor(pubKey, message) {
  return ethers.keccak256(ethers.concat([pubKey, message]));
}
// A genuine MOCK envelope binding to `message` under `pubKey` (the mock at 0x0AE1..4 accepts it).
// Envelope shape matches the LIVE precompiles: Falcon = nonce(40)||esig, fixed = raw sig.
function genuineMock(scheme, pubKey, message, rng) {
  const m = META[scheme];
  const commitment = commitmentFor(pubKey, message);
  if (m.falcon) {
    const nonce = rng.bytes(40);
    const esig = ethers.concat([new Uint8Array([m.esigHeader]), commitment]);
    return ethers.hexlify(ethers.concat([nonce, esig]));
  }
  const sig = new Uint8Array(m.sigLen);
  sig.set(ethers.getBytes(commitment), 0);
  return ethers.hexlify(sig);
}
// Flip one byte in the commitment region -> a forgery (the mock rejects it).
function tamperMock(scheme, sigHex) {
  const b = ethers.getBytes(sigHex);
  b[META[scheme].falcon ? 41 : 0] ^= 0x01;
  return ethers.hexlify(b);
}
// PackedUserOperation (v0.7) tuple with `sig` in the signature slot; only `signature` matters here.
function userOp(sender, sig) {
  return [sender, 0, "0x", "0x", ZERO, 0, ZERO, "0x", sig];
}

describe("AerePQCTxAccount (post-quantum-authorized ERC-4337 account)", function () {
  let deployer, ep, other, relayer, stranger;
  let mockCode;

  before(async function () {
    [deployer, ep, other, relayer, stranger] = await ethers.getSigners();
    mockCode = {};
    const F = await ethers.getContractFactory("MockPQCPrecompile");
    for (const s of [1, 2, 3, 4]) {
      const m = await F.deploy(s);
      await m.waitForDeployment();
      mockCode[s] = await ethers.provider.getCode(await m.getAddress());
    }
  });

  async function installMocks() {
    for (const s of [1, 2, 3, 4]) {
      await network.provider.send("hardhat_setCode", [PRECOMPILE[s], mockCode[s]]);
    }
  }

  async function deployAccount(scheme, rng, fund = true) {
    await installMocks();
    const pk = makePubKey(scheme, rng);
    const F = await ethers.getContractFactory("AerePQCTxAccount");
    const account = await F.deploy(ep.address, scheme, pk);
    await account.waitForDeployment();
    if (fund) {
      await (await deployer.sendTransaction({ to: await account.getAddress(), value: ethers.parseEther("1") })).wait();
    }
    return { account, pk, F };
  }

  // =========================================================================
  //  Construction / config
  // =========================================================================
  describe("construction", function () {
    it("wires entryPoint, scheme, pubkey and the right precompile for each scheme", async function () {
      const rng = makeRng("ctor-wire");
      for (const scheme of [1, 2, 3, 4]) {
        const { account, pk } = await deployAccount(scheme, rng, false);
        expect(await account.entryPoint()).to.equal(ep.address);
        expect(await account.pqcScheme()).to.equal(scheme);
        expect((await account.pqcPubKey()).toLowerCase()).to.equal(pk.toLowerCase());
        expect(await account.precompileAddress()).to.equal(PRECOMPILE[scheme]);
        expect(await account.execNonce()).to.equal(0n);
      }
    });

    it("rejects a zero EntryPoint", async function () {
      const rng = makeRng("ctor-zeroep");
      const F = await ethers.getContractFactory("AerePQCTxAccount");
      await expect(F.deploy(ethers.ZeroAddress, 1, makePubKey(1, rng)))
        .to.be.revertedWithCustomError(F, "InvalidEntryPoint");
    });

    it("rejects a malformed public key (wrong length or wrong Falcon header)", async function () {
      const rng = makeRng("ctor-badpk");
      const F = await ethers.getContractFactory("AerePQCTxAccount");
      // Falcon-512 wrong length
      await expect(F.deploy(ep.address, 1, "0x09" + "ab".repeat(10)))
        .to.be.revertedWithCustomError(F, "InvalidPqcPubKey");
      // Falcon-512 wrong header (0x08)
      await expect(F.deploy(ep.address, 1, "0x08" + "ab".repeat(896)))
        .to.be.revertedWithCustomError(F, "InvalidPqcPubKey");
      // ML-DSA-44 wrong length
      await expect(F.deploy(ep.address, 3, "0x" + "ab".repeat(100)))
        .to.be.revertedWithCustomError(F, "InvalidPqcPubKey");
    });

    it("rejects an unknown scheme", async function () {
      const F = await ethers.getContractFactory("AerePQCTxAccount");
      await expect(F.deploy(ep.address, 5, "0x" + "ab".repeat(32)))
        .to.be.revertedWithCustomError(F, "InvalidScheme");
    });
  });

  // =========================================================================
  //  validateUserOp — mock verifier, ALL FOUR schemes (accept + reject)
  // =========================================================================
  describe("validateUserOp (mock verifier, all schemes)", function () {
    for (const scheme of [1, 2, 3, 4]) {
      it(`ACCEPT/REJECT for ${META[scheme].name} (scheme ${scheme}) via 0x0AE${scheme}`, async function () {
        const rng = makeRng(`vuo-${scheme}`);
        const { account, pk } = await deployAccount(scheme, rng, false);
        const userOpHash = ethers.id(`userop-${scheme}`);
        const message = domained(USEROP_DOMAIN, userOpHash);

        // Valid PQC signature over USEROP_DOMAIN||userOpHash -> validationData 0.
        const good = genuineMock(scheme, pk, message, rng);
        const okVal = await account
          .connect(ep)
          .validateUserOp.staticCall(userOp(await account.getAddress(), good), userOpHash, 0);
        expect(okVal).to.equal(0n);

        // Tampered signature -> SIG_VALIDATION_FAILED (1), NOT a revert.
        const bad = tamperMock(scheme, good);
        const badVal = await account
          .connect(ep)
          .validateUserOp.staticCall(userOp(await account.getAddress(), bad), userOpHash, 0);
        expect(badVal).to.equal(1n);
      });
    }

    it("REJECT: a signature over a DIFFERENT userOpHash does not validate (binds the op)", async function () {
      const rng = makeRng("vuo-wronghash");
      const { account, pk } = await deployAccount(1, rng, false);
      const signed = genuineMock(1, pk, domained(USEROP_DOMAIN, ethers.id("op-A")), rng);
      // Submit it against op-B's hash.
      const val = await account
        .connect(ep)
        .validateUserOp.staticCall(userOp(await account.getAddress(), signed), ethers.id("op-B"), 0);
      expect(val).to.equal(1n);
    });

    it("validateUserOp reverts if not called by the EntryPoint", async function () {
      const rng = makeRng("vuo-auth");
      const { account, pk } = await deployAccount(1, rng, false);
      const h = ethers.id("op");
      const sig = genuineMock(1, pk, domained(USEROP_DOMAIN, h), rng);
      await expect(
        account.connect(other).validateUserOp(userOp(await account.getAddress(), sig), h, 0)
      ).to.be.revertedWithCustomError(account, "Unauthorized");
    });

    it("validateUserOp forwards missingAccountFunds to the EntryPoint as prefund", async function () {
      const rng = makeRng("vuo-prefund");
      const { account, pk } = await deployAccount(1, rng, true); // account funded with 1 ETH
      const h = ethers.id("op-prefund");
      const sig = genuineMock(1, pk, domained(USEROP_DOMAIN, h), rng);
      const prefund = ethers.parseEther("0.2");
      const before = await ethers.provider.getBalance(ep.address);
      // Real (state-changing) call from the EntryPoint: returns 0 AND sends the prefund.
      const rc = await (
        await account.connect(ep).validateUserOp(userOp(await account.getAddress(), sig), h, prefund)
      ).wait();
      const after = await ethers.provider.getBalance(ep.address);
      const gasCost = rc.gasUsed * rc.gasPrice;
      // ep received `prefund` minus the gas it paid to send the tx.
      expect(after - before + gasCost).to.equal(prefund);
    });
  });

  // =========================================================================
  //  executeWithPQC (self-relay) + EIP-1271 + domain separation
  // =========================================================================
  describe("executeWithPQC / EIP-1271 / domain separation (mock verifier)", function () {
    it("executeWithPQC moves value on a valid signature and bumps execNonce", async function () {
      const rng = makeRng("exec-happy");
      const { account, pk } = await deployAccount(1, rng, true);
      const target = other.address;
      const value = ethers.parseEther("0.25");
      const data = "0x";
      const n = await account.execNonce();
      const challenge = await account.getExecuteChallenge(target, value, data, n);
      const sig = genuineMock(1, pk, domained(EXEC_DOMAIN, challenge), rng);

      const before = await ethers.provider.getBalance(target);
      await (await account.connect(relayer).executeWithPQC(sig, target, value, data)).wait();
      expect((await ethers.provider.getBalance(target)) - before).to.equal(value);
      expect(await account.execNonce()).to.equal(n + 1n);
    });

    it("executeWithPQC reverts on an invalid signature; no value moves, nonce unchanged", async function () {
      const rng = makeRng("exec-bad");
      const { account, pk } = await deployAccount(3, rng, true); // ML-DSA-44 owner
      const target = other.address;
      const value = ethers.parseEther("0.1");
      const n = await account.execNonce();
      const challenge = await account.getExecuteChallenge(target, value, "0x", n);
      const bad = tamperMock(3, genuineMock(3, pk, domained(EXEC_DOMAIN, challenge), rng));
      await expect(account.executeWithPQC(bad, target, value, "0x"))
        .to.be.revertedWithCustomError(account, "InvalidSignature");
      expect(await account.execNonce()).to.equal(n);
    });

    it("executeWithPQC legs are single-use: replaying the same signature fails (nonce advanced)", async function () {
      const rng = makeRng("exec-replay");
      const { account, pk } = await deployAccount(1, rng, true);
      const target = other.address;
      const value = ethers.parseEther("0.05");
      const n = await account.execNonce();
      const challenge = await account.getExecuteChallenge(target, value, "0x", n);
      const sig = genuineMock(1, pk, domained(EXEC_DOMAIN, challenge), rng);
      await (await account.executeWithPQC(sig, target, value, "0x")).wait();
      // execNonce is now n+1, so the same signature (bound to nonce n) no longer authorizes.
      await expect(account.executeWithPQC(sig, target, value, "0x"))
        .to.be.revertedWithCustomError(account, "InvalidSignature");
    });

    it("EIP-1271 returns the magic value for a valid raw-hash signature, 0 for a tampered one", async function () {
      const rng = makeRng("1271");
      const { account, pk } = await deployAccount(1, rng, false);
      const hash = ethers.id("post-quantum login to dapp.example");
      const sig = genuineMock(1, pk, hash, rng); // raw 32-byte message
      expect(await account.isValidSignature(hash, sig)).to.equal("0x1626ba7e");
      expect(await account.isValidSignature(hash, tamperMock(1, sig))).to.equal("0x00000000");
    });

    it("DOMAIN SEPARATION: an EIP-1271 login signature cannot be replayed as an executeWithPQC spend", async function () {
      const rng = makeRng("domain-sep");
      const { account, pk } = await deployAccount(1, rng, true);
      const target = other.address;
      const value = ethers.parseEther("0.5");
      const data = "0x";
      const n = await account.execNonce();
      // Attacker computes the exact execute challenge (all inputs public; execNonce is public).
      const challenge = await account.getExecuteChallenge(target, value, data, n);
      // User is phished into signing `challenge` as a RAW login hash (the EIP-1271 surface).
      const loginSig = genuineMock(1, pk, challenge, rng); // commitment over the RAW 32-byte challenge
      // That login signature IS accepted by the raw-hash isValidSignature surface...
      expect(await account.isValidSignature(challenge, loginSig)).to.equal("0x1626ba7e");
      // ...but must NOT authorize a spend: executeWithPQC verifies over (EXEC_DOMAIN || challenge) = 64 bytes.
      await expect(account.executeWithPQC(loginSig, target, value, data))
        .to.be.revertedWithCustomError(account, "InvalidSignature");
      expect(await account.execNonce()).to.equal(n); // nothing moved
    });
  });

  // =========================================================================
  //  EntryPoint-gated dispatch
  // =========================================================================
  describe("executeFromEntryPoint / executeBatchFromEntryPoint", function () {
    it("executeFromEntryPoint is EntryPoint-gated and dispatches when the EntryPoint calls", async function () {
      const rng = makeRng("efe");
      const { account } = await deployAccount(1, rng, true);
      await expect(account.connect(other).executeFromEntryPoint(other.address, 0, "0x"))
        .to.be.revertedWithCustomError(account, "Unauthorized");
      const value = ethers.parseEther("0.3");
      const before = await ethers.provider.getBalance(other.address);
      await (await account.connect(ep).executeFromEntryPoint(other.address, value, "0x")).wait();
      expect((await ethers.provider.getBalance(other.address)) - before).to.equal(value);
    });

    it("executeBatchFromEntryPoint dispatches every call and is EntryPoint-gated", async function () {
      const rng = makeRng("ebe");
      const { account } = await deployAccount(1, rng, true);
      const calls = [
        [stranger.address, ethers.parseEther("0.1"), "0x"],
        [other.address, ethers.parseEther("0.2"), "0x"],
      ];
      await expect(account.connect(other).executeBatchFromEntryPoint(calls))
        .to.be.revertedWithCustomError(account, "Unauthorized");
      const b1 = await ethers.provider.getBalance(stranger.address);
      const b2 = await ethers.provider.getBalance(other.address);
      await (await account.connect(ep).executeBatchFromEntryPoint(calls)).wait();
      expect((await ethers.provider.getBalance(stranger.address)) - b1).to.equal(ethers.parseEther("0.1"));
      expect((await ethers.provider.getBalance(other.address)) - b2).to.equal(ethers.parseEther("0.2"));
    });
  });

  // =========================================================================
  //  REAL Falcon-512 crypto: genuine signatures via @noble, verified by the
  //  NIST-KAT-proven AereFalcon512Verifier behind a precompile-format adapter.
  // =========================================================================
  describe("real Falcon-512 authorization (KAT-proven verifier, no mock)", function () {
    let falcon512, adapterCode;
    // Pure-Solidity Falcon-512 verify (stand-in for the native precompile) is ~10.7M gas; give the
    // heavy nested-staticcall paths ample headroom (the hardhat network permits a large block gas
    // limit). TEST-HARNESS accommodation for the expensive stand-in, not a contract concern.
    const HEAVY = { gasLimit: 400_000_000 };

    before(async function () {
      const nobleFalcon = path.resolve(__dirname, "../../sdk-js/node_modules/@noble/post-quantum/falcon.js");
      ({ falcon512 } = await import(pathToFileURL(nobleFalcon).href));

      const V = await ethers.getContractFactory("AereFalcon512Verifier");
      const verifier = await V.deploy();
      await verifier.waitForDeployment();
      const A = await ethers.getContractFactory("RealFalcon512PrecompileAdapter");
      const adapter = await A.deploy(await verifier.getAddress());
      await adapter.waitForDeployment();
      adapterCode = await ethers.provider.getCode(await adapter.getAddress());
    });

    // noble detached (0x39 || nonce(40) || compressed) -> AERE envelope (nonce(40) || 0x29 || compressed).
    function realEnvelope(detached) {
      const nonce = detached.slice(1, 41);
      const compressed = detached.slice(41);
      const esig = new Uint8Array(1 + compressed.length);
      esig[0] = 0x29;
      esig.set(compressed, 1);
      return ethers.hexlify(ethers.concat([nonce, esig]));
    }

    async function installRealFalconAt0AE1() {
      await network.provider.send("hardhat_setCode", [PRECOMPILE[1], adapterCode]);
    }

    async function deployRealAccount(seedFill) {
      await installRealFalconAt0AE1();
      const kp = falcon512.keygen(new Uint8Array(48).fill(seedFill));
      const pk = ethers.hexlify(kp.publicKey);
      const F = await ethers.getContractFactory("AerePQCTxAccount");
      const account = await F.deploy(ep.address, 1, pk);
      await account.waitForDeployment();
      return { account, kp };
    }

    it("validateUserOp returns 0 for a GENUINE Falcon-512 signature over the domained userOpHash", async function () {
      const { account, kp } = await deployRealAccount(0x11);
      const userOpHash = ethers.id("real-userop-accept");
      const message = ethers.getBytes(domained(USEROP_DOMAIN, userOpHash)); // 64 bytes
      // Falcon signing is randomized; re-sign until the full account->adapter->verifier path accepts.
      let accepted = false;
      for (let attempt = 0; attempt < 8 && !accepted; attempt++) {
        const env = realEnvelope(falcon512.sign(message, kp.secretKey));
        const val = await account
          .connect(ep)
          .validateUserOp.staticCall(userOp(await account.getAddress(), env), userOpHash, 0, HEAVY);
        if (val === 0n) accepted = true;
      }
      expect(accepted, "a genuine Falcon-512 signature should validate").to.equal(true);
    });

    it("validateUserOp returns 1 (SIG_VALIDATION_FAILED) for a tampered Falcon-512 signature", async function () {
      const { account, kp } = await deployRealAccount(0x22);
      const userOpHash = ethers.id("real-userop-reject");
      const message = ethers.getBytes(domained(USEROP_DOMAIN, userOpHash));
      const env = ethers.getBytes(realEnvelope(falcon512.sign(message, kp.secretKey)));
      env[50] ^= 0x01; // flip a byte inside the compressed signature body -> forgery
      const val = await account
        .connect(ep)
        .validateUserOp.staticCall(userOp(await account.getAddress(), ethers.hexlify(env)), userOpHash, 0, HEAVY);
      expect(val).to.equal(1n);
    });
  });
});