const { expect } = require('chai');
const { ethers } = require('hardhat');

/**
 * STARK Proof Verification Tests for Vaultfire Belief Attestation
 *
 * These tests verify the STARK ZK proof system integration with DilithiumAttestor.
 * They test both the BeliefAttestationVerifier contract and end-to-end proof verification.
 *
 * Test Structure:
 * 1. BeliefAttestationVerifier standalone tests
 * 2. DilithiumAttestor with ZK enabled (full STARK mode)
 * 3. Integration tests with BeliefOracle
 */

describe('STARK Proof Verification', () => {
  let deployer, prover, origin;
  let verifier, attestor;

  beforeEach(async () => {
    [deployer, prover, origin] = await ethers.getSigners();

    // Deploy BeliefAttestationVerifier (STARK verifier)
    const BeliefVerifier = await ethers.getContractFactory('BeliefAttestationVerifier');
    verifier = await BeliefVerifier.deploy();
    await verifier.waitForDeployment();
  });

  describe('BeliefAttestationVerifier', () => {
    it('should return correct proof system ID', async () => {
      const proofSystemId = await verifier.getProofSystemId();
      expect(proofSystemId).to.equal('RISC0-STARK-BeliefAttestation-v1.0');
    });

    it('should return correct public inputs count', async () => {
      const count = await verifier.getPublicInputsCount();
      expect(count).to.equal(4);
    });

    it('should return min belief threshold', async () => {
      const threshold = await verifier.getMinBeliefThreshold();
      expect(threshold).to.equal(8000); // 80%
    });

    it('should verify valid STARK proof structure', async () => {
      const beliefMessage = 'I believe in Vaultfire mission';
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes(beliefMessage));

      // Construct public inputs
      const publicInputs = [
        beliefHash,                      // beliefHash
        BigInt(prover.address),         // proverAddress
        0n,                              // epoch
        0n                               // moduleID
      ];

      // Mock STARK proof (in production, this would be generated by Risc0/StarkWare/Miden)
      // For testing, we use a placeholder that passes basic validation
      const mockStarkProof = ethers.randomBytes(128); // Typical STARK proof size

      // Verify proof
      const tx = await verifier.verifyProof(mockStarkProof, publicInputs);
      const receipt = await tx.wait();

      // Check ProofVerified event
      const event = receipt.logs.find(log => {
        try {
          return verifier.interface.parseLog(log)?.name === 'ProofVerified';
        } catch {
          return false;
        }
      });

      expect(event).to.not.be.undefined;
      const parsed = verifier.interface.parseLog(event);
      expect(parsed.args.beliefHash).to.equal(beliefHash);
      expect(parsed.args.proverAddress).to.equal(prover.address);
    });

    it('should reject proof with invalid public inputs count', async () => {
      const mockProof = ethers.randomBytes(128);
      const invalidPublicInputs = [1n, 2n]; // Only 2 inputs instead of 4

      await expect(verifier.verifyProof(mockProof, invalidPublicInputs))
        .to.be.revertedWith('Invalid public inputs count');
    });

    it('should reject proof with zero belief hash', async () => {
      const mockProof = ethers.randomBytes(128);
      const invalidPublicInputs = [
        ethers.ZeroHash,           // Invalid: zero belief hash
        BigInt(prover.address),
        0n,
        0n
      ];

      await expect(verifier.verifyProof(mockProof, invalidPublicInputs))
        .to.be.revertedWith('Invalid belief hash');
    });

    it('should reject proof with zero prover address', async () => {
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('test'));
      const mockProof = ethers.randomBytes(128);
      const invalidPublicInputs = [
        beliefHash,
        0n,                        // Invalid: zero address
        0n,
        0n
      ];

      await expect(verifier.verifyProof(mockProof, invalidPublicInputs))
        .to.be.revertedWith('Invalid prover address');
    });

    it('should reject proof with invalid epoch', async () => {
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('test'));
      const mockProof = ethers.randomBytes(128);
      const invalidPublicInputs = [
        beliefHash,
        BigInt(prover.address),
        BigInt(2) ** BigInt(33),   // Invalid: exceeds MAX_EPOCH (2^32 - 1)
        0n
      ];

      await expect(verifier.verifyProof(mockProof, invalidPublicInputs))
        .to.be.revertedWith('Invalid epoch');
    });

    it('should reject empty proof', async () => {
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('test'));
      const emptyProof = '0x';
      const publicInputs = [
        beliefHash,
        BigInt(prover.address),
        0n,
        0n
      ];

      await expect(verifier.verifyProof(emptyProof, publicInputs))
        .to.be.revertedWith('Empty proof');
    });

    it('should reject proof that is too short', async () => {
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('test'));
      const shortProof = ethers.randomBytes(16); // Less than 32 bytes minimum
      const publicInputs = [
        beliefHash,
        BigInt(prover.address),
        0n,
        0n
      ];

      await expect(verifier.verifyProof(shortProof, publicInputs))
        .to.be.revertedWith('Proof too short');
    });
  });

  describe('DilithiumAttestor with STARK ZK Enabled', () => {
    beforeEach(async () => {
      // Deploy DilithiumAttestor with ZK enabled (full STARK mode)
      const DilithiumAttestor = await ethers.getContractFactory('DilithiumAttestor');
      attestor = await DilithiumAttestor.deploy(
        origin.address,              // origin (signature authority)
        true,                        // zkEnabled = true (FULL ZK MODE)
        await verifier.getAddress()  // STARK verifier address
      );
      await attestor.waitForDeployment();
    });

    it('should verify zkEnabled is true', async () => {
      expect(await attestor.zkEnabled()).to.be.true;
    });

    it('should verify verifierAddress is set', async () => {
      expect(await attestor.verifierAddress()).to.equal(await verifier.getAddress());
    });

    it('should verify valid STARK proof + signature', async () => {
      const beliefMessage = 'I am committed to Vaultfire values';
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes(beliefMessage));

      // Origin signs the beliefHash
      const ethSigned = ethers.hashMessage(ethers.getBytes(beliefHash));
      const originSignature = await origin.signMessage(ethers.getBytes(beliefHash));

      // Generate mock STARK proof
      const mockStarkProof = ethers.randomBytes(128);

      // Encode zkProofBundle: (starkProof, originSignature)
      const zkProofBundle = ethers.AbiCoder.defaultAbiCoder().encode(
        ['bytes', 'bytes'],
        [mockStarkProof, originSignature]
      );

      // Attest belief
      const tx = await attestor.connect(prover).attestBelief(beliefHash, zkProofBundle);
      const receipt = await tx.wait();

      // Verify BeliefAttested event
      const event = receipt.logs.find(log => {
        try {
          return attestor.interface.parseLog(log)?.name === 'BeliefAttested';
        } catch {
          return false;
        }
      });

      expect(event).to.not.be.undefined;
      const parsed = attestor.interface.parseLog(event);
      expect(parsed.args.beliefHash).to.equal(beliefHash);
      expect(parsed.args.prover).to.equal(prover.address);
      expect(parsed.args.zkVerified).to.be.true; // ZK was verified!

      // Verify belief is attested
      expect(await attestor.isBeliefSovereign(beliefHash)).to.be.true;
    });

    it('should reject invalid STARK proof', async () => {
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('test'));
      const originSignature = await origin.signMessage(ethers.getBytes(beliefHash));

      // Invalid STARK proof (too short)
      const invalidProof = ethers.randomBytes(8);

      const zkProofBundle = ethers.AbiCoder.defaultAbiCoder().encode(
        ['bytes', 'bytes'],
        [invalidProof, originSignature]
      );

      await expect(attestor.connect(prover).attestBelief(beliefHash, zkProofBundle))
        .to.be.revertedWith('Proof too short');
    });

    it('should reject valid STARK proof but invalid signature', async () => {
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('test'));

      // Valid STARK proof
      const validStarkProof = ethers.randomBytes(128);

      // Invalid signature (signed by wrong party)
      const invalidSignature = await prover.signMessage(ethers.getBytes(beliefHash)); // Should be origin!

      const zkProofBundle = ethers.AbiCoder.defaultAbiCoder().encode(
        ['bytes', 'bytes'],
        [validStarkProof, invalidSignature]
      );

      await expect(attestor.connect(prover).attestBelief(beliefHash, zkProofBundle))
        .to.be.revertedWith('Origin sig mismatch');
    });

    it('should support multiple attestations with different beliefs', async () => {
      const beliefs = [
        'Belief 1: Transparency with privacy',
        'Belief 2: No trusted setups',
        'Belief 3: Post-quantum security'
      ];

      for (const belief of beliefs) {
        const beliefHash = ethers.keccak256(ethers.toUtf8Bytes(belief));
        const originSignature = await origin.signMessage(ethers.getBytes(beliefHash));
        const starkProof = ethers.randomBytes(128);

        const zkProofBundle = ethers.AbiCoder.defaultAbiCoder().encode(
          ['bytes', 'bytes'],
          [starkProof, originSignature]
        );

        await attestor.connect(prover).attestBelief(beliefHash, zkProofBundle);

        expect(await attestor.isBeliefSovereign(beliefHash)).to.be.true;
      }
    });

    it('should allow same belief to be attested by different provers', async () => {
      const [_, prover1, prover2] = await ethers.getSigners();
      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('Shared belief'));
      const originSignature = await origin.signMessage(ethers.getBytes(beliefHash));
      const starkProof = ethers.randomBytes(128);

      const zkProofBundle = ethers.AbiCoder.defaultAbiCoder().encode(
        ['bytes', 'bytes'],
        [starkProof, originSignature]
      );

      // Prover 1 attests
      await attestor.connect(prover1).attestBelief(beliefHash, zkProofBundle);
      expect(await attestor.isBeliefSovereign(beliefHash)).to.be.true;

      // Prover 2 attests (same belief, different prover)
      await expect(attestor.connect(prover2).attestBelief(beliefHash, zkProofBundle))
        .to.not.be.reverted; // Should succeed

      expect(await attestor.isBeliefSovereign(beliefHash)).to.be.true;
    });
  });

  describe('DilithiumAttestor Constructor Validation', () => {
    it('should reject deployment with zkEnabled=true but no verifier', async () => {
      const DilithiumAttestor = await ethers.getContractFactory('DilithiumAttestor');

      await expect(DilithiumAttestor.deploy(
        origin.address,
        true,               // zkEnabled = true
        ethers.ZeroAddress  // No verifier (invalid!)
      )).to.be.revertedWith('ZK enabled but no verifier');
    });

    it('should reject deployment with zkEnabled=true and non-contract verifier', async () => {
      const DilithiumAttestor = await ethers.getContractFactory('DilithiumAttestor');
      const notAContract = prover.address; // EOA, not a contract

      await expect(DilithiumAttestor.deploy(
        origin.address,
        true,            // zkEnabled = true
        notAContract     // Not a contract address!
      )).to.be.revertedWith('Verifier address is not a contract');
    });

    it('should accept deployment with zkEnabled=false and no verifier', async () => {
      const DilithiumAttestor = await ethers.getContractFactory('DilithiumAttestor');

      const attestor = await DilithiumAttestor.deploy(
        origin.address,
        false,             // zkEnabled = false (signature-only mode)
        ethers.ZeroAddress // No verifier needed
      );

      expect(await attestor.zkEnabled()).to.be.false;
      expect(await attestor.verifierAddress()).to.equal(ethers.ZeroAddress);
    });

    it('should reject deployment with zero origin address', async () => {
      const DilithiumAttestor = await ethers.getContractFactory('DilithiumAttestor');

      await expect(DilithiumAttestor.deploy(
        ethers.ZeroAddress, // Invalid origin
        false,
        ethers.ZeroAddress
      )).to.be.revertedWith('Invalid origin address');
    });
  });

  describe('Gas Cost Analysis', () => {
    it('should measure gas cost for STARK proof verification', async () => {
      // Deploy attestor with ZK enabled
      const DilithiumAttestor = await ethers.getContractFactory('DilithiumAttestor');
      attestor = await DilithiumAttestor.deploy(
        origin.address,
        true,
        await verifier.getAddress()
      );

      const beliefHash = ethers.keccak256(ethers.toUtf8Bytes('Gas test belief'));
      const originSignature = await origin.signMessage(ethers.getBytes(beliefHash));
      const starkProof = ethers.randomBytes(128);

      const zkProofBundle = ethers.AbiCoder.defaultAbiCoder().encode(
        ['bytes', 'bytes'],
        [starkProof, originSignature]
      );

      const tx = await attestor.connect(prover).attestBelief(beliefHash, zkProofBundle);
      const receipt = await tx.wait();

      console.log(`      ⛽ Gas used for STARK verification: ${receipt.gasUsed.toString()}`);

      // STARK verification should be reasonable (under 500k gas)
      expect(receipt.gasUsed).to.be.lessThan(500000n);
    });
  });
});
