from __future__ import annotations

import asyncio
from collections.abc import Callable
from typing import Any

from eth_utils import to_checksum_address

from wayfinder_paths.core.adapters.BaseAdapter import BaseAdapter, require_wallet
from wayfinder_paths.core.constants.base import MAX_UINT256, SECONDS_PER_YEAR
from wayfinder_paths.core.constants.chains import CHAIN_ID_ETHEREUM
from wayfinder_paths.core.constants.erc20_abi import ERC20_ABI
from wayfinder_paths.core.constants.ethena_abi import ETHENA_SUSDE_VAULT_ABI
from wayfinder_paths.core.constants.ethena_contracts import (
    ETHENA_SUSDE_VAULT_MAINNET,
    ETHENA_USDE_MAINNET,
    ethena_tokens_by_chain_id,
)
from wayfinder_paths.core.utils.interest import apr_to_apy
from wayfinder_paths.core.utils.multicall import (
    Call,
    read_only_calls_multicall_or_gather,
)
from wayfinder_paths.core.utils.tokens import ensure_allowance
from wayfinder_paths.core.utils.transaction import encode_call, send_transaction
from wayfinder_paths.core.utils.web3 import web3_from_chain_id

VESTING_PERIOD_S = 8 * 60 * 60  # 8 hours


