// SPDX-License-Identifier: MIT
pragma solidity 0.8.23;

import "@openzeppelin/contracts/access/Ownable.sol";

/// @dev Canonical SP1 verifier / SP1VerifierGateway ABI: verifyProof RETURNS
///      NOTHING and REVERTS on an invalid proof.
interface ISp1Verifier {
    function verifyProof(
        bytes32 programVKey,
        bytes calldata publicValues,
        bytes calldata proof
    ) external view;
}

/**
 * @title AereProofAggregatorV2 — on-chain anchor for RECURSIVE SP1 proof aggregation.
 *         CORRECTED fork of AereProofAggregator 0x6a260238…F84C5 (see THE FIX below).
 *
 * @notice Verifies a single SP1 Groth16 "proof of proofs": the aggregation guest
 *         recursively verified N inner SP1 proofs inside the zkVM via
 *         `verify_sp1_proof`, and committed one composite digest binding all N
 *         (vkey, publicValues) pairs. This contract verifies that Groth16 proof
 *         through the SP1 gateway, recomputes the composite from the supplied
 *         inner data to bind it to the proof, REQUIRES every folded inner program
 *         to be an ALLOWED program, and records the aggregation.
 *
 *         COMPOSITE (recomputed here byte-for-byte with the guest):
 *           composite = sha256(
 *               uint32_be(n)
 *               || (innerVKeyDigest_0 || innerPvHash_0)
 *               || ...
 *               || (innerVKeyDigest_{n-1} || innerPvHash_{n-1})
 *           )
 *         where
 *           innerVKeyDigest_i : the KoalaBear (SP1Field) verifying-key digest the
 *                               recursion verifier uses (SP1VerifyingKey::hash_u32,
 *                               big-endian). This is what `verify_sp1_proof`
 *                               actually checked, so binding it here is maximally
 *                               sound.
 *           innerPvHash_i     : sha256 of inner proof i's public values.
 *
 * @dev    ------------------------------------------------------------------
 *         THE FIX (vs V1 0x6a260238…F84C5) — INNER VKEYS COUNTED, NEVER ENFORCED.
 *
 *         THE GUEST IS UNCHANGED AND WAS NEVER THE PROBLEM. It is genuinely
 *         recursive, and the composite it commits ALREADY binds every inner vkey
 *         digest. V1 even recomputed that composite correctly, so V1 always KNEW
 *         exactly which programs had been folded.
 *
 *         V1's flaw was what it DID with that knowledge: nothing. It merely COUNTED
 *         how many inner vkeys happened to be registered into a `recognized` field
 *         and recorded the aggregation regardless:
 *             for (i) { if (innerPrograms[digests[i]].known) recognized++; }
 *         An unregistered inner program was not an error — just a lower number.
 *
 *         THE ATTACK. An attacker writes an SP1 guest with NO CONSTRAINTS AT ALL:
 *             fn main() { sp1_zkvm::io::commit_slice(&attacker_chosen_bytes); }
 *         and has it commit, say, the exact AereZKScreen public-values tuple. They
 *         prove it (trivially — the program asserts nothing), fold it with the real
 *         aggregation guest (which honestly verifies it, because it IS a valid SP1
 *         proof of that vacuous program), and submit. V1 records an aggregation that
 *         `exists`, carries the pinned aggVKey, and looks legitimate — differing from
 *         a genuine one only by `recognized < count`, a field nothing enforced.
 *
 *         V2 REQUIRES every inner vkey digest to be an ALLOWED program and REVERTS
 *         with `UnregisteredInnerProgram(digest, index)` otherwise. The attacker's
 *         vacuous guest has a different vkey (a vkey IS a commitment to the program),
 *         so it can never be in the allowed set. Both `recordAggregation` and the
 *         `verifyAggregation` view enforce it — mirroring the AereZkEthLightClientV2
 *         lesson that a view used as an oracle must bind exactly like the stateful path.
 *
 *         THE ALLOWED SET IS SEEDED AT DEPLOY (constructor), so this contract is
 *         FUNCTIONAL ON ARRIVAL with the same three programs V1 carried — it needs no
 *         post-deploy signature to be safe. The Foundation may ADD programs, and
 *         {sealRegistry} is a ONE-WAY latch that makes the set permanently IMMUTABLE.
 *         `recognized` is retained for read-shape compatibility but is now an
 *         INVARIANT: it always equals `count`, because an unregistered inner reverts.
 *         ------------------------------------------------------------------
 *
 * @dev    HONEST SCOPE / RESIDUAL TRUST — read before quoting.
 *         * What this proves: "N proofs, each from an ALLOWED program, exist with
 *           these public values, and one Groth16 proof attests all N were verified
 *           recursively in-zkVM". That is a real, strong statement — and it is the
 *           statement V1 failed to make.
 *         * What this does NOT prove: the SEMANTICS of the inner public values. The
 *           contract does not know what a zkscreen tuple means; it only knows the
 *           program that produced it is allowed. A consumer must still interpret the
 *           inner public values against the program's own spec.
 *         * REGISTRATION IS A TRUST ACTION. Allowing a program asserts that program's
 *           constraints are sound. A bad registration is as damaging as the V1 bug.
 *           This is why {sealRegistry} exists: once the set is final, seal it and the
 *           admin surface is gone forever. Until sealed, the owner is trusted.
 *         * VKEY ENCODINGS. `innerVKeyDigest` is the KoalaBear digest the recursion
 *           verifier checks; `onChainVKey` is the BN254 bytes32 published for the same
 *           program. They are two encodings of one key; the registry records the
 *           mapping so each folded program can be labelled and cross-checked. The
 *           ENFORCED value is the KoalaBear digest — the one the proof actually binds.
 *         * NOT QUANTUM-SAFE, TRUST-MINIMIZED NOT TRUSTLESS: SP1 Groth16 over BN254 is
 *           classical. This is an APPLICATION-layer contract; it does not touch AERE
 *           consensus (still classical ECDSA QBFT).
 *
 *         IMMUTABILITY: SP1_VERIFIER is fixed at deploy. Recorded aggregations cannot
 *         be altered or removed.
 */
