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

/**
 * @title MockPQCPrecompile, TEST-ONLY faithful mock of the AERE native PQC precompiles
 *
 * @notice Hardhat has no PQC precompile, so tests install this at 0x0AE1..0x0AE4 via
 *         hardhat_setCode. It parses the input EXACTLY as the live Besu precompiles do
 *         (per their Java sources: AereFalconSupport / MLDSA44 / SLHDSA128s) and returns
 *         a RAW 32-byte word 0x..01 iff the embedded signature "commitment" matches
 *         keccak256(pk || message). This models the real semantics: valid iff a genuine
 *         signature by THIS key over THIS message. Any malformed or non-matching input
 *         returns a raw 32-byte 0x..00 (invalid). The raw (non-ABI-encoded) return byte
 *         layout is identical to the live precompiles, so the caller's
 *         `ret.length >= 32 && ret[31] == 0x01` check behaves as it does on mainnet.
 *
 *         The parse offsets here are the precompiles' spec offsets, NOT copied from the
 *         contract under test, so a positive attest passing here proves the contract
 *         built the input at the correct offsets. The real encoding is separately proven
 *         against the LIVE mainnet precompile via eth_call.
 */
contract MockPQCPrecompile {
    uint8 private immutable _scheme;

    constructor(uint8 scheme_) {
        _scheme = scheme_;
    }

    // No public functions other than fallback, so an input whose leading bytes look
    // like a selector can never dispatch anywhere but the verifier.
    fallback() external {
        bool valid = _check(msg.data);
        bytes32 w = valid ? bytes32(uint256(1)) : bytes32(0);
        assembly {
            mstore(0x0, w)
            return(0x0, 32)
        }
    }

    function _check(bytes calldata input) internal view returns (bool) {
        if (_scheme == 1) return _checkFalcon(input, 897, 0x29);
        if (_scheme == 2) return _checkFalcon(input, 1793, 0x2A);
        if (_scheme == 3) return _checkFixed(input, 1312, 2420);
        if (_scheme == 4) return _checkFixed(input, 32, 7856);
        return false;
    }

    // Falcon: input = pk(pkLen) || sm ; sm = sigLen(2 BE) || nonce(40) || message || esig ;
    // esig = esigHeader || compressedSig ; sigLen == esig.length. (mirrors AereFalconSupport.verify)
    function _checkFalcon(bytes calldata input, uint256 pkLen, uint8 esigHeader) internal pure returns (bool) {
        if (input.length <= pkLen) return false;
        bytes calldata pk = input[0:pkLen];
        bytes calldata sm = input[pkLen:];
        if (sm.length < 2 + 40 + 2) return false;
        uint256 sigLen = (uint256(uint8(sm[0])) << 8) | uint256(uint8(sm[1]));
        if (sigLen < 2 || 2 + 40 + sigLen > sm.length) return false;
        uint256 msgLen = sm.length - 2 - 40 - sigLen;
        bytes calldata esig = sm[sm.length - sigLen:];
        if (esig.length < 33) return false;
        if (uint8(esig[0]) != esigHeader) return false;
        bytes calldata message = sm[42:42 + msgLen];
        bytes32 commitment = _load(esig, 1);
        return commitment == keccak256(abi.encodePacked(pk, message));
    }

    // ML-DSA-44 / SLH-DSA-128s: input = pk(pkLen) || sig(sigLen) || message(rest).
    function _checkFixed(bytes calldata input, uint256 pkLen, uint256 sigLen) internal pure returns (bool) {
        if (input.length < pkLen + sigLen + 32) return false;
        bytes calldata pk = input[0:pkLen];
        bytes calldata sig = input[pkLen:pkLen + sigLen];
        bytes calldata message = input[pkLen + sigLen:];
        bytes32 commitment = _load(sig, 0);
        return commitment == keccak256(abi.encodePacked(pk, message));
    }

    function _load(bytes calldata s, uint256 off) internal pure returns (bytes32 r) {
        assembly {
            r := calldataload(add(s.offset, off))
        }
    }
}