"""Axon AutoGen tools — FunctionTool wrappers for vault operations."""

from __future__ import annotations

from autogen_core.tools import FunctionTool

# Lazy-initialized shared client (set by get_axon_tools)
_client = None
_chain_id = None


def _get_client():
    if _client is None:
        raise RuntimeError("Axon tools not initialized. Call get_axon_tools() first.")
    return _client


def _format_result(result) -> str:
    """Format a PaymentResult into an LLM-friendly string."""
    if result.status == "approved":
        return f"Approved! TX: {result.tx_hash}"
    elif result.status == "pending_review":
        return f"Under review (request ID: {result.request_id}). Poll for status later."
    else:
        return f"Rejected: {result.reason}"


# ── Tool functions ───────────────────────────────────────────────────────────


def pay(to: str, amount: str, memo: str = "") -> str:
    """Send a payment from the Axon vault.

    Args:
        to: Recipient address (0x...).
        amount: Human-readable USDC amount (e.g. "5.00" for 5 USDC).
        memo: Optional payment description.

    Returns:
        Transaction result: approved with TX hash, pending review, or rejected with reason.
    """
    client = _get_client()
    kwargs: dict = {"to": to, "token": "USDC", "amount": float(amount)}
    if memo:
        kwargs["memo"] = memo
    result = client.pay(**kwargs)
    return _format_result(result)


def swap(from_token: str, to_token: str, amount: str) -> str:
    """Swap tokens within the Axon vault (in-vault rebalancing).

    The output tokens stay in the vault. Use this to convert between tokens
    (e.g. swap USDC for WETH) without sending funds externally.

    Args:
        from_token: Source token symbol (e.g. "USDC", "WETH").
        to_token: Target token symbol (e.g. "WETH", "USDC").
        amount: Human-readable amount of source token to spend (e.g. "100.00").

    Returns:
        Transaction result: approved with TX hash, pending review, or rejected with reason.
    """
    client = _get_client()
    result = client.swap(
        from_token=from_token,
        to_token=to_token,
        max_from_amount=float(amount),
        min_to_amount=0,
    )
    return _format_result(result)


def execute_protocol(target: str, calldata: str) -> str:
    """Execute a DeFi protocol interaction through the Axon vault.

    The vault approves tokens, executes the call, then revokes approval.
    The target protocol must be approved on the vault.

    Args:
        target: Protocol contract address (0x...).
        calldata: ABI-encoded function calldata (0x...).

    Returns:
        Transaction result: approved with TX hash, pending review, or rejected with reason.
    """
    client = _get_client()
    result = client.execute(protocol=target, call_data=calldata)
    return _format_result(result)


def get_balance(token: str = "USDC") -> str:
    """Check the vault balance for a specific token.

    Args:
        token: Token symbol (e.g. "USDC", "WETH"). Defaults to "USDC".

    Returns:
        Human-readable balance (e.g. "Vault holds 1250.50 USDC").
    """
    from axonfi import KNOWN_TOKENS, resolve_token

    client = _get_client()
    token_address = resolve_token(token, _chain_id)
    balance_raw = client.get_balance(token_address)

    info = KNOWN_TOKENS.get(token.upper())
    if info:
        human = balance_raw / (10**info.decimals)
        return f"Vault holds {human:.6g} {token.upper()}"
    return f"Vault holds {balance_raw} base units of {token}"


def get_vault_value() -> str:
    """Get the total USD value of the vault with a per-token breakdown.

    Returns:
        Total vault value in USD and each token's balance, price, and value.
    """
    client = _get_client()
    value = client.get_vault_value()
    lines = [f"Total vault value: ${value.total_value_usd:.2f}"]
    for t in value.tokens:
        human_balance = int(t.balance) / (10**t.decimals)
        lines.append(f"  {t.symbol}: {human_balance:.6g} (${t.value_usd:.2f} @ ${t.price_usd:.4g})")
    return "\n".join(lines)


# ── Entry point ──────────────────────────────────────────────────────────────


def get_axon_tools(
    relayer_url: str,
    bot_private_key: str,
    vault_address: str,
    chain_id: int,
) -> list[FunctionTool]:
    """Create Axon tools for use with AutoGen agents.

    Initializes a shared AxonClientSync and returns a list of FunctionTool
    instances ready to be passed to an AutoGen AssistantAgent.

    Args:
        relayer_url: Axon relayer URL (e.g. "https://relay.axonfi.xyz").
        bot_private_key: Bot's private key (0x-prefixed hex string).
        vault_address: Axon vault contract address (0x...).
        chain_id: Chain ID (e.g. 8453 for Base, 42161 for Arbitrum).

    Returns:
        List of 5 FunctionTool instances: pay, swap, execute_protocol,
        get_balance, get_vault_value.
    """
    from axonfi import AxonClientSync

    global _client, _chain_id
    _client = AxonClientSync(
        vault_address=vault_address,
        chain_id=chain_id,
        bot_private_key=bot_private_key,
        relayer_url=relayer_url,
    )
    _chain_id = chain_id

    return [
        FunctionTool(pay, description="Send a USDC payment from the Axon vault to a recipient address"),
        FunctionTool(swap, description="Swap tokens within the vault (e.g. USDC to WETH). Output stays in the vault"),
        FunctionTool(
            execute_protocol,
            description="Execute a DeFi protocol call through the vault (approve, call, revoke pattern)",
        ),
        FunctionTool(get_balance, description="Check the vault balance for a specific token (default: USDC)"),
        FunctionTool(
            get_vault_value,
            description="Get the total USD value of the vault with a per-token breakdown",
        ),
    ]
