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

interface IAerePQCKeyRegistry {
    function getKey(uint256 keyId)
        external
        view
        returns (
            address owner,
            uint8 scheme,
            uint8 status,
            uint64 registeredBlock,
            uint256 rotatedFromId,
            uint256 rotatedToId,
            bytes memory pubKey
        );

    function verifyWithKey(uint256 keyId, bytes32 message, bytes calldata signature) external view returns (bool);

    function keyCount() external view returns (uint256);
}

/**
 * @title AereAgentDID, post-quantum-rooted decentralized identity for AI agents
 *        (AERE chain 2800)
 *
 * @notice A DID for autonomous/AI agents whose ROOT authority is a quantum-durable
 *         Falcon key held in AerePQCKeyRegistry, and whose day-to-day work is signed by
 *         cheap, short-lived, revocable secp256k1 SESSION keys.
 *
 *         WHY THE SPLIT. The Falcon root is a cold, high-value, quantum-safe authority:
 *         possession was proven on-chain when the key was registered, and it is the only
 *         thing that can authorize a new session (issuance is Falcon-proof-of-possession
 *         gated, verified in full on-chain by AERE's live native precompile). Sessions
 *         are hot keys for high-frequency signing: each agent action is a plain secp256k1
 *         signature recovered with ecrecover (~3k gas), bounded by a per-session SCOPE
 *         and cumulative SPEND CAP, expiring at a deadline, and revocable at any time.
 *         A leaked session key is contained by its scope, cap, expiry, and the on-chain
 *         revocation list; the quantum-durable root is never exposed to hot use.
 *
 *         ROOT LIFECYCLE IS ENFORCED. A session is only valid while its root key is still
 *         ACTIVE in the registry: if the Falcon root is ROTATED or REVOKED, every session
 *         under it stops validating immediately (checked on each read via the registry).
 *
 *         HONEST SCOPE. Application-layer identity + authorization. This does NOT change
 *         AERE consensus (blocks are still Besu QBFT ECDSA). Falcon verification is done
 *         by the live native precompile through AerePQCKeyRegistry.verifyWithKey. This
 *         contract holds no funds: the SPEND CAP is an accounting policy the contract
 *         enforces on recorded actions, not custody of value.
 */