class EthenaVaultAdapter(BaseAdapter):
    """
    Ethena sUSDe staking vault adapter (canonical vault on Ethereum mainnet).

    - Deposit: ERC-4626 `deposit` (stake USDe -> receive sUSDe shares)
    - Withdraw: two-step cooldown (`cooldownShares`/`cooldownAssets`) then `unstake`
    """

    adapter_type = "ETHENA"

    def __init__(
        self,
        config: dict[str, Any] | None = None,
        sign_callback: Callable | None = None,
        wallet_address: str | None = None,
    ) -> None:
        super().__init__("ethena_vault_adapter", config)
        self.sign_callback = sign_callback
        self.wallet_address: str | None = (
            to_checksum_address(wallet_address) if wallet_address else None
        )

    async def get_apy(self) -> tuple[bool, float | str]:
        """
        Compute a "spot" supply APY from Ethena's linear vesting model.

        Ethena rewards vest linearly over ~8 hours. We estimate the current
        per-second asset growth as:

            vesting_rate_assets_per_sec = unvested / remaining
            growth_per_sec = vesting_rate_assets_per_sec / total_assets
        """
        try:
            async with web3_from_chain_id(CHAIN_ID_ETHEREUM) as web3:
                vault = web3.eth.contract(
                    address=ETHENA_SUSDE_VAULT_MAINNET,
                    abi=ETHENA_SUSDE_VAULT_ABI,
                )
                (vals, block) = await asyncio.gather(
                    read_only_calls_multicall_or_gather(
                        web3=web3,
                        chain_id=CHAIN_ID_ETHEREUM,
                        calls=[
                            Call(vault, "getUnvestedAmount", postprocess=int),
                            Call(vault, "lastDistributionTimestamp", postprocess=int),
                            Call(vault, "totalAssets", postprocess=int),
                        ],
                        block_identifier="pending",
                    ),
                    web3.eth.get_block("latest"),
                )
                unvested, last_dist, total_assets = vals

                if unvested <= 0 or total_assets <= 0:
                    return True, 0.0

                now_ts = block.get("timestamp") or 0
                elapsed = max(0, now_ts - last_dist)
                remaining = max(0, VESTING_PERIOD_S - elapsed)
                if remaining <= 0:
                    return True, 0.0

                vesting_rate_assets_per_s = unvested / float(remaining)
                apr = (vesting_rate_assets_per_s / float(total_assets)) * float(
                    SECONDS_PER_YEAR
                )
                apy = apr_to_apy(apr)
                return True, float(apy)
        except Exception as exc:
            return False, str(exc)

    async def get_cooldown(
        self,
        *,
        account: str,
    ) -> tuple[bool, dict[str, int] | str]:
        acct = to_checksum_address(account)
        try:
            async with web3_from_chain_id(CHAIN_ID_ETHEREUM) as web3:
                vault = web3.eth.contract(
                    address=ETHENA_SUSDE_VAULT_MAINNET,
                    abi=ETHENA_SUSDE_VAULT_ABI,
                )
                cooldown_end, underlying_amount = await vault.functions.cooldowns(
                    acct
                ).call(block_identifier="pending")
                return True, {
                    "cooldownEnd": cooldown_end or 0,
                    "underlyingAmount": underlying_amount or 0,
                }
        except Exception as exc:
            return False, str(exc)

    async def get_full_user_state(
        self,
        *,
        account: str,
        chain_id: int = CHAIN_ID_ETHEREUM,
        include_apy: bool = False,
        include_zero_positions: bool = False,
        block_identifier: int | str = "pending",
    ) -> tuple[bool, dict[str, Any] | str]:
        acct = to_checksum_address(account)
        cid = int(chain_id)

        try:
            token_addrs = ethena_tokens_by_chain_id(cid)

            if cid == CHAIN_ID_ETHEREUM:
                async with web3_from_chain_id(CHAIN_ID_ETHEREUM) as web3:
                    vault = web3.eth.contract(
                        address=ETHENA_SUSDE_VAULT_MAINNET,
                        abi=ETHENA_SUSDE_VAULT_ABI,
                    )
                    usde = web3.eth.contract(
                        address=to_checksum_address(token_addrs["usde"]),
                        abi=ERC20_ABI,
                    )
                    susde = web3.eth.contract(
                        address=to_checksum_address(token_addrs["susde"]),
                        abi=ERC20_ABI,
                    )

                    if block_identifier == "pending":
                        (
                            usde_balance,
                            susde_balance,
                            cooldown_raw,
                        ) = await read_only_calls_multicall_or_gather(
                            web3=web3,
                            chain_id=CHAIN_ID_ETHEREUM,
                            calls=[
                                Call(
                                    usde,
                                    "balanceOf",
                                    args=(acct,),
                                    postprocess=int,
                                ),
                                Call(
                                    susde,
                                    "balanceOf",
                                    args=(acct,),
                                    postprocess=int,
                                ),
                                Call(vault, "cooldowns", args=(acct,)),
                            ],
                            block_identifier="pending",
                        )
                    else:
                        (
                            usde_balance,
                            susde_balance,
                        ) = await read_only_calls_multicall_or_gather(
                            web3=web3,
                            chain_id=CHAIN_ID_ETHEREUM,
                            calls=[
                                Call(
                                    usde,
                                    "balanceOf",
                                    args=(acct,),
                                    postprocess=int,
                                ),
                                Call(
                                    susde,
                                    "balanceOf",
                                    args=(acct,),
                                    postprocess=int,
                                ),
                            ],
                            block_identifier=block_identifier,
                        )
                        cooldown_raw = await vault.functions.cooldowns(acct).call(
                            block_identifier="pending"
                        )
                    cooldown = {
                        "cooldownEnd": cooldown_raw[0] or 0,
                        "underlyingAmount": cooldown_raw[1] or 0,
                    }
                    shares = susde_balance or 0
                    usde_equivalent = 0
                    if shares > 0:
                        usde_equivalent = (
                            await vault.functions.convertToAssets(shares).call(
                                block_identifier="pending"
                            )
                            or 0
                        )
            else:
                # Balances on the target chain, vault reads on mainnet.
                async with web3_from_chain_id(cid) as web3:
                    usde = web3.eth.contract(
                        address=to_checksum_address(token_addrs["usde"]),
                        abi=ERC20_ABI,
                    )
                    susde = web3.eth.contract(
                        address=to_checksum_address(token_addrs["susde"]),
                        abi=ERC20_ABI,
                    )
                    (
                        usde_balance,
                        susde_balance,
                    ) = await read_only_calls_multicall_or_gather(
                        web3=web3,
                        chain_id=cid,
                        calls=[
                            Call(
                                usde,
                                "balanceOf",
                                args=(acct,),
                                postprocess=int,
                            ),
                            Call(
                                susde,
                                "balanceOf",
                                args=(acct,),
                                postprocess=int,
                            ),
                        ],
                        block_identifier=block_identifier,
                    )

                shares = susde_balance or 0

                async with web3_from_chain_id(CHAIN_ID_ETHEREUM) as web3_hub:
                    vault = web3_hub.eth.contract(
                        address=ETHENA_SUSDE_VAULT_MAINNET,
                        abi=ETHENA_SUSDE_VAULT_ABI,
                    )
                    mc_calls = [Call(vault, "cooldowns", args=(acct,))]
                    if shares > 0:
                        mc_calls.append(
                            Call(
                                vault,
                                "convertToAssets",
                                args=(int(shares),),
                                postprocess=int,
                            )
                        )
                    results = await read_only_calls_multicall_or_gather(
                        web3=web3_hub,
                        chain_id=CHAIN_ID_ETHEREUM,
                        calls=mc_calls,
                        block_identifier="pending",
                    )
                    cooldown_raw = results[0]
                    cooldown = {
                        "cooldownEnd": cooldown_raw[0] or 0,
                        "underlyingAmount": cooldown_raw[1] or 0,
                    }
                    usde_equivalent = int(results[1] or 0) if shares > 0 else 0

            cd_underlying = cooldown.get("underlyingAmount", 0)

            apy_supply: float | None = None
            if include_apy:
                ok_apy, apy_val = await self.get_apy()
                if ok_apy and isinstance(apy_val, (float, int)):
                    apy_supply = float(apy_val)

            include_position = include_zero_positions or shares > 0 or cd_underlying > 0

            positions: list[dict[str, Any]] = []
            if include_position:
                positions.append(
                    {
                        "chainId": cid,
                        "usde": token_addrs["usde"],
                        "susde": token_addrs["susde"],
                        "usdeBalance": usde_balance or 0,
                        "susdeBalance": shares,
                        "usdeEquivalent": usde_equivalent,
                        "cooldown": cooldown,
                        "apySupply": apy_supply,
                        "apyBorrow": None,
                    }
                )

            return (
                True,
                {
                    "protocol": "ethena",
                    "hubChainId": CHAIN_ID_ETHEREUM,
                    "chainId": cid,
                    "account": acct,
                    "positions": positions,
                },
            )
        except Exception as exc:
            return False, str(exc)

    @require_wallet
    async def deposit_usde(
        self,
        *,
        amount_assets: int,
        receiver: str | None = None,
    ) -> tuple[bool, Any]:
        if amount_assets <= 0:
            return False, "amount_assets must be positive"

        recv = to_checksum_address(receiver) if receiver else self.wallet_address

        try:
            approved = await ensure_allowance(
                token_address=ETHENA_USDE_MAINNET,
                owner=self.wallet_address,
                spender=ETHENA_SUSDE_VAULT_MAINNET,
                amount=amount_assets,
                chain_id=CHAIN_ID_ETHEREUM,
                signing_callback=self.sign_callback,
                approval_amount=MAX_UINT256,
            )
            if not approved[0]:
                return approved

            tx = await encode_call(
                target=ETHENA_SUSDE_VAULT_MAINNET,
                abi=ETHENA_SUSDE_VAULT_ABI,
                fn_name="deposit",
                args=[amount_assets, recv],
                from_address=self.wallet_address,
                chain_id=CHAIN_ID_ETHEREUM,
            )
            txn_hash = await send_transaction(tx, self.sign_callback)
            return True, txn_hash
        except Exception as exc:
            return False, str(exc)

    @require_wallet
    async def request_withdraw_by_shares(
        self,
        *,
        shares: int,
    ) -> tuple[bool, Any]:
        if shares <= 0:
            return False, "shares must be positive"

        try:
            tx = await encode_call(
                target=ETHENA_SUSDE_VAULT_MAINNET,
                abi=ETHENA_SUSDE_VAULT_ABI,
                fn_name="cooldownShares",
                args=[shares],
                from_address=self.wallet_address,
                chain_id=CHAIN_ID_ETHEREUM,
            )
            txn_hash = await send_transaction(tx, self.sign_callback)
            return True, txn_hash
        except Exception as exc:
            return False, str(exc)

    @require_wallet
    async def request_withdraw_by_assets(
        self,
        *,
        assets: int,
    ) -> tuple[bool, Any]:
        if assets <= 0:
            return False, "assets must be positive"

        try:
            tx = await encode_call(
                target=ETHENA_SUSDE_VAULT_MAINNET,
                abi=ETHENA_SUSDE_VAULT_ABI,
                fn_name="cooldownAssets",
                args=[assets],
                from_address=self.wallet_address,
                chain_id=CHAIN_ID_ETHEREUM,
            )
            txn_hash = await send_transaction(tx, self.sign_callback)
            return True, txn_hash
        except Exception as exc:
            return False, str(exc)

    @require_wallet
    async def claim_withdraw(
        self,
        *,
        receiver: str | None = None,
        require_matured: bool = True,
    ) -> tuple[bool, Any]:
        recv = to_checksum_address(receiver) if receiver else self.wallet_address

        try:
            ok_cd, cd = await self.get_cooldown(account=self.wallet_address)
            if not ok_cd:
                return False, str(cd)
            if not isinstance(cd, dict):
                return False, "unexpected cooldown payload"

            cooldown_end = cd.get("cooldownEnd") or 0
            underlying_amount = cd.get("underlyingAmount") or 0
            if underlying_amount <= 0:
                return True, "no pending cooldown"

            if require_matured and cooldown_end > 0:
                async with web3_from_chain_id(CHAIN_ID_ETHEREUM) as web3:
                    block = await web3.eth.get_block("latest")
                    now_ts = block.get("timestamp") or 0
                if now_ts < cooldown_end:
                    return (
                        False,
                        f"Cooldown not finished (now={now_ts}, cooldownEnd={cooldown_end})",
                    )

            tx = await encode_call(
                target=ETHENA_SUSDE_VAULT_MAINNET,
                abi=ETHENA_SUSDE_VAULT_ABI,
                fn_name="unstake",
                args=[recv],
                from_address=self.wallet_address,
                chain_id=CHAIN_ID_ETHEREUM,
            )
            txn_hash = await send_transaction(tx, self.sign_callback)
            return True, txn_hash
        except Exception as exc:
            return False, str(exc)
