"""Drain manager mixin."""

from typing import Any, Dict, Optional, Tuple

from loguru import logger

from iwa.core.contracts.erc20 import ERC20Contract
from iwa.plugins.olas.constants import OLAS_TOKEN_ADDRESS_GNOSIS
from iwa.plugins.olas.contracts.staking import StakingContract, StakingState
from iwa.web.cache import response_cache

# Retry configuration for drain operations after claiming rewards
# Max retries when waiting for claimed rewards to appear in RPC
DRAIN_RETRY_MAX_ATTEMPTS = 6
# Delay in seconds between retry attempts
DRAIN_RETRY_DELAY_SECONDS = 3


class DrainManagerMixin:
    """Mixin for draining and service token management."""

    def claim_rewards(  # noqa: C901
        self, staking_contract: Optional[StakingContract] = None
    ) -> Tuple[bool, int]:
        """Claim staking rewards for the active service.

        The claimed OLAS tokens will be sent to the service's multisig (Safe).

        Args:
            staking_contract: Optional pre-loaded StakingContract. If not provided,
                              it will be loaded from the service's staking_contract_address.

        Returns:
            Tuple of (success, claimed_amount_wei).

        """
        if not self.service:
            logger.error("No active service")
            return False, 0

        if not self.service.staking_contract_address:
            logger.error("Service is not staked")
            return False, 0

        # Load staking contract if not provided
        if not staking_contract:
            try:
                staking_contract = StakingContract(
                    str(self.service.staking_contract_address),
                    chain_name=self.chain_name,
                )
            except Exception as e:
                logger.error(f"Failed to load staking contract: {e}")
                return False, 0

        service_id = self.service.service_id

        # Check if actually staked
        if staking_contract.get_staking_state(service_id) != StakingState.STAKED:
            logger.info("Service not staked, skipping claim")
            return False, 0

        # Check claimable rewards using calculate_staking_reward for accurate value
        # (get_accrued_rewards returns stored value which may be outdated)
        try:
            claimable_rewards = staking_contract.calculate_staking_reward(service_id)
        except Exception:
            # Fallback to stored value if calculation fails
            claimable_rewards = staking_contract.get_accrued_rewards(service_id)

        if claimable_rewards == 0:
            logger.info("No rewards to claim")
            return False, 0

        logger.info(f"Claiming ~{claimable_rewards / 1e18:.4f} OLAS rewards for service {service_id}")

        # Use service owner which holds the reward rights (not necessarily master)
        owner_address = self.service.service_owner_address or self.wallet.master_account.address

        # Prepare and send claim transaction
        claim_tx = staking_contract.prepare_claim_tx(
            from_address=owner_address,
            service_id=service_id,
        )

        if not claim_tx:
            logger.error("Failed to prepare claim transaction")
            return False, 0

        # Simulate transaction to catch revert before sending
        try:
            staking_contract.chain_interface.web3.eth.call(claim_tx)
        except Exception as e:
            logger.warning(f"Claim would revert, skipping: {e}")
            return False, 0

        success, receipt = self.wallet.sign_and_send_transaction(
            claim_tx,
            signer_address_or_tag=owner_address,
            chain_name=self.chain_name,
            tags=["olas_claim_rewards"],
        )
        if not success:
            logger.error("Failed to send claim transaction")
            return False, 0

        events = staking_contract.extract_events(receipt)

        # Extract actual claimed amount from RewardClaimed event
        claimed_amount = claimable_rewards  # Default to estimated
        for event in events:
            if event["name"] == "RewardClaimed":
                # RewardClaimed event has 'amount' or 'reward' field
                claimed_amount = event["args"].get("amount", event["args"].get("reward", claimed_amount))
                break
        else:
            logger.warning("RewardClaimed event not found, using estimated amount")

        logger.info(f"Successfully claimed {claimed_amount / 1e18:.4f} OLAS rewards")

        # Log claim with OLAS amount and EUR price for tax reporting
        try:
            from iwa.core.db import log_transaction
            from iwa.core.pricing import PriceService

            olas_price = PriceService().get_token_price("autonolas", "eur")
            olas_amount = claimed_amount / 1e18
            value_eur = olas_amount * olas_price if olas_price else None

            tx_hash = receipt.get("transactionHash") or getattr(receipt, "transactionHash", None)
            if tx_hash and hasattr(tx_hash, "hex"):
                tx_hash = tx_hash.hex()

            if tx_hash:
                # Resolve to_tag from multisig address
                multisig_tag = (
                    self.wallet.account_service.get_tag_by_address(
                        str(self.service.multisig_address)
                    )
                    or self.service.service_name
                    or None
                )

                log_transaction(
                    tx_hash=tx_hash,
                    from_addr=str(self.service.staking_contract_address),
                    to_addr=str(self.service.multisig_address),
                    token="OLAS",
                    amount_wei=claimed_amount,
                    chain=self.chain_name,
                    from_tag="staking_contract",
                    to_tag=multisig_tag,
                    price_eur=olas_price,
                    value_eur=value_eur,
                    tags=["olas_claim_rewards", "staking_reward"],
                )
        except Exception as e:
            logger.warning(f"Failed to log claim with pricing: {e}")

        # Invalidate caches - rewards and balances changed
        response_cache.invalidate(f"staking_status:{self.service.key}")
        response_cache.invalidate(f"balances:{self.service.key}")

        return True, claimed_amount

    def withdraw_rewards(self) -> Tuple[bool, float]:
        """Withdraw OLAS from the service Safe to the configured withdrawal address.

        The OLAS tokens are transferred from the service's multisig to the
        withdrawal_address configured in the OlasConfig.

        Returns:
            Tuple of (success, olas_amount_transferred).

        """
        if not self.service:
            logger.error("No active service")
            return False, 0

        if not self.service.multisig_address:
            logger.error("Service has no multisig address")
            return False, 0

        withdrawal_address = (
            str(self.olas_config.withdrawal_address)
            if self.olas_config.withdrawal_address
            else str(self.wallet.master_account.address)
        )
        multisig_address = str(self.service.multisig_address)

        if not self.olas_config.withdrawal_address:
            logger.warning(
                "No withdrawal address configured. Defaulting to master account: "
                f"{withdrawal_address}"
            )

        # Get OLAS balance of the Safe
        olas_token = ERC20Contract(
            str(OLAS_TOKEN_ADDRESS_GNOSIS),
            chain_name=self.chain_name,
        )

        olas_balance = olas_token.balance_of_wei(multisig_address)
        if olas_balance == 0:
            logger.info("No OLAS balance to withdraw")
            return False, 0

        olas_amount = olas_balance / 1e18
        withdrawal_tag = (
            self.wallet.account_service.get_tag_by_address(withdrawal_address)
            or withdrawal_address
        )
        multisig_tag = (
            self.wallet.account_service.get_tag_by_address(multisig_address)
            or multisig_address
        )

        logger.info(f"Withdrawing {olas_amount:.4f} OLAS from {multisig_tag} to {withdrawal_tag}")

        # Transfer from Safe to withdrawal address
        tx_hash = self.wallet.send(
            from_address_or_tag=multisig_address,
            to_address_or_tag=withdrawal_address,
            amount_wei=olas_balance,
            token_address_or_name=str(OLAS_TOKEN_ADDRESS_GNOSIS),
            chain_name=self.chain_name,
        )

        if not tx_hash:
            logger.error("Failed to transfer OLAS")
            return False, 0

        # Invalidate balances cache - OLAS moved from Safe
        response_cache.invalidate(f"balances:{self.service.key}")

        logger.info(f"Withdrew {olas_amount:.4f} OLAS to {withdrawal_tag}")
        return True, olas_amount

    def drain_service(
        self,
        target_address: Optional[str] = None,
        claim_rewards: bool = True,
    ) -> Dict[str, Dict[str, float]]:
        """Drain all service accounts to a target address.

        This method:
        1. Claims any pending staking rewards (if staked and claim_rewards=True)
        2. Drains the Safe (multisig) - native + OLAS tokens
        3. Drains the Agent account - native + OLAS tokens
        4. Drains the Owner account - native + OLAS tokens

        All assets are transferred to the target address (defaults to master account).

        Args:
            target_address: Address to receive drained funds. Defaults to master account.
            claim_rewards: Whether to claim staking rewards before draining.

        Returns:
            Dict with drained amounts per account.

        """
        if not self.service:
            logger.error("No active service")
            return {}

        target = target_address or self.wallet.master_account.address
        chain = self.chain_name
        drained: Dict[str, Any] = {}

        logger.info(f"Draining service {self.service.key} to {target}")

        # Step 1: Claim rewards if staked
        claimed_rewards = self._claim_rewards_if_needed(claim_rewards)

        # Step 2: Drain the Safe
        safe_result = self._drain_safe_account(target, chain, claimed_rewards)
        if safe_result:
            drained["safe"] = safe_result

        # Step 3: Drain the Agent account
        agent_result = self._drain_agent_account(target, chain)
        if agent_result:
            drained["agent"] = agent_result

        # Step 4: Drain the Owner account
        owner_result = self._drain_owner_account(target, chain)
        if owner_result:
            drained["owner"] = owner_result

        # Handle partial success (rewards claimed but no drain)
        if not drained and claimed_rewards > 0:
            logger.info("Drain returned empty but rewards were claimed. Reporting partial success.")
            drained["safe_rewards_only"] = {"olas": claimed_rewards / 1e18}

        # Invalidate balances cache - multiple accounts drained
        response_cache.invalidate(f"balances:{self.service.key}")

        logger.info(f"Drain complete. Accounts drained: {list(drained.keys())}")
        return drained

    def _claim_rewards_if_needed(self, claim_rewards: bool) -> int:
        """Claim rewards if applicable."""
        if claim_rewards and self.service.staking_contract_address:
            try:
                success, amount = self.claim_rewards()
                if success and amount > 0:
                    logger.info(f"Claimed {amount / 1e18:.4f} OLAS rewards")
                    return amount
            except Exception as e:
                logger.warning(f"Could not claim rewards: {e}")
        return 0

    def _drain_safe_account(self, target: str, chain: str, claimed_rewards: int) -> Optional[Any]:
        """Drain the Safe account with retry logic for rewards."""
        if not self.service.multisig_address:
            return None

        safe_addr = str(self.service.multisig_address)
        logger.info(f"Attempting to drain Safe: {safe_addr}")

        # Retry loop if we claimed rewards to allow for RPC indexing
        max_retries = DRAIN_RETRY_MAX_ATTEMPTS if claimed_rewards > 0 else 1

        for attempt in range(max_retries):
            try:
                result = self.wallet.drain(
                    from_address_or_tag=safe_addr,
                    to_address_or_tag=target,
                    chain_name=chain,
                )
                logger.info(f"Safe drain result (attempt {attempt + 1}): {result}")

                normalized_result = self._normalize_drain_result(result)
                if normalized_result:
                    logger.info(f"Drained Safe: {normalized_result}")
                    return normalized_result

                if attempt < max_retries - 1:
                    logger.info(
                        f"Waiting for rewards to appear in balance (attempt {attempt + 1})..."
                    )
                    import time

                    time.sleep(DRAIN_RETRY_DELAY_SECONDS)

            except Exception as e:
                logger.warning(f"Could not drain Safe: {e}")
                import traceback

                logger.warning(f"Safe traceback: {traceback.format_exc()}")
                if attempt < max_retries - 1:
                    import time

                    time.sleep(DRAIN_RETRY_DELAY_SECONDS)
        return None

    def _drain_agent_account(self, target: str, chain: str) -> Optional[Any]:
        """Drain the Agent account."""
        if not self.service.agent_address:
            return None

        agent_addr = str(self.service.agent_address)
        logger.info(f"Attempting to drain Agent: {agent_addr}")
        try:
            result = self.wallet.drain(
                from_address_or_tag=agent_addr,
                to_address_or_tag=target,
                chain_name=chain,
            )
            logger.info(f"Agent drain result: {result}")
            normalized = self._normalize_drain_result(result)
            if normalized:
                logger.info(f"Drained Agent: {normalized}")
                return normalized
            else:
                logger.warning("Agent drain returned None/empty")
        except Exception as e:
            logger.warning(f"Could not drain Agent: {e}")
            import traceback

            logger.warning(f"Agent traceback: {traceback.format_exc()}")
        return None

    def _drain_owner_account(self, target: str, chain: str) -> Optional[Any]:
        """Drain the Owner account."""
        if not self.service.service_owner_address:
            return None

        owner_addr = str(self.service.service_owner_address)

        # Skip if owner == target (owner is already the destination, e.g., master)
        if owner_addr.lower() == target.lower():
            logger.info("Skipping owner drain: owner is already the target address")
            return None

        logger.info(f"Attempting to drain Owner: {owner_addr}")
        try:
            result = self.wallet.drain(
                from_address_or_tag=owner_addr,
                to_address_or_tag=target,
                chain_name=chain,
            )
            logger.info(f"Owner drain result: {result}")
            normalized = self._normalize_drain_result(result)
            if normalized:
                logger.info(f"Drained Owner: {normalized}")
                return normalized
            else:
                logger.warning("Owner drain returned None/empty")
        except Exception as e:
            logger.warning(f"Could not drain Owner: {e}")
            import traceback

            logger.warning(f"Owner traceback: {traceback.format_exc()}")
        return None

    def _normalize_drain_result(self, result: Any) -> Any:
        """Normalize the result from wallet.drain to a transaction hash string or dict."""
        if not result:
            return None

        # Handle Tuple[bool, dict] from EOA/TransactionService
        if isinstance(result, tuple) and len(result) >= 2:
            success, receipt = result
            if success:
                tx_hash = receipt.get("transactionHash")
                if hasattr(tx_hash, "hex"):
                    return tx_hash.hex()
                return str(tx_hash)
            return None

        return result
