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

import {Test} from "forge-std/Test.sol";
import {AerePQCAttestation} from "../../contracts/pqc/AerePQCAttestation.sol";

/**
 * @title  Symbolic (halmos) proofs for the LIVE AerePQCAttestation contract.
 *
 * WHAT IS PROVEN (see PROVEN-PROPERTIES.md for the precise status of each):
 *   P1  a key's nonce only ever increases, and by EXACTLY 1, on a successful attest;
 *       a non-successful attest leaves the nonce unchanged.
 *   P2  an attestation slot (keyId, nonce) is written at most once (no overwrite):
 *       two consecutive successful attests land in DISTINCT slots and the first
 *       slot's recorded data is not mutated by the second.
 *   P3  a reverting / failed attest leaves ALL observable state unchanged
 *       (nonce, attestationCount, and the target attestation slot).
 *   P4  registerKey rejects malformed pubKeys: registration can succeed ONLY for a
 *       (scheme, length, header) triple that exactly matches the scheme's spec.
 *
 * CENTRAL ASSUMPTION -- PRECOMPILE AS AN UNINTERPRETED ORACLE.
 *   halmos cannot execute AERE's native PQC precompiles (0x0AE1..0x0AE4): they are
 *   Java/Bouncy-Castle native code on the Besu fork, not EVM bytecode. So we do NOT
 *   model Falcon/ML-DSA/SLH-DSA verification. Instead we ETCH a tiny EVM stub at the
 *   scheme's precompile address whose 32-byte return word is read from its storage
 *   slot 0, and we drive that word with a SYMBOLIC value (`oracleWord`). halmos then
 *   explores BOTH the "precompile says valid" and "precompile says invalid" branches.
 *   This is strictly MORE adversarial than a faithful precompile: it proves the
 *   safety properties hold no matter what the precompile returns, on every call.
 *   It therefore proves NOTHING about the cryptographic soundness of the precompile
 *   itself (that is out of scope for symbolic execution and is covered separately by
 *   the NIST-KAT tests and live eth_call checks).
 *
 * OTHER ASSUMPTIONS:
 *   - keccak256 is treated as injective (halmos's default: no hash collisions). P2's
 *     "distinct (keyId,nonce) => distinct storage slot" rests on this.
 *   - We fix the key under test to a single registered Falcon-512 key (scheme 1,
 *     897-byte pubkey). The nonce/slot logic is scheme-independent, so this is a
 *     representative instantiation, not a whole-contract proof over all schemes.
 *   - Symbolic `bytes` lengths come from --default-bytes-lengths (passed on the CLI).
 */
