// 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 AereProofAggregator: on-chain anchor for RECURSIVE SP1 proof aggregation
 * @notice Verifies a single SP1 Groth16 "proof of proofs": the aggregation guest
 *         recursively verified N inner SP1 proofs (each from a real, distinct AERE
 *         production program on chain 2800) 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, 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.
 *
 *         Because the on-chain (BN254) vkey and the KoalaBear digest are two
 *         different encodings of the same key, the Foundation registers the
 *         mapping (registerInnerProgram) so each folded program can be labelled
 *         and cross-checked against its published on-chain vkey.
 *
 *         The aggregation program itself is registered by the Foundation
 *         (registerAggProgram) and is the only vkey routed to the gateway.
 *
 *         IMMUTABILITY: SP1_VERIFIER is fixed at deploy. Recorded aggregations
 *         cannot be altered or removed.
 */
contract AereProofAggregator 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.
    mapping(bytes32 => InnerProgram) public innerPrograms;

    struct Aggregation {
        bool exists;
        bytes32 aggVKey;
        uint32 count;
        uint32 recognized;
        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 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();

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

    constructor(address sp1Verifier) {
        if (sp1Verifier == address(0)) revert ZeroAddress();
        SP1_VERIFIER = sp1Verifier;
    }

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

    function registerAggProgram(bytes32 aggVKey, string calldata description) external onlyOwner {
        isAggProgram[aggVKey] = true;
        aggProgramDesc[aggVKey] = description;
        emit AggProgramRegistered(aggVKey, description);
    }

    function registerInnerProgram(
        bytes32 innerVKeyDigest,
        bytes32 onChainVKey,
        string calldata label
    ) external onlyOwner {
        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.
    /// @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) {
        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);

        // 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();
        }

        // count how many folded inner programs are registered/recognized.
        uint32 recognized;
        for (uint256 i; i < n; ++i) {
            if (innerPrograms[innerVKeyDigests[i]].known) recognized++;
        }

        aggregations[composite] = Aggregation({
            exists: true,
            aggVKey: aggVKey,
            count: uint32(n),
            recognized: recognized,
            at: uint64(block.timestamp),
            submitter: msg.sender
        });
        aggregationCount += 1;

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

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

    /// @notice Stateless verification: recomputes the composite, checks it equals
    ///         the committed public values, 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) {
        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);

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

        ISp1Verifier(SP1_VERIFIER).verifyProof(aggVKey, 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);
    }

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

    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);
    }
}