// 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 AereZKMLVerifier — on-chain anchor for zero-knowledge ML inference
 * @notice Verifies SP1 zkVM proofs that a registered machine-learning model
 *         classified a PRIVATE input, WITHOUT revealing the input. The reference
 *         model is a trained, quantized MNIST handwritten-digit classifier
 *         (784 -> 256 -> 10 MLP) whose forward pass runs inside the zkVM on a
 *         private 28x28 image; only the predicted digit is made public.
 *
 *         Each model is registered by the Foundation with:
 *           - programVKey  : SP1 verification key binding the exact guest ELF
 *                            (which embeds the quantized weights).
 *           - modelHash    : keccak256 of the quantized weight blob, recomputed
 *                            and committed INSIDE the zkVM, so a verifier knows
 *                            the classification came from THIS exact model.
 *           - numClasses   : output cardinality (10 for MNIST digits).
 *           - accuracyBps  : the model's measured test-set accuracy in basis
 *                            points (e.g. 9798 = 97.98% on the 10,000-image
 *                            MNIST test set) recorded for transparency.
 *
 *         PUBLIC-VALUES CONVENTION (exactly 128 bytes):
 *           abi.encode(uint256 chainId, address user, bytes32 modelHash, uint256 predictedClass)
 *
 *         submitInference verifies the proof through the SP1 gateway, then binds
 *         chainId == this chain, user == msg.sender, and modelHash == the
 *         registered model, and records the user's latest verified prediction.
 *
 *         IMMUTABILITY: SP1_VERIFIER is fixed at deploy; a model's programVKey,
 *         modelHash and numClasses are immutable after registration.
 */
contract AereZKMLVerifier is Ownable {

    address public immutable SP1_VERIFIER;

    struct Model {
        bool registered;
        uint16 numClasses;
        uint32 accuracyBps;     // measured test-set accuracy, basis points
        bytes32 programVKey;
        bytes32 modelHash;      // keccak256 of the quantized weight blob
        string description;
    }

    struct Inference {
        bool exists;
        uint16 predictedClass;
        uint64 at;              // block timestamp of the verified inference
    }

    /// @notice programVKey => Model
    mapping(bytes32 => Model) public models;

    /// @notice user => programVKey => latest verified inference
    mapping(address => mapping(bytes32 => Inference)) public inferences;

    /// @notice programVKey => count of verified inferences
    mapping(bytes32 => uint256) public inferenceCount;

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

    event ModelRegistered(
        bytes32 indexed programVKey,
        bytes32 indexed modelHash,
        uint16 numClasses,
        uint32 accuracyBps,
        string description
    );
    event InferenceVerified(
        address indexed user,
        bytes32 indexed programVKey,
        bytes32 indexed modelHash,
        uint256 predictedClass,
        uint256 at
    );

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

    error UnknownModel();
    error AlreadyRegistered();
    error InvalidProof();
    error WrongChain(uint256 inProof, uint256 onChain);
    error UserMismatch(address inProof, address msgSender);
    error ModelHashMismatch(bytes32 inProof, bytes32 registered);
    error ClassOutOfRange(uint256 predicted, uint16 numClasses);
    error BadPublicValuesLength(uint256 got);
    error ZeroAddress();
    error ZeroModelHash();

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

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

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

    function registerModel(
        bytes32 programVKey,
        bytes32 modelHash,
        uint16 numClasses,
        uint32 accuracyBps,
        string calldata description
    ) external onlyOwner {
        if (models[programVKey].registered) revert AlreadyRegistered();
        if (modelHash == bytes32(0)) revert ZeroModelHash();
        models[programVKey] = Model({
            registered: true,
            numClasses: numClasses,
            accuracyBps: accuracyBps,
            programVKey: programVKey,
            modelHash: modelHash,
            description: description
        });
        emit ModelRegistered(programVKey, modelHash, numClasses, accuracyBps, description);
    }

    /* ------------------------------ proof submit ---------------------------- */

    /// @notice Verify an SP1 proof that the model `programVKey` classified the
    ///         caller's private input, and record the predicted class. The proof
    ///         must commit (chainId, user, modelHash, predictedClass) with
    ///         chainId == this chain, user == msg.sender, and modelHash equal to
    ///         the registered model's hash.
    function submitInference(
        bytes32 programVKey,
        bytes calldata publicValues,
        bytes calldata proof
    ) external returns (uint256 predictedClass) {
        Model memory m = models[programVKey];
        if (!m.registered) revert UnknownModel();
        // 4 x 32 = 128 bytes for (uint256, address, bytes32, uint256).
        if (publicValues.length != 128) revert BadPublicValuesLength(publicValues.length);

        try ISp1Verifier(SP1_VERIFIER).verifyProof(programVKey, publicValues, proof) {
            // verified — did not revert
        } catch {
            revert InvalidProof();
        }

        (uint256 chainId, address user, bytes32 modelHash, uint256 predicted) =
            abi.decode(publicValues, (uint256, address, bytes32, uint256));

        if (chainId != block.chainid) revert WrongChain(chainId, block.chainid);
        if (user != msg.sender) revert UserMismatch(user, msg.sender);
        if (modelHash != m.modelHash) revert ModelHashMismatch(modelHash, m.modelHash);
        if (predicted >= m.numClasses) revert ClassOutOfRange(predicted, m.numClasses);

        inferences[user][programVKey] = Inference({
            exists: true,
            predictedClass: uint16(predicted),
            at: uint64(block.timestamp)
        });
        inferenceCount[programVKey] += 1;

        emit InferenceVerified(user, programVKey, modelHash, predicted, block.timestamp);
        return predicted;
    }

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

    /// @notice Stateless verification: reverts if the proof is invalid for the
    ///         given programVKey, otherwise returns the decoded predicted class.
    ///         Does not check chainId/user/modelHash and records nothing.
    function verify(
        bytes32 programVKey,
        bytes calldata publicValues,
        bytes calldata proof
    ) external view returns (uint256 predictedClass) {
        if (publicValues.length != 128) revert BadPublicValuesLength(publicValues.length);
        ISp1Verifier(SP1_VERIFIER).verifyProof(programVKey, publicValues, proof);
        (, , , predictedClass) = abi.decode(publicValues, (uint256, address, bytes32, uint256));
    }

    function latestPrediction(address user, bytes32 programVKey)
        external
        view
        returns (bool exists, uint16 predictedClass, uint64 at)
    {
        Inference memory inf = inferences[user][programVKey];
        return (inf.exists, inf.predictedClass, inf.at);
    }
}