"""KeyVault manager — on-chain spending policy enforcement for AI agent wallets."""

from dataclasses import dataclass

from web3 import Web3

from .constants import EXPLORER_URLS, KEYVAULT_ABI, VKP_FACTORY_ABI


@dataclass
class VaultPolicy:
    """On-chain spending policy for a KeyVault."""

    max_per_tx: int  # wei
    max_hourly: int  # wei
    max_daily: int  # wei
    emergency_threshold: int  # wei
    frozen: bool


@dataclass
class VaultStatus:
    """Current status of a KeyVault."""

    address: str
    owner: str
    balance_wei: int
    balance_eth: float
    remaining_hourly_wei: int
    remaining_daily_wei: int
    policy: VaultPolicy


class KeyVaultManager:
    """Manage KeyVault instances — create, execute, freeze, withdraw.

    A KeyVault wraps an agent's wallet with on-chain spending policies.
    No prompt injection or client-side bypass can override them.

    Example:
        >>> from vkp import VKPClient
        >>> client = VKPClient(rpc_url, network_id)
        >>> vault_mgr = client.keyvault
        >>> vault_address = vault_mgr.create_vault(
        ...     max_per_tx="0.1",
        ...     max_hourly="0.5",
        ...     max_daily="1.0",
        ...     emergency_threshold="5.0",
        ... )
    """

    def __init__(self, w3: Web3, network_id: str, factory_address: str | None = None):
        self._w3 = w3
        self._network_id = network_id
        self._factory_address = factory_address
        self._explorer = EXPLORER_URLS.get(network_id, "")

    def _explorer_link(self, tx_hash: str) -> str:
        if self._explorer:
            return f"{self._explorer}/tx/{tx_hash}"
        return tx_hash

    def create_vault(
        self,
        from_address: str,
        max_per_tx: str = "0.1",
        max_hourly: str = "0.5",
        max_daily: str = "1.0",
        emergency_threshold: str = "5.0",
    ) -> dict:
        """Create a new KeyVault with spending policies.

        Args:
            from_address: The agent's address (must be ERC-8004 registered).
            max_per_tx: Max ETH per single transaction.
            max_hourly: Max ETH per rolling hour.
            max_daily: Max ETH per rolling 24 hours.
            emergency_threshold: Freeze vault if total spent exceeds this.

        Returns:
            dict with vault_address, tx_hash, and explorer_link.
        """
        if not self._factory_address:
            raise ValueError("VKP Factory address not configured for this network.")

        factory = self._w3.eth.contract(
            address=Web3.to_checksum_address(self._factory_address),
            abi=VKP_FACTORY_ABI,
        )

        tx = factory.functions.createVault(
            Web3.to_wei(max_per_tx, "ether"),
            Web3.to_wei(max_hourly, "ether"),
            Web3.to_wei(max_daily, "ether"),
            Web3.to_wei(emergency_threshold, "ether"),
        ).build_transaction({
            "from": Web3.to_checksum_address(from_address),
            "nonce": self._w3.eth.get_transaction_count(
                Web3.to_checksum_address(from_address)
            ),
        })

        return {"unsigned_tx": tx, "description": "Create KeyVault"}

    def get_status(self, vault_address: str) -> VaultStatus:
        """Get the current status of a KeyVault.

        Args:
            vault_address: The KeyVault contract address.

        Returns:
            VaultStatus with balance, policy, and remaining allowances.
        """
        vault = self._w3.eth.contract(
            address=Web3.to_checksum_address(vault_address),
            abi=KEYVAULT_ABI,
        )

        policy_data = vault.functions.policy().call()
        balance = vault.functions.balance().call()
        owner = vault.functions.owner().call()
        remaining_hourly = vault.functions.remainingHourly().call()
        remaining_daily = vault.functions.remainingDaily().call()

        policy = VaultPolicy(
            max_per_tx=policy_data[0],
            max_daily=policy_data[1],
            max_hourly=policy_data[2],
            emergency_threshold=policy_data[3],
            frozen=policy_data[4],
        )

        return VaultStatus(
            address=vault_address,
            owner=owner,
            balance_wei=balance,
            balance_eth=float(Web3.from_wei(balance, "ether")),
            remaining_hourly_wei=remaining_hourly,
            remaining_daily_wei=remaining_daily,
            policy=policy,
        )

    def execute(
        self,
        vault_address: str,
        from_address: str,
        to: str,
        value: str = "0",
        data: bytes = b"",
    ) -> dict:
        """Execute a transaction through the KeyVault with policy enforcement.

        Args:
            vault_address: The KeyVault contract address.
            from_address: The vault owner's address.
            to: Target contract or address.
            value: Native token value in ETH.
            data: Calldata for the target.

        Returns:
            dict with unsigned_tx for signing.
        """
        vault = self._w3.eth.contract(
            address=Web3.to_checksum_address(vault_address),
            abi=KEYVAULT_ABI,
        )

        tx = vault.functions.execute(
            Web3.to_checksum_address(to),
            Web3.to_wei(value, "ether"),
            data,
        ).build_transaction({
            "from": Web3.to_checksum_address(from_address),
            "nonce": self._w3.eth.get_transaction_count(
                Web3.to_checksum_address(from_address)
            ),
        })

        return {"unsigned_tx": tx, "description": f"Execute via KeyVault → {to}"}

    def freeze(self, vault_address: str, from_address: str) -> dict:
        """Emergency freeze — stops all transactions immediately.

        Args:
            vault_address: The KeyVault contract address.
            from_address: The vault owner's address.

        Returns:
            dict with unsigned_tx for signing.
        """
        vault = self._w3.eth.contract(
            address=Web3.to_checksum_address(vault_address),
            abi=KEYVAULT_ABI,
        )

        tx = vault.functions.freeze().build_transaction({
            "from": Web3.to_checksum_address(from_address),
            "nonce": self._w3.eth.get_transaction_count(
                Web3.to_checksum_address(from_address)
            ),
        })

        return {"unsigned_tx": tx, "description": "Emergency freeze KeyVault"}

    def unfreeze(self, vault_address: str, from_address: str) -> dict:
        """Unfreeze vault and reset total spending counter.

        Args:
            vault_address: The KeyVault contract address.
            from_address: The vault owner's address.

        Returns:
            dict with unsigned_tx for signing.
        """
        vault = self._w3.eth.contract(
            address=Web3.to_checksum_address(vault_address),
            abi=KEYVAULT_ABI,
        )

        tx = vault.functions.unfreeze().build_transaction({
            "from": Web3.to_checksum_address(from_address),
            "nonce": self._w3.eth.get_transaction_count(
                Web3.to_checksum_address(from_address)
            ),
        })

        return {"unsigned_tx": tx, "description": "Unfreeze KeyVault"}

    def emergency_withdraw(self, vault_address: str, from_address: str) -> dict:
        """Withdraw all funds from vault (emergency exit).

        Args:
            vault_address: The KeyVault contract address.
            from_address: The vault owner's address.

        Returns:
            dict with unsigned_tx for signing.
        """
        vault = self._w3.eth.contract(
            address=Web3.to_checksum_address(vault_address),
            abi=KEYVAULT_ABI,
        )

        tx = vault.functions.emergencyWithdraw().build_transaction({
            "from": Web3.to_checksum_address(from_address),
            "nonce": self._w3.eth.get_transaction_count(
                Web3.to_checksum_address(from_address)
            ),
        })

        return {"unsigned_tx": tx, "description": "Emergency withdraw all funds"}

    def allow_contract(self, vault_address: str, from_address: str, target: str) -> dict:
        """Add a contract to the vault's allowlist.

        Args:
            vault_address: The KeyVault contract address.
            from_address: The vault owner's address.
            target: The contract address to allow.

        Returns:
            dict with unsigned_tx for signing.
        """
        vault = self._w3.eth.contract(
            address=Web3.to_checksum_address(vault_address),
            abi=KEYVAULT_ABI,
        )

        tx = vault.functions.allowContract(
            Web3.to_checksum_address(target)
        ).build_transaction({
            "from": Web3.to_checksum_address(from_address),
            "nonce": self._w3.eth.get_transaction_count(
                Web3.to_checksum_address(from_address)
            ),
        })

        return {"unsigned_tx": tx, "description": f"Allow contract {target}"}
