// SPDX-License-Identifier: MIT
pragma solidity ^0.8.26;

import "./IVaultfireIdentity.sol";

/// @title KeyVault — On-chain spending policy enforcement for AI agent wallets
/// @author Vaultfire Protocol
/// @notice Wraps an agent's wallet with enforced spending limits, contract
///         allowlists, and emergency freeze. Policies are on-chain — no prompt
///         injection or client-side bypass can override them.
/// @dev Requires the owner to be registered in the Vaultfire Identity Registry (ERC-8004).
contract KeyVault {
    // ═══════════════════════════════════════════════════
    // REENTRANCY GUARD
    // ═══════════════════════════════════════════════════

    uint256 private constant _NOT_ENTERED = 1;
    uint256 private constant _ENTERED = 2;
    uint256 private _reentrancyStatus = _NOT_ENTERED;

    modifier nonReentrant() {
        require(_reentrancyStatus != _ENTERED, "ReentrancyGuard: reentrant call");
        _reentrancyStatus = _ENTERED;
        _;
        _reentrancyStatus = _NOT_ENTERED;
    }

    // ═══════════════════════════════════════════════════
    // STRUCTS
    // ═══════════════════════════════════════════════════

    struct Policy {
        uint256 maxPerTx;           // Max native token per single transaction (wei)
        uint256 maxDaily;           // Max native token per rolling 24h window (wei)
        uint256 maxHourly;          // Max native token per rolling 1h window (wei)
        uint256 emergencyThreshold; // Freeze if total spent exceeds this (wei)
        bool frozen;                // Emergency freeze flag
    }

    struct SpendingWindow {
        uint256 hourlySpent;
        uint256 dailySpent;
        uint256 totalSpent;
        uint256 hourlyResetAt;
        uint256 dailyResetAt;
    }

    // ═══════════════════════════════════════════════════
    // STATE
    // ═══════════════════════════════════════════════════

    address public immutable owner;
    IVaultfireIdentity public immutable identityRegistry;

    Policy public policy;
    SpendingWindow public spending;

    mapping(address => bool) public allowedContracts;
    uint256 public allowedContractCount;

    bool public allowlistEnabled;

    // ═══════════════════════════════════════════════════
    // EVENTS
    // ═══════════════════════════════════════════════════

    event VaultCreated(address indexed owner, uint256 maxDaily, uint256 maxPerTx);
    event TransactionExecuted(address indexed to, uint256 value, bytes data);
    event PolicyUpdated(uint256 maxPerTx, uint256 maxDaily, uint256 maxHourly);
    event ContractAllowed(address indexed target);
    event ContractRemoved(address indexed target);
    event EmergencyFreeze(uint256 totalSpent, uint256 threshold);
    event EmergencyUnfreeze(address indexed owner);
    event FundsDeposited(address indexed from, uint256 amount);
    event FundsWithdrawn(address indexed to, uint256 amount);

    // ═══════════════════════════════════════════════════
    // ERRORS
    // ═══════════════════════════════════════════════════

    error NotOwner();
    error NotRegistered();
    error VaultFrozen();
    error ExceedsPerTxLimit(uint256 requested, uint256 limit);
    error ExceedsHourlyLimit(uint256 wouldSpend, uint256 limit);
    error ExceedsDailyLimit(uint256 wouldSpend, uint256 limit);
    error ContractNotAllowed(address target);
    error TransactionFailed();
    error ZeroAddress();

    // ═══════════════════════════════════════════════════
    // MODIFIERS
    // ═══════════════════════════════════════════════════

    modifier onlyOwner() {
        if (msg.sender != owner) revert NotOwner();
        _;
    }

    modifier notFrozen() {
        if (policy.frozen) revert VaultFrozen();
        _;
    }

    // ═══════════════════════════════════════════════════
    // CONSTRUCTOR
    // ═══════════════════════════════════════════════════

    /// @notice Create a new KeyVault for a registered Vaultfire agent
    /// @param _owner The agent address that will own this vault (must be ERC-8004 registered)
    /// @param _identityRegistry Address of the Vaultfire Identity Registry
    /// @param _maxPerTx Maximum native token per transaction (wei)
    /// @param _maxHourly Maximum native token per rolling hour (wei)
    /// @param _maxDaily Maximum native token per rolling 24 hours (wei)
    /// @param _emergencyThreshold Freeze vault if total spent exceeds this (wei)
    constructor(
        address _owner,
        address _identityRegistry,
        uint256 _maxPerTx,
        uint256 _maxHourly,
        uint256 _maxDaily,
        uint256 _emergencyThreshold
    ) {
        if (_owner == address(0)) revert ZeroAddress();
        if (_identityRegistry == address(0)) revert ZeroAddress();

        identityRegistry = IVaultfireIdentity(_identityRegistry);

        // Verify owner is a registered ERC-8004 agent
        (string memory agentURI,,,) = identityRegistry.getAgent(_owner);
        if (bytes(agentURI).length == 0) revert NotRegistered();

        owner = _owner;
        policy = Policy({
            maxPerTx: _maxPerTx,
            maxDaily: _maxDaily,
            maxHourly: _maxHourly,
            emergencyThreshold: _emergencyThreshold,
            frozen: false
        });

        spending = SpendingWindow({
            hourlySpent: 0,
            dailySpent: 0,
            totalSpent: 0,
            hourlyResetAt: block.timestamp + 1 hours,
            dailyResetAt: block.timestamp + 1 days
        });

        allowlistEnabled = false;

        emit VaultCreated(_owner, _maxDaily, _maxPerTx);
    }

    // ═══════════════════════════════════════════════════
    // RECEIVE / DEPOSIT
    // ═══════════════════════════════════════════════════

    receive() external payable {
        emit FundsDeposited(msg.sender, msg.value);
    }

    /// @notice Deposit native tokens into the vault
    function deposit() external payable {
        emit FundsDeposited(msg.sender, msg.value);
    }

    // ═══════════════════════════════════════════════════
    // EXECUTE — Policy-enforced transaction execution
    // ═══════════════════════════════════════════════════

    /// @notice Execute a transaction through the vault with policy enforcement
    /// @param to Target address
    /// @param value Native token value to send (wei)
    /// @param data Calldata for the target contract
    /// @return result The return data from the call
    function execute(
        address to,
        uint256 value,
        bytes calldata data
    ) external onlyOwner notFrozen nonReentrant returns (bytes memory result) {
        if (to == address(0)) revert ZeroAddress();

        // Check contract allowlist
        if (allowlistEnabled && !allowedContracts[to]) {
            revert ContractNotAllowed(to);
        }

        // Check per-transaction limit
        if (value > policy.maxPerTx) {
            revert ExceedsPerTxLimit(value, policy.maxPerTx);
        }

        // Reset rolling windows if expired
        _resetWindows();

        // Check hourly limit
        if (spending.hourlySpent + value > policy.maxHourly) {
            revert ExceedsHourlyLimit(spending.hourlySpent + value, policy.maxHourly);
        }

        // Check daily limit
        if (spending.dailySpent + value > policy.maxDaily) {
            revert ExceedsDailyLimit(spending.dailySpent + value, policy.maxDaily);
        }

        // Update spending (effects before interactions — CEI pattern)
        spending.hourlySpent += value;
        spending.dailySpent += value;
        spending.totalSpent += value;

        // Check emergency threshold
        if (policy.emergencyThreshold > 0 && spending.totalSpent > policy.emergencyThreshold) {
            policy.frozen = true;
            emit EmergencyFreeze(spending.totalSpent, policy.emergencyThreshold);
        }

        // Emit event BEFORE external call (checks-effects-interactions)
        emit TransactionExecuted(to, value, data);

        // Execute the actual transaction (interaction — last step)
        bool success;
        (success, result) = to.call{value: value}(data);
        if (!success) revert TransactionFailed();
    }

    // ═══════════════════════════════════════════════════
    // POLICY MANAGEMENT
    // ═══════════════════════════════════════════════════

    /// @notice Update spending limits (only owner)
    function updatePolicy(
        uint256 _maxPerTx,
        uint256 _maxHourly,
        uint256 _maxDaily,
        uint256 _emergencyThreshold
    ) external onlyOwner {
        policy.maxPerTx = _maxPerTx;
        policy.maxHourly = _maxHourly;
        policy.maxDaily = _maxDaily;
        policy.emergencyThreshold = _emergencyThreshold;
        emit PolicyUpdated(_maxPerTx, _maxDaily, _maxHourly);
    }

    /// @notice Add a contract to the allowlist
    function allowContract(address target) external onlyOwner {
        if (target == address(0)) revert ZeroAddress();
        if (!allowedContracts[target]) {
            allowedContracts[target] = true;
            allowedContractCount++;
            allowlistEnabled = true;
            emit ContractAllowed(target);
        }
    }

    /// @notice Remove a contract from the allowlist
    function removeContract(address target) external onlyOwner {
        if (allowedContracts[target]) {
            allowedContracts[target] = false;
            allowedContractCount--;
            if (allowedContractCount == 0) {
                allowlistEnabled = false;
            }
            emit ContractRemoved(target);
        }
    }

    /// @notice Disable the contract allowlist (allow all targets)
    function disableAllowlist() external onlyOwner {
        allowlistEnabled = false;
    }

    /// @notice Enable the contract allowlist
    function enableAllowlist() external onlyOwner {
        allowlistEnabled = true;
    }

    // ═══════════════════════════════════════════════════
    // EMERGENCY CONTROLS
    // ═══════════════════════════════════════════════════

    /// @notice Emergency freeze — stops all transactions immediately
    function freeze() external onlyOwner {
        policy.frozen = true;
        emit EmergencyFreeze(spending.totalSpent, policy.emergencyThreshold);
    }

    /// @notice Unfreeze the vault and reset total spending counter
    function unfreeze() external onlyOwner {
        policy.frozen = false;
        spending.totalSpent = 0;
        emit EmergencyUnfreeze(owner);
    }

    /// @notice Withdraw all funds (emergency exit — only owner)
    function emergencyWithdraw() external onlyOwner nonReentrant {
        uint256 currentBalance = address(this).balance;
        emit FundsWithdrawn(owner, currentBalance);
        (bool success, ) = owner.call{value: currentBalance}("");
        if (!success) revert TransactionFailed();
    }

    // ═══════════════════════════════════════════════════
    // VIEW FUNCTIONS
    // ═══════════════════════════════════════════════════

    /// @notice Get remaining hourly allowance
    function remainingHourly() external view returns (uint256) {
        if (block.timestamp >= spending.hourlyResetAt) return policy.maxHourly;
        if (spending.hourlySpent >= policy.maxHourly) return 0;
        return policy.maxHourly - spending.hourlySpent;
    }

    /// @notice Get remaining daily allowance
    function remainingDaily() external view returns (uint256) {
        if (block.timestamp >= spending.dailyResetAt) return policy.maxDaily;
        if (spending.dailySpent >= policy.maxDaily) return 0;
        return policy.maxDaily - spending.dailySpent;
    }

    /// @notice Check if the vault is currently frozen
    function isFrozen() external view returns (bool) {
        return policy.frozen;
    }

    /// @notice Get the vault's balance
    function balance() external view returns (uint256) {
        return address(this).balance;
    }

    // ═══════════════════════════════════════════════════
    // INTERNAL
    // ═══════════════════════════════════════════════════

    function _resetWindows() internal {
        if (block.timestamp >= spending.hourlyResetAt) {
            spending.hourlySpent = 0;
            spending.hourlyResetAt = block.timestamp + 1 hours;
        }
        if (block.timestamp >= spending.dailyResetAt) {
            spending.dailySpent = 0;
            spending.dailyResetAt = block.timestamp + 1 days;
        }
    }
}