contract AereProofAggregatorV2 is Ownable {
    address public immutable SP1_VERIFIER;

    /// @notice Registered aggregation-program vkeys (BN254 bytes32) allowed to be
    ///         routed to the gateway, with a human description.
    mapping(bytes32 => bool) public isAggProgram;
    mapping(bytes32 => string) public aggProgramDesc;

    struct InnerProgram {
        bool known;
        bytes32 onChainVKey; // BN254 bytes32 published for this program
        string label;        // e.g. "zkscreen", "over18", "zkml-mnist"
    }

    /// @notice KoalaBear inner vkey digest => registered program metadata.
    ///         An inner proof whose digest is NOT `known` here is REJECTED (the V2 fix).
    mapping(bytes32 => InnerProgram) public innerPrograms;

    /// @notice Once true, the allowed set is IMMUTABLE forever: no program can be
    ///         added or changed, by anyone, including the owner. One-way latch.
    bool public registrySealed;

    struct Aggregation {
        bool exists;
        bytes32 aggVKey;
        uint32 count;
        uint32 recognized; // INVARIANT in V2: always == count (unregistered inners revert)
        uint64 at;
        address submitter;
    }

    /// @notice composite digest => recorded aggregation.
    mapping(bytes32 => Aggregation) public aggregations;
    uint256 public aggregationCount;

    /* --------------------------------- events ------------------------------- */

    event AggProgramRegistered(bytes32 indexed aggVKey, string description);
    event InnerProgramRegistered(bytes32 indexed innerVKeyDigest, bytes32 onChainVKey, string label);
    event RegistrySealed(address indexed by, uint256 at);
    event AggregationVerified(
        bytes32 indexed composite,
        bytes32 indexed aggVKey,
        uint32 count,
        uint32 recognized,
        bytes32[] innerVKeyDigests,
        address indexed submitter,
        uint256 at
    );

    /* --------------------------------- errors ------------------------------- */

    error ZeroAddress();
    error UnknownAggProgram(bytes32 aggVKey);
    error LengthMismatch(uint256 vkeys, uint256 pvHashes);
    error EmptyAggregation();
    error BadPublicValuesLength(uint256 got);
    error CompositeMismatch(bytes32 recomputed, bytes32 committed);
    error InvalidProof();
    error VerifierHasNoCode();
    error ZeroVKey();
    error RegistryIsSealed();
    error SeedLengthMismatch();
    /// @notice THE V2 FIX: a folded inner program that is not in the allowed set.
    error UnregisteredInnerProgram(bytes32 innerVKeyDigest, uint256 index);

    /* ----------------------------- constructor ------------------------------ */

    /**
     * @param sp1Verifier       the live SP1VerifierGateway.
     * @param initialOwner      the account that may extend the allowed set (Foundation
     *                          or a Timelock). Set to address(0) ONLY if `seal` is true.
     * @param aggVKeys          aggregation-program vkeys to allow at deploy.
     * @param aggDescs          matching descriptions.
     * @param innerDigests      KoalaBear inner vkey digests to allow at deploy.
     * @param innerOnChainVKeys matching BN254 on-chain vkeys.
     * @param innerLabels       matching labels.
     * @param seal              if true, seal the registry immediately: the seeded set
     *                          becomes IMMUTABLE and ownership is renounced in the same
     *                          transaction (no admin ever exists).
     */
    constructor(
        address sp1Verifier,
        address initialOwner,
        bytes32[] memory aggVKeys,
        string[] memory aggDescs,
        bytes32[] memory innerDigests,
        bytes32[] memory innerOnChainVKeys,
        string[] memory innerLabels,
        bool seal
    ) {
        if (sp1Verifier == address(0)) revert ZeroAddress();
        // A code-less verifier would make verifyProof (which returns no data) succeed
        // silently, accepting every "proof". Removes a deploy-time footgun.
        if (sp1Verifier.code.length == 0) revert VerifierHasNoCode();
        if (aggVKeys.length != aggDescs.length) revert SeedLengthMismatch();
        if (innerDigests.length != innerOnChainVKeys.length || innerDigests.length != innerLabels.length) {
            revert SeedLengthMismatch();
        }
        SP1_VERIFIER = sp1Verifier;

        for (uint256 i; i < aggVKeys.length; ++i) {
            _registerAggProgram(aggVKeys[i], aggDescs[i]);
        }
        for (uint256 i; i < innerDigests.length; ++i) {
            _registerInnerProgram(innerDigests[i], innerOnChainVKeys[i], innerLabels[i]);
        }

        if (seal) {
            registrySealed = true;
            emit RegistrySealed(msg.sender, block.timestamp);
            _transferOwnership(address(0)); // no admin, ever
        } else {
            if (initialOwner == address(0)) revert ZeroAddress();
            _transferOwnership(initialOwner);
        }
    }

    /* ------------------------------ governance ------------------------------ */

    function registerAggProgram(bytes32 aggVKey, string calldata description) external onlyOwner {
        if (registrySealed) revert RegistryIsSealed();
        _registerAggProgram(aggVKey, description);
    }

    function registerInnerProgram(
        bytes32 innerVKeyDigest,
        bytes32 onChainVKey,
        string calldata label
    ) external onlyOwner {
        if (registrySealed) revert RegistryIsSealed();
        _registerInnerProgram(innerVKeyDigest, onChainVKey, label);
    }

    /// @notice Permanently freeze the allowed set. ONE-WAY: after this, no program can
    ///         ever be added or changed, and ownership is renounced. The contract
    ///         becomes fully immutable — the strongest form of the V2 fix.
    function sealRegistry() external onlyOwner {
        if (registrySealed) revert RegistryIsSealed();
        registrySealed = true;
        emit RegistrySealed(msg.sender, block.timestamp);
        _transferOwnership(address(0));
    }

    function _registerAggProgram(bytes32 aggVKey, string memory description) internal {
        if (aggVKey == bytes32(0)) revert ZeroVKey();
        isAggProgram[aggVKey] = true;
        aggProgramDesc[aggVKey] = description;
        emit AggProgramRegistered(aggVKey, description);
    }

    function _registerInnerProgram(bytes32 innerVKeyDigest, bytes32 onChainVKey, string memory label) internal {
        // A zero digest must never be allowed: it is the default value of an
        // unregistered slot and would turn the enforcement gate into a no-op.
        if (innerVKeyDigest == bytes32(0)) revert ZeroVKey();
        innerPrograms[innerVKeyDigest] = InnerProgram({known: true, onChainVKey: onChainVKey, label: label});
        emit InnerProgramRegistered(innerVKeyDigest, onChainVKey, label);
    }

    /* --------------------------- record aggregation ------------------------- */

    /// @notice Verify a recursive aggregation Groth16 proof and record it.
    ///         REVERTS unless EVERY folded inner program is in the allowed set (V2 fix).
    /// @param aggVKey            registered aggregation-program vkey (BN254 bytes32)
    /// @param innerVKeyDigests   KoalaBear vkey digests of the folded inner proofs
    /// @param innerPvHashes      sha256 of each inner proof's public values
    /// @param publicValues       the aggregation proof's committed 32-byte composite
    /// @param proof              the SP1 Groth16 proof bytes (selector 0x4388a21c)
    /// @return composite         the verified composite digest
    function recordAggregation(
        bytes32 aggVKey,
        bytes32[] calldata innerVKeyDigests,
        bytes32[] calldata innerPvHashes,
        bytes calldata publicValues,
        bytes calldata proof
    ) external returns (bytes32 composite) {
        composite = _check(aggVKey, innerVKeyDigests, innerPvHashes, publicValues, proof);

        uint32 n = uint32(innerVKeyDigests.length);
        aggregations[composite] = Aggregation({
            exists: true,
            aggVKey: aggVKey,
            count: n,
            recognized: n, // invariant: enforcement above guarantees all are registered
            at: uint64(block.timestamp),
            submitter: msg.sender
        });
        aggregationCount += 1;

        emit AggregationVerified(composite, aggVKey, n, n, innerVKeyDigests, msg.sender, block.timestamp);
    }

    /* ----------------------------------- views ------------------------------ */

    /// @notice Stateless verification: recomputes the composite, checks it equals the
    ///         committed public values, ENFORCES that every inner program is allowed
    ///         (the V2 fix — a view used as an oracle must bind like the stateful path),
    ///         and verifies the Groth16 proof. Reverts if anything is invalid; records nothing.
    function verifyAggregation(
        bytes32 aggVKey,
        bytes32[] calldata innerVKeyDigests,
        bytes32[] calldata innerPvHashes,
        bytes calldata publicValues,
        bytes calldata proof
    ) external view returns (bytes32 composite) {
        composite = _check(aggVKey, innerVKeyDigests, innerPvHashes, publicValues, proof);
    }

    function getAggregation(bytes32 composite)
        external
        view
        returns (bool exists, bytes32 aggVKey, uint32 count, uint32 recognized, uint64 at, address submitter)
    {
        Aggregation memory a = aggregations[composite];
        return (a.exists, a.aggVKey, a.count, a.recognized, a.at, a.submitter);
    }

    /// @notice True iff every supplied inner vkey digest is an allowed program.
    ///         Convenience for callers that want to check before spending gas.
    function areInnerProgramsAllowed(bytes32[] calldata innerVKeyDigests) external view returns (bool) {
        for (uint256 i; i < innerVKeyDigests.length; ++i) {
            if (!innerPrograms[innerVKeyDigests[i]].known) return false;
        }
        return true;
    }

    /* ---------------------------------- internal ---------------------------- */

    /// @dev The single gate both entry points go through: agg-program allowed,
    ///      lengths sane, EVERY inner program allowed (the V2 fix), composite bound
    ///      to the committed public values, and the Groth16 proof valid.
    function _check(
        bytes32 aggVKey,
        bytes32[] calldata innerVKeyDigests,
        bytes32[] calldata innerPvHashes,
        bytes calldata publicValues,
        bytes calldata proof
    ) internal view returns (bytes32 composite) {
        if (!isAggProgram[aggVKey]) revert UnknownAggProgram(aggVKey);
        uint256 n = innerVKeyDigests.length;
        if (n != innerPvHashes.length) revert LengthMismatch(n, innerPvHashes.length);
        if (n == 0) revert EmptyAggregation();
        if (publicValues.length != 32) revert BadPublicValuesLength(publicValues.length);

        // THE V2 FIX: every folded inner program must be ALLOWED. V1 only counted these.
        // Done BEFORE the proof verify so an unregistered program is rejected cheaply and
        // with a precise, debuggable error rather than an opaque gateway revert.
        for (uint256 i; i < n; ++i) {
            if (!innerPrograms[innerVKeyDigests[i]].known) {
                revert UnregisteredInnerProgram(innerVKeyDigests[i], i);
            }
        }

        // recompute the composite exactly as the aggregation guest committed it.
        composite = _composite(innerVKeyDigests, innerPvHashes);

        bytes32 committed;
        assembly {
            committed := calldataload(publicValues.offset)
        }
        if (composite != committed) revert CompositeMismatch(composite, committed);

        // verify the Groth16 proof of proofs through the SP1 gateway.
        try ISp1Verifier(SP1_VERIFIER).verifyProof(aggVKey, publicValues, proof) {
            // verified: did not revert
        } catch {
            revert InvalidProof();
        }
    }

    function _composite(bytes32[] calldata innerVKeyDigests, bytes32[] calldata innerPvHashes)
        internal
        pure
        returns (bytes32)
    {
        uint256 n = innerVKeyDigests.length;
        bytes memory packed = abi.encodePacked(uint32(n));
        for (uint256 i; i < n; ++i) {
            packed = bytes.concat(packed, innerVKeyDigests[i], innerPvHashes[i]);
        }
        return sha256(packed);
    }
}