contract AerePQCAttestationSymbolic is Test {
    AerePQCAttestation internal att;

    // Falcon-512 precompile address the contract staticcalls for scheme 1.
    address internal constant ORACLE = address(0x0AE1);

    // Runtime that returns storage slot 0 as the 32-byte result word:
    // PUSH1 0 SLOAD PUSH1 0 MSTORE PUSH1 32 PUSH1 0 RETURN
    bytes internal constant ORACLE_CODE = hex"60005460005260206000f3";

    function setUp() public {
        att = new AerePQCAttestation();

        // Register one well-formed Falcon-512 key => keyId 0.
        bytes memory pk = new bytes(897);
        pk[0] = bytes1(uint8(0x09)); // FALCON512_PK_HEADER
        att.registerKey(1, pk);

        // Stand up the uninterpreted precompile oracle at 0x0AE1.
        vm.etch(ORACLE, ORACLE_CODE);
    }

    // Falcon envelope must satisfy 40 < len and esig length <= uint16 max. We bound
    // the symbolic signature length to a valid band so the success branch is reachable;
    // shorter lengths still exercise the failure/revert branch.
    function _assumeFalconSigLen(uint256 len) internal pure {
        // 41..1024 keeps esig in (0, 984], well under type(uint16).max.
        // (Concrete symbolic lengths are chosen by --default-bytes-lengths.)
        vm.assume(len == 0 || len == 41 || len == 65 || len == 200 || len == 1024);
    }

    // ---------------------------------------------------------------------
    // P1 + P3 : nonce increments by exactly 1 on success, unchanged otherwise;
    //           attestationCount and the target slot move in lock-step, and a
    //           failed/reverting attest is a complete no-op.
    // ---------------------------------------------------------------------
    function check_nonceAndStateMoveOnlyOnSuccess(
        bytes32 messageHash,
        bytes calldata signature,
        bytes32 oracleWord
    ) external {
        // Drive the uninterpreted precompile result (valid iff low byte == 0x01).
        vm.store(ORACLE, bytes32(0), oracleWord);

        uint64 preNonce = att.nonceOf(0);
        uint256 preCount = att.attestationCount();
        (bool preExists,,,) = att.getAttestation(0, preNonce);

        (bool ok,) = address(att).call(
            abi.encodeCall(att.attest, (0, messageHash, signature))
        );

        uint64 postNonce = att.nonceOf(0);
        uint256 postCount = att.attestationCount();
        (bool postExists, bytes32 postMsg,,) = att.getAttestation(0, preNonce);

        if (ok) {
            // Exactly one step forward, and the slot for the consumed nonce now exists
            // committing to messageHash. (Before this call it did not exist.)
            assert(postNonce == preNonce + 1);
            assert(postCount == preCount + 1);
            assert(!preExists);
            assert(postExists);
            assert(postMsg == messageHash);
        } else {
            // Total no-op: nonce, count, and the slot are all untouched.
            assert(postNonce == preNonce);
            assert(postCount == preCount);
            assert(postExists == preExists);
        }
    }

    // ---------------------------------------------------------------------
    // P2 : no overwrite. Two consecutive SUCCESSFUL attests use distinct slots,
    //      and the second does not mutate the first slot's recorded data.
    // ---------------------------------------------------------------------
    function check_noOverwriteAcrossTwoAttests(
        bytes32 m0,
        bytes32 m1,
        bytes calldata s0,
        bytes calldata s1
    ) external {
        _assumeFalconSigLen(s0.length);
        _assumeFalconSigLen(s1.length);
        vm.assume(s0.length > 40 && s1.length > 40); // force the reachable success band

        // Force the precompile to accept (low byte 0x01): both attests succeed.
        vm.store(ORACLE, bytes32(0), bytes32(uint256(1)));

        att.attest(0, m0, s0); // consumes nonce 0
        (bool e0, bytes32 mh0,,) = att.getAttestation(0, 0);
        assert(e0 && mh0 == m0);

        att.attest(0, m1, s1); // consumes nonce 1

        // Distinct storage ids for the two (keyId, nonce) pairs (keccak injectivity).
        assert(att.attestationId(0, 0) != att.attestationId(0, 1));

        // Slot 1 recorded m1; slot 0 is UNCHANGED (still m0, still the original record).
        (bool e1, bytes32 mh1,,) = att.getAttestation(0, 1);
        assert(e1 && mh1 == m1);
        (bool e0b, bytes32 mh0b,,) = att.getAttestation(0, 0);
        assert(e0b && mh0b == m0);
    }

    // ---------------------------------------------------------------------
    // P4 : registerKey rejects malformed pubKeys. Registration may succeed ONLY
    //      for a (scheme, length, header) triple matching the scheme's spec; an
    //      unknown scheme can never register.
    // ---------------------------------------------------------------------
    function check_registerKeyRejectsMalformedPubKey(uint8 scheme, bytes calldata pubKey) external {
        (bool ok,) = address(att).call(abi.encodeCall(att.registerKey, (scheme, pubKey)));
        if (!ok) return; // rejected: nothing to prove on this path

        if (scheme == 1) {
            assert(pubKey.length == 897 && uint8(pubKey[0]) == 0x09);
        } else if (scheme == 2) {
            assert(pubKey.length == 1793 && uint8(pubKey[0]) == 0x0A);
        } else if (scheme == 3) {
            assert(pubKey.length == 1312);
        } else if (scheme == 4) {
            assert(pubKey.length == 32);
        } else {
            assert(false); // unknown scheme must be impossible to register
        }
    }
}