contract AereAgentDID {
    // ----- Falcon scheme ids in AerePQCKeyRegistry (root must be one of these) -----
    uint8 internal constant SCHEME_FALCON512 = 1;
    uint8 internal constant SCHEME_FALCON1024 = 2;
    uint8 internal constant REGISTRY_STATUS_ACTIVE = 1;

    // ----- domain separators -----
    bytes32 public constant SESSION_DOMAIN = keccak256("AereAgentDID.v1.issueSession");
    bytes32 public constant ACTION_DOMAIN = keccak256("AereAgentDID.v1.action");

    // ----- immutable wiring -----
    IAerePQCKeyRegistry public immutable keyRegistry;
    // Read-only cross-links to the agent economic/reputation layer (may be zero if the
    // deployer does not wire them). Informational: consumers read these to locate an
    // agent's bond / reputation contracts. This contract never calls into them.
    address public immutable agentBond; // AereAgentBond (slashable agent bond) or 0
    address public immutable aiReputation; // AereAIReputation (composable reputation) or 0

    // ----- agents: one DID per Falcon root key (agentId == rootKeyId) -----
    struct Agent {
        uint256 rootKeyId; // the Falcon root key in AerePQCKeyRegistry
        address controller; // the registry owner of rootKeyId at creation (fast-revoke authority)
        uint64 createdBlock;
        uint64 sessionNonce; // per-agent nonce making each issuance challenge unique
        bool exists;
    }

    mapping(uint256 => Agent) private _agents; // rootKeyId => Agent

    // ----- sessions (global incrementing id -> on-chain revocation list) -----
    struct Session {
        uint256 rootKeyId; // the agent this session belongs to
        address sessionAddr; // secp256k1 address that signs actions (ecrecover target)
        bytes32 scopeHash; // opaque scope tag the session is locked to
        uint256 spendCap; // cumulative spend allowed under this session
        uint256 spent; // cumulative spend recorded so far
        uint64 issuedBlock;
        uint64 expiry; // unix timestamp after which the session is invalid
        uint64 actionNonce; // per-session action counter (anti-replay for authorize)
        bool revoked;
        bool exists;
    }

    Session[] private _sessions; // sessionId is the index (0-based)
    mapping(uint256 => uint256[]) private _agentSessions; // rootKeyId => sessionIds

    // ----- events -----
    event AgentCreated(uint256 indexed rootKeyId, address indexed controller, uint8 rootScheme);
    event SessionIssued(
        uint256 indexed sessionId,
        uint256 indexed rootKeyId,
        address indexed sessionAddr,
        bytes32 scopeHash,
        uint256 spendCap,
        uint64 expiry
    );
    event SessionRevoked(uint256 indexed sessionId, uint256 indexed rootKeyId, address by);
    event ActionAuthorized(
        uint256 indexed sessionId, uint256 indexed rootKeyId, bytes32 indexed actionHash, uint256 amount, uint256 spent
    );

    error ZeroRegistry();
    error RootKeyNotFalcon(uint256 rootKeyId, uint8 scheme);
    error RootKeyNotActive(uint256 rootKeyId);
    error NotRootController(uint256 rootKeyId);
    error AgentExists(uint256 rootKeyId);
    error UnknownAgent(uint256 rootKeyId);
    error UnknownSession(uint256 sessionId);
    error SessionAuthFailed();
    error SessionExpiredOrInvalid(uint256 sessionId);
    error SessionScopeViolation(uint256 sessionId);
    error SessionSpendExceeded(uint256 sessionId);
    error SessionSigInvalid(uint256 sessionId);
    error InvalidSessionParams();
    error NotSessionRevoker(uint256 sessionId);

    constructor(address keyRegistry_, address agentBond_, address aiReputation_) {
        if (keyRegistry_ == address(0)) revert ZeroRegistry();
        keyRegistry = IAerePQCKeyRegistry(keyRegistry_);
        agentBond = agentBond_;
        aiReputation = aiReputation_;
    }

    // ==========================================================================
    //                              Agent creation
    // ==========================================================================

    /**
     * @notice Create the DID for the Falcon root key `rootKeyId`. Callable only by the
     *         registry owner of that key (the party who proved Falcon possession at
     *         registration). The root must be an ACTIVE Falcon-512/1024 key. One DID per
     *         root key; the agent id IS the rootKeyId.
     */
    function createAgent(uint256 rootKeyId) external {
        (address rootOwner, uint8 scheme, uint8 status) = _rootKeyState(rootKeyId);
        if (scheme != SCHEME_FALCON512 && scheme != SCHEME_FALCON1024) revert RootKeyNotFalcon(rootKeyId, scheme);
        if (status != REGISTRY_STATUS_ACTIVE) revert RootKeyNotActive(rootKeyId);
        if (msg.sender != rootOwner) revert NotRootController(rootKeyId);
        if (_agents[rootKeyId].exists) revert AgentExists(rootKeyId);

        _agents[rootKeyId] = Agent({
            rootKeyId: rootKeyId,
            controller: rootOwner,
            createdBlock: uint64(block.number),
            sessionNonce: 0,
            exists: true
        });

        emit AgentCreated(rootKeyId, rootOwner, scheme);
    }

    // ==========================================================================
    //                             Session issuance
    // ==========================================================================

    /**
     * @notice Issue a short-lived secp256k1 session key for an agent. Authorized by a
     *         Falcon PROOF-OF-POSSESSION: `popSig` must be a valid signature by the agent's
     *         root key over the issuance challenge for the current session nonce. May be
     *         relayed by any msg.sender (only the root-key holder can produce popSig). The
     *         root key must still be ACTIVE Falcon at issuance time.
     * @param rootKeyId   the agent (root Falcon key) issuing the session.
     * @param sessionAddr the secp256k1 address that will sign this session's actions.
     * @param scopeHash   opaque scope tag the session is locked to (enforced in authorize).
     * @param spendCap    cumulative spend allowed under this session.
     * @param expiry      unix timestamp after which the session is invalid (must be future).
     * @param popSig      the Falcon proof-of-possession envelope over the challenge.
     * @return sessionId  the id of the newly issued session.
     */
    function issueSession(
        uint256 rootKeyId,
        address sessionAddr,
        bytes32 scopeHash,
        uint256 spendCap,
        uint64 expiry,
        bytes calldata popSig
    ) external returns (uint256 sessionId) {
        Agent storage a = _agents[rootKeyId];
        if (!a.exists) revert UnknownAgent(rootKeyId);
        if (sessionAddr == address(0) || expiry <= block.timestamp) revert InvalidSessionParams();

        // Root must still be an ACTIVE Falcon key (it could have been rotated/revoked).
        (, uint8 scheme, uint8 status) = _rootKeyState(rootKeyId);
        if (scheme != SCHEME_FALCON512 && scheme != SCHEME_FALCON1024) revert RootKeyNotFalcon(rootKeyId, scheme);
        if (status != REGISTRY_STATUS_ACTIVE) revert RootKeyNotActive(rootKeyId);

        uint64 nonce = a.sessionNonce;
        bytes32 challenge = _sessionChallenge(rootKeyId, nonce, sessionAddr, scopeHash, spendCap, expiry);

        // Full on-chain Falcon verification via the registry (live native precompile).
        if (!keyRegistry.verifyWithKey(rootKeyId, challenge, popSig)) revert SessionAuthFailed();

        // Effects.
        a.sessionNonce = nonce + 1;
        sessionId = _sessions.length;
        _sessions.push(
            Session({
                rootKeyId: rootKeyId,
                sessionAddr: sessionAddr,
                scopeHash: scopeHash,
                spendCap: spendCap,
                spent: 0,
                issuedBlock: uint64(block.number),
                expiry: expiry,
                actionNonce: 0,
                revoked: false,
                exists: true
            })
        );
        _agentSessions[rootKeyId].push(sessionId);

        emit SessionIssued(sessionId, rootKeyId, sessionAddr, scopeHash, spendCap, expiry);
    }

    // ==========================================================================
    //                             Session revocation
    // ==========================================================================

    /**
     * @notice Revoke a session immediately. Fast path for incident response: callable by
     *         the agent's controller (the registry owner of the root key) with a cheap
     *         ECDSA tx, no Falcon signature required. Idempotent-safe: revoking an already
     *         revoked session is a no-op success. Sessions also auto-expire at `expiry`.
     */
    function revokeSession(uint256 sessionId) external {
        if (sessionId >= _sessions.length) revert UnknownSession(sessionId);
        Session storage s = _sessions[sessionId];
        // Only the current registry owner of the root key may fast-revoke.
        address controller = _currentController(s.rootKeyId);
        if (msg.sender != controller) revert NotSessionRevoker(sessionId);

        if (!s.revoked) {
            s.revoked = true;
            emit SessionRevoked(sessionId, s.rootKeyId, msg.sender);
        }
    }

    // ==========================================================================
    //                          Action authorization
    // ==========================================================================

    /**
     * @notice Authorize (and record) an agent action signed by a session key. The hot
     *         path: verifies a secp256k1 signature by the session address over the action
     *         digest (ecrecover), enforces the session SCOPE and SPEND CAP, and records
     *         the spend. Reverts (recording nothing) if the session is invalid/expired/
     *         revoked, its root key is no longer ACTIVE Falcon, the scope does not match,
     *         the cap would be exceeded, or the signature is not by the session key.
     * @param sessionId  the session authorizing the action.
     * @param scope      the scope tag for this action (must equal the session's scopeHash).
     * @param actionHash a 32-byte commitment to the action being authorized.
     * @param amount     the spend to charge against the session's cap.
     * @param sig        65-byte secp256k1 signature (r||s||v) by the session address over
     *                   actionDigest(sessionId, scope, actionHash, amount, actionNonce).
     */
    function authorize(uint256 sessionId, bytes32 scope, bytes32 actionHash, uint256 amount, bytes calldata sig)
        external
        returns (bool)
    {
        if (sessionId >= _sessions.length) revert UnknownSession(sessionId);
        Session storage s = _sessions[sessionId];
        if (!_isSessionLive(s)) revert SessionExpiredOrInvalid(sessionId);
        if (scope != s.scopeHash) revert SessionScopeViolation(sessionId);
        if (s.spent + amount < s.spent || s.spent + amount > s.spendCap) revert SessionSpendExceeded(sessionId);

        uint64 nonce = s.actionNonce;
        bytes32 digest = _actionDigest(sessionId, scope, actionHash, amount, nonce);
        address signer = _recover(digest, sig);
        if (signer == address(0) || signer != s.sessionAddr) revert SessionSigInvalid(sessionId);

        // Effects.
        s.actionNonce = nonce + 1;
        s.spent += amount;

        emit ActionAuthorized(sessionId, s.rootKeyId, actionHash, amount, s.spent);
        return true;
    }

    // ==========================================================================
    //                          Validity / digest views
    // ==========================================================================

    /**
     * @notice True iff the session exists, is not revoked, has not expired, and its root
     *         key is still ACTIVE Falcon in the registry.
     */
    function isSessionValid(uint256 sessionId) external view returns (bool) {
        if (sessionId >= _sessions.length) return false;
        return _isSessionLive(_sessions[sessionId]);
    }

    function _isSessionLive(Session storage s) internal view returns (bool) {
        if (!s.exists || s.revoked || block.timestamp > s.expiry) return false;
        (, uint8 scheme, uint8 status) = _rootKeyStateNoRevert(s.rootKeyId);
        if (status != REGISTRY_STATUS_ACTIVE) return false;
        if (scheme != SCHEME_FALCON512 && scheme != SCHEME_FALCON1024) return false;
        return true;
    }

    /// @notice The exact Falcon issuance challenge for a session's parameters at `nonce`.
    function sessionChallenge(
        uint256 rootKeyId,
        uint64 nonce,
        address sessionAddr,
        bytes32 scopeHash,
        uint256 spendCap,
        uint64 expiry
    ) external view returns (bytes32) {
        return _sessionChallenge(rootKeyId, nonce, sessionAddr, scopeHash, spendCap, expiry);
    }

    /// @notice The current issuance challenge for an agent (uses its current sessionNonce).
    function currentSessionChallenge(
        uint256 rootKeyId,
        address sessionAddr,
        bytes32 scopeHash,
        uint256 spendCap,
        uint64 expiry
    ) external view returns (bytes32) {
        return _sessionChallenge(rootKeyId, _agents[rootKeyId].sessionNonce, sessionAddr, scopeHash, spendCap, expiry);
    }

    /// @notice The digest a session key must sign to authorize an action at `actionNonce`.
    function actionDigest(uint256 sessionId, bytes32 scope, bytes32 actionHash, uint256 amount, uint64 actionNonce)
        external
        view
        returns (bytes32)
    {
        return _actionDigest(sessionId, scope, actionHash, amount, actionNonce);
    }

    function _sessionChallenge(
        uint256 rootKeyId,
        uint64 nonce,
        address sessionAddr,
        bytes32 scopeHash,
        uint256 spendCap,
        uint64 expiry
    ) internal view returns (bytes32) {
        return keccak256(
            abi.encode(
                SESSION_DOMAIN, block.chainid, address(this), rootKeyId, nonce, sessionAddr, scopeHash, spendCap, expiry
            )
        );
    }

    function _actionDigest(uint256 sessionId, bytes32 scope, bytes32 actionHash, uint256 amount, uint64 actionNonce)
        internal
        view
        returns (bytes32)
    {
        return keccak256(
            abi.encode(ACTION_DOMAIN, block.chainid, address(this), sessionId, scope, actionHash, amount, actionNonce)
        );
    }

    // ==========================================================================
    //                                Registry reads
    // ==========================================================================

    function _rootKeyState(uint256 rootKeyId) internal view returns (address owner, uint8 scheme, uint8 status) {
        (owner, scheme, status,,,,) = keyRegistry.getKey(rootKeyId);
    }

    /// @dev Like _rootKeyState but never reverts (an out-of-range key on the registry
    ///      would revert getKey). Returns status 0 (not ACTIVE) so callers fail closed.
    function _rootKeyStateNoRevert(uint256 rootKeyId)
        internal
        view
        returns (address owner, uint8 scheme, uint8 status)
    {
        try keyRegistry.getKey(rootKeyId) returns (
            address o, uint8 sc, uint8 st, uint64, uint256, uint256, bytes memory
        ) {
            return (o, sc, st);
        } catch {
            return (address(0), 0, 0);
        }
    }

    /// @dev The address currently allowed to fast-revoke: the registry owner of the root
    ///      key (falls back to the cached controller if the registry read reverts).
    function _currentController(uint256 rootKeyId) internal view returns (address) {
        // A rotated/revoked root still lets its current registry owner fast-revoke.
        (address owner,,) = _rootKeyStateNoRevert(rootKeyId);
        if (owner != address(0)) return owner;
        // Registry read failed; fall back to the controller cached at creation.
        return _agents[rootKeyId].controller;
    }

    // ==========================================================================
    //                              ECDSA recovery
    // ==========================================================================

    /// @dev Minimal ecrecover with malleability guard (low-s, v in {27,28}). Returns
    ///      address(0) on any malformed input so callers treat it as an invalid signature.
    function _recover(bytes32 digest, bytes calldata sig) internal pure returns (address) {
        if (sig.length != 65) return address(0);
        bytes32 r;
        bytes32 s;
        uint8 v;
        assembly {
            r := calldataload(sig.offset)
            s := calldataload(add(sig.offset, 32))
            v := byte(0, calldataload(add(sig.offset, 64)))
        }
        // Reject high-s (EIP-2) to prevent signature malleability.
        if (uint256(s) > 0x7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF5D576E7357A4501DDFE92F46681B20A0) {
            return address(0);
        }
        if (v != 27 && v != 28) return address(0);
        return ecrecover(digest, v, r, s);
    }

    // ==========================================================================
    //                                  Views
    // ==========================================================================

    /// @notice Total sessions ever issued (sessionIds run 0..sessionCount-1).
    function sessionCount() external view returns (uint256) {
        return _sessions.length;
    }

    /// @notice Read an agent (DID). Reverts if the agent does not exist.
    function getAgent(uint256 rootKeyId)
        external
        view
        returns (address controller, uint64 createdBlock, uint64 sessionNonce, uint256 sessionCountForAgent)
    {
        Agent storage a = _agents[rootKeyId];
        if (!a.exists) revert UnknownAgent(rootKeyId);
        return (a.controller, a.createdBlock, a.sessionNonce, _agentSessions[rootKeyId].length);
    }

    /// @notice True iff a DID exists for `rootKeyId`.
    function agentExists(uint256 rootKeyId) external view returns (bool) {
        return _agents[rootKeyId].exists;
    }

    /// @notice Read a session. Reverts on an unknown sessionId.
    function getSession(uint256 sessionId)
        external
        view
        returns (
            uint256 rootKeyId,
            address sessionAddr,
            bytes32 scopeHash,
            uint256 spendCap,
            uint256 spent,
            uint64 issuedBlock,
            uint64 expiry,
            uint64 actionNonce,
            bool revoked
        )
    {
        if (sessionId >= _sessions.length) revert UnknownSession(sessionId);
        Session storage s = _sessions[sessionId];
        return (
            s.rootKeyId,
            s.sessionAddr,
            s.scopeHash,
            s.spendCap,
            s.spent,
            s.issuedBlock,
            s.expiry,
            s.actionNonce,
            s.revoked
        );
    }

    /// @notice The remaining spend under a session (0 if invalid/exhausted).
    function remainingSpend(uint256 sessionId) external view returns (uint256) {
        if (sessionId >= _sessions.length) return 0;
        Session storage s = _sessions[sessionId];
        if (!_isSessionLive(s) || s.spent >= s.spendCap) return 0;
        return s.spendCap - s.spent;
    }

    /// @notice All sessionIds ever issued for an agent (append-only, includes revoked).
    function sessionsOf(uint256 rootKeyId) external view returns (uint256[] memory) {
        return _agentSessions[rootKeyId];
    }
}
