Wallet

# wallet skill module

"""
Wallet skill exports — for use in task scripts via core.skill_tools.

Usage:
    from core.skill_tools import wallet
    info = wallet.wallet_info()
    bal = wallet.wallet_balance(chain="base")
    all_bal = wallet.wallet_get_all_balances()
"""

import asyncio

try:
    import nest_asyncio
    nest_asyncio.apply()
except ImportError:
    pass

from .tools.common import is_fly_machine, wallet_request, get_wallet_addresses
from .tools.info import get_wallet_info as _get_info
from .tools.balance import get_evm_balance as _get_evm_bal, get_sol_balance as _get_sol_bal, get_all_balances as _get_all
from .tools.transfer import (
    evm_transfer as _evm_transfer,
    evm_sign_transaction as _evm_sign_tx,
    evm_sign_message as _evm_sign_msg,
    evm_sign_typed_data as _evm_sign_typed,
    evm_transactions as _evm_txs,
    sol_transfer as _sol_transfer,
    sol_sign_transaction as _sol_sign_tx,
    sol_sign_message as _sol_sign_msg,
    sol_transactions as _sol_txs,
)
from .tools.policy import get_policy as _get_policy


def _run(coro):
    """Run async in sync context (works even if event loop is running)."""
    try:
        loop = asyncio.get_event_loop()
        return loop.run_until_complete(coro)
    except RuntimeError:
        return asyncio.run(coro)


# ── Info ─────────────────────────────────────────────────────────────────────

def wallet_info():
    """Get all wallet addresses."""
    return _run(_get_info())


# ── Balances ─────────────────────────────────────────────────────────────────

def wallet_balance(chain: str, address: str = "", asset: str = ""):
    """Get EVM balance on a chain."""
    return _run(_get_evm_bal(chain, address, asset))

def wallet_sol_balance(address: str = "", asset: str = ""):
    """Get Solana balance."""
    return _run(_get_sol_bal(address, asset))

def wallet_get_all_balances(evm_address: str = "", sol_address: str = ""):
    """Get balances across all chains."""
    return _run(_get_all(evm_address, sol_address))


# ── Transfers ────────────────────────────────────────────────────────────────

def wallet_transfer(to, amount, chain_id=1, data="", **kw):
    """Sign and broadcast EVM tx."""
    return _run(_evm_transfer(to, amount, chain_id, data, **kw))

def wallet_sign_transaction(to, amount, chain_id=1, data="", **kw):
    """Sign EVM tx without broadcasting."""
    return _run(_evm_sign_tx(to, amount, chain_id, data, **kw))

def wallet_sign(message):
    """EIP-191 personal_sign."""
    return _run(_evm_sign_msg(message))

def wallet_sign_typed_data(domain, types, primaryType, message):
    """Sign EIP-712 typed data."""
    return _run(_evm_sign_typed(domain, types, primaryType, message))

def wallet_transactions(chain="ethereum", asset="eth", limit=20):
    """EVM tx history."""
    return _run(_evm_txs(chain, asset, limit))


# ── Solana ───────────────────────────────────────────────────────────────────

def wallet_sol_transfer(transaction, caip2="solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp"):
    """Sign and broadcast Solana tx."""
    return _run(_sol_transfer(transaction, caip2))

def wallet_sol_sign_transaction(transaction):
    """Sign Solana tx without broadcasting."""
    return _run(_sol_sign_tx(transaction))

def wallet_sol_sign(message):
    """Sign message with Solana wallet."""
    return _run(_sol_sign_msg(message))

def wallet_sol_transactions(chain="solana", asset="sol", limit=20):
    """Solana tx history."""
    return _run(_sol_txs(chain, asset, limit))


# ── Policy ───────────────────────────────────────────────────────────────────

def wallet_get_policy(chain_type="ethereum"):
    """Get current policy status."""
    return _run(_get_policy(chain_type))

# wallet tools module

"""
Balance tools — EVM (DeBank) + Solana (Birdeye) + all-chains aggregation.
"""

import asyncio
import os
import requests

from .common import (
    is_fly_machine, wallet_request, get_wallet_addresses,
    proxied_get_with_retry, DEBANK_CHAIN_MAP, EVM_CHAINS, logger,
)


async def get_evm_balance(chain: str, address: str = "", asset: str = "") -> dict:
    """Get EVM wallet balance on a specific chain."""
    if chain not in EVM_CHAINS:
        return {"error": f"Invalid chain '{chain}'. Must be one of: {', '.join(EVM_CHAINS)}"}

    debank_key = os.environ.get("DEBANK_API_KEY", "")

    if debank_key:
        evm_address = address
        if not evm_address:
            if not is_fly_machine():
                return {"error": "No address and not on Fly Machine"}
            addrs = await get_wallet_addresses()
            evm_address = addrs.get("evm", "")
        if not evm_address:
            return {"error": "Could not determine EVM wallet address"}

        debank_chain_id = DEBANK_CHAIN_MAP.get(chain)
        try:
            resp = proxied_get_with_retry(
                "https://pro-openapi.debank.com/v1/user/token_list",
                params={"id": evm_address, "chain_id": debank_chain_id, "is_all": "false"},
                headers={"AccessKey": debank_key},
            )
            return {"address": evm_address, "chain": chain, "tokens": resp.json(), "source": "debank"}
        except Exception as e:
            return {"error": f"DeBank request failed: {e}"}
    else:
        if not is_fly_machine():
            return {"error": "Not on Fly Machine — wallet unavailable"}
        params = [f"chain_type=ethereum&chain={chain}"]
        if asset:
            params.append(f"asset={asset}")
        qs = "?" + "&".join(params)
        return await wallet_request("GET", f"/agent/balance{qs}")


async def get_sol_balance(address: str = "", asset: str = "") -> dict:
    """Get Solana wallet balance."""
    birdeye_key = os.environ.get("BIRDEYE_API_KEY", "")

    if birdeye_key:
        sol_address = address
        if not sol_address:
            if not is_fly_machine():
                return {"error": "No address and not on Fly Machine"}
            addrs = await get_wallet_addresses()
            sol_address = addrs.get("sol", "")
        if not sol_address:
            return {"error": "Could not determine Solana wallet address"}

        try:
            resp = proxied_get_with_retry(
                "https://public-api.birdeye.so/wallet/v2/net-worth",
                params={"wallet": sol_address},
                headers={"X-API-KEY": birdeye_key, "x-chain": "solana", "accept": "application/json"},
            )
            return {"address": sol_address, "source": "birdeye", "data": resp.json()}
        except Exception as e:
            return {"error": f"Birdeye request failed: {e}"}
    else:
        if not is_fly_machine():
            return {"error": "Not on Fly Machine — wallet unavailable"}
        params = ["chain_type=solana"]
        if asset:
            params.append(f"asset={asset}")
        qs = "?" + "&".join(params)
        return await wallet_request("GET", f"/agent/balance{qs}")


async def get_all_balances(evm_address: str = "", sol_address: str = "") -> dict:
    """Get balances across ALL chains concurrently."""
    if not evm_address or not sol_address:
        if is_fly_machine():
            try:
                addrs = await get_wallet_addresses()
                evm_address = evm_address or addrs.get("evm", "")
                sol_address = sol_address or addrs.get("sol", "")
            except Exception:
                pass

    result = {}
    errors = []
    evm_usd = 0.0
    sol_usd = 0.0

    async def _fetch_evm():
        nonlocal evm_usd
        if not evm_address:
            return
        debank_key = os.environ.get("DEBANK_API_KEY", "")
        if not debank_key:
            errors.append("No DEBANK_API_KEY")
            return
        try:
            resp = proxied_get_with_retry(
                "https://pro-openapi.debank.com/v1/user/all_token_list",
                params={"id": evm_address, "is_all": "true"},
                headers={"AccessKey": debank_key},
            )
            tokens = resp.json()
            by_chain = {}
            for t in tokens:
                chain = t.get("chain", "unknown")
                if chain not in by_chain:
                    by_chain[chain] = {"tokens": [], "total_usd": 0.0}
                usd = t.get("price", 0) * t.get("amount", 0)
                by_chain[chain]["tokens"].append(t)
                by_chain[chain]["total_usd"] = round(by_chain[chain]["total_usd"] + usd, 2)
                evm_usd += usd
            result["evm"] = {"address": evm_address, "chains": by_chain, "total_usd": round(evm_usd, 2), "source": "debank"}
        except Exception as e:
            errors.append(f"DeBank: {e}")

    async def _fetch_sol():
        nonlocal sol_usd
        if not sol_address:
            return
        birdeye_key = os.environ.get("BIRDEYE_API_KEY", "")
        if not birdeye_key:
            errors.append("No BIRDEYE_API_KEY")
            return
        try:
            resp = proxied_get_with_retry(
                "https://public-api.birdeye.so/wallet/v2/net-worth",
                params={"wallet": sol_address},
                headers={"X-API-KEY": birdeye_key, "x-chain": "solana", "accept": "application/json"},
            )
            data = resp.json()
            sol_usd = data.get("data", {}).get("totalUsd", 0)
            result["solana"] = {"address": sol_address, "source": "birdeye", "data": data, "total_usd": round(sol_usd, 2)}
        except Exception as e:
            errors.append(f"Birdeye: {e}")

    await asyncio.gather(_fetch_evm(), _fetch_sol())

    result["total_usd_value"] = round(evm_usd + sol_usd, 2)
    if errors:
        result["errors"] = errors

    has_data = "evm" in result or "solana" in result
    if not has_data and errors:
        return {"error": "All balance queries failed: " + "; ".join(errors)}
    return result

"""
Wallet Tools — shared infrastructure.

Auth: Fly OIDC tokens from /.fly/api Unix Socket.
Each machine has one wallet per chain_type (ethereum, solana).
"""

import asyncio
import logging
import os
import time
from typing import Optional

import aiohttp

# ── Fix aiohttp 3.13.x + Brotli 1.1.0 incompatibility ──────────────────────
try:
    from typing import cast
    import aiohttp.compression_utils as _cu
    _orig = getattr(_cu.BrotliDecompressor, 'decompress_sync', None)
    if _orig:
        def _fixed_decompress_sync(self, data, max_length=0):
            if hasattr(self._obj, 'decompress'):
                return cast(bytes, self._obj.decompress(data))
            return cast(bytes, self._obj.process(data))
        _cu.BrotliDecompressor.decompress_sync = _fixed_decompress_sync
except Exception:
    pass

import requests
from core.http_client import proxied_get

logger = logging.getLogger(__name__)

WALLET_SERVICE_URL = os.environ.get(
    "WALLET_SERVICE_URL", "https://wallet-service-dev.fly.dev"
)
OIDC_AUDIENCE = os.environ.get("WALLET_OIDC_AUDIENCE", WALLET_SERVICE_URL)
FLY_API_SOCKET = "/.fly/api"

DEBANK_CHAIN_MAP = {
    "ethereum": "eth", "base": "base", "arbitrum": "arb",
    "optimism": "op", "polygon": "matic", "linea": "linea",
}

EVM_CHAINS = list(DEBANK_CHAIN_MAP.keys())

# ── OIDC token cache ─────────────────────────────────────────────────────────
_cached_token: Optional[str] = None
_cached_token_exp: float = 0


async def _fetch_oidc_token() -> str:
    conn = aiohttp.UnixConnector(path=FLY_API_SOCKET)
    async with aiohttp.ClientSession(connector=conn) as session:
        async with session.post(
            "http://localhost/v1/tokens/oidc",
            json={"aud": OIDC_AUDIENCE},
            timeout=aiohttp.ClientTimeout(total=5),
        ) as resp:
            resp.raise_for_status()
            return (await resp.text()).strip()


async def ensure_token() -> str:
    global _cached_token, _cached_token_exp
    if not _cached_token or time.time() > _cached_token_exp - 120:
        _cached_token = await _fetch_oidc_token()
        _cached_token_exp = time.time() + 600
    return _cached_token


def is_fly_machine() -> bool:
    return os.path.exists(FLY_API_SOCKET)


async def wallet_request(method: str, path: str, json_body: dict = None) -> dict:
    """Authenticated request to privy-wallet Agent API."""
    token = await ensure_token()
    headers = {"Authorization": f"Bearer {token}", "Accept-Encoding": "gzip, deflate"}
    url = f"{WALLET_SERVICE_URL}{path}"

    async with aiohttp.ClientSession() as session:
        async with session.request(
            method, url,
            headers=headers,
            json=json_body if method in ("POST", "PUT", "PATCH") else None,
            timeout=aiohttp.ClientTimeout(total=30),
        ) as resp:
            body = await resp.json()
            if resp.status >= 400:
                detail = body.get("detail") or body.get("error") or str(body)
                raise Exception(f"Wallet service error ({resp.status}): {detail}")
            return body


# ── Wallet address cache ─────────────────────────────────────────────────────
_addr_cache: dict = {}
_addr_cache_exp: float = 0


async def get_wallet_addresses() -> dict:
    """Return {"evm": "0x...", "sol": "..."} from wallet service, cached 5 min."""
    global _addr_cache, _addr_cache_exp
    if _addr_cache and time.time() < _addr_cache_exp:
        return _addr_cache

    data = await wallet_request("GET", "/agent/wallet")
    wallets = data if isinstance(data, list) else data.get("wallets", [])
    result = {}
    for w in wallets:
        if not isinstance(w, dict):
            continue
        ct = w.get("chain_type", "")
        addr = w.get("wallet_address", "") or w.get("address", "")
        if ct == "ethereum" and addr and "evm" not in result:
            result["evm"] = addr
        elif ct == "solana" and addr and "sol" not in result:
            result["sol"] = addr
    _addr_cache = result
    _addr_cache_exp = time.time() + 300
    return result


def proxied_get_with_retry(url, params=None, headers=None, timeout=30, max_retries=3):
    """proxied_get with retry on timeout / 429 / 5xx."""
    last_exc = None
    for attempt in range(max_retries):
        try:
            resp = proxied_get(url, params=params, headers=headers, timeout=timeout)
            if resp.status_code == 429 or resp.status_code >= 500:
                if attempt < max_retries - 1:
                    time.sleep(1 * (attempt + 1))
                    continue
            resp.raise_for_status()
            return resp
        except requests.exceptions.Timeout as e:
            last_exc = e
            if attempt < max_retries - 1:
                time.sleep(1)
        except requests.exceptions.ConnectionError as e:
            last_exc = e
            if attempt < max_retries - 1:
                time.sleep(1)
        except requests.exceptions.HTTPError:
            raise
        except Exception as e:
            last_exc = e
            break
    raise last_exc or requests.exceptions.RequestException("Max retries exceeded")

"""
Wallet info — addresses and chain info.
"""

from .common import wallet_request


async def get_wallet_info() -> dict:
    """Get all wallet addresses for this agent."""
    return await wallet_request("GET", "/agent/wallet")

"""
Wallet policy — get and propose policy updates.
Policy validation logic for Privy TEE rules.
"""

from .common import wallet_request, logger

# ── Privy Policy Constants ───────────────────────────────────────────────────

VALID_EVM_METHODS = {
    "eth_sendTransaction", "eth_signTransaction", "eth_signTypedData_v4",
    "eth_signUserOperation", "eth_sign7702Authorization",
}
VALID_SOLANA_METHODS = {
    "signTransaction", "signAndSendTransaction", "signTransactionBytes",
}
VALID_ANY_METHODS = {"exportPrivateKey", "*"}
VALID_METHODS = VALID_EVM_METHODS | VALID_SOLANA_METHODS | VALID_ANY_METHODS

EVM_FIELD_SOURCES = {
    "ethereum_transaction", "ethereum_calldata",
    "ethereum_typed_data_domain", "ethereum_typed_data_message",
    "ethereum_7702_authorization",
}
SOLANA_FIELD_SOURCES = {
    "solana_program_instruction", "solana_system_program_instruction",
    "solana_token_program_instruction",
}
ANY_FIELD_SOURCES = {"system"}
VALID_FIELD_SOURCES = EVM_FIELD_SOURCES | SOLANA_FIELD_SOURCES | ANY_FIELD_SOURCES

VALID_OPERATORS = {"eq", "gt", "gte", "lt", "lte", "in", "in_condition_set"}

RULE_KEYS = {"name", "method", "conditions", "action"}
CONDITION_KEYS = {"field_source", "field", "operator", "value", "abi", "typed_data"}

FIELD_SOURCE_VALID_FIELDS = {
    "ethereum_typed_data_domain": {"chainId", "verifyingContract", "chain_id", "verifying_contract"},
    "ethereum_transaction": {"to", "value", "chain_id", "data", "gas", "gas_price", "max_fee_per_gas", "max_priority_fee_per_gas", "nonce"},
}

METHODS_REQUIRING_CONDITIONS = {"eth_signTypedData_v4"}

METHOD_FIELD_SOURCES = {
    "eth_sendTransaction":       {"ethereum_transaction", "ethereum_calldata", "system"},
    "eth_signTransaction":       {"ethereum_transaction", "ethereum_calldata", "system"},
    "eth_signUserOperation":     {"ethereum_transaction", "ethereum_calldata", "system"},
    "eth_signTypedData_v4":      {"ethereum_typed_data_domain", "ethereum_typed_data_message", "system"},
    "eth_sign7702Authorization": {"ethereum_7702_authorization", "system"},
    "signTransaction":           {"solana_program_instruction", "solana_system_program_instruction", "solana_token_program_instruction", "system"},
    "signAndSendTransaction":    {"solana_program_instruction", "solana_system_program_instruction", "solana_token_program_instruction", "system"},
    "signTransactionBytes":      {"solana_program_instruction", "solana_system_program_instruction", "solana_token_program_instruction", "system"},
    "exportPrivateKey":          {"system"},
    "*":                         {"system"},
}


def validate_and_clean_rules(rules: list, chain_type: str) -> tuple:
    """Validate rules against Privy schema. Returns (cleaned_rules, errors)."""
    errors = []
    cleaned = []

    for i, rule in enumerate(rules):
        prefix = f"Rule #{i + 1}"

        if not isinstance(rule, dict):
            errors.append(f"{prefix}: must be an object")
            continue

        name = rule.get("name")
        if not name or not isinstance(name, str):
            errors.append(f"{prefix}: 'name' is required")

        method = rule.get("method")
        if not method or not isinstance(method, str):
            errors.append(f"{prefix}: 'method' is required")
        elif method not in VALID_METHODS:
            errors.append(f"{prefix}: unknown method '{method}'")
        else:
            if chain_type == "ethereum" and method in VALID_SOLANA_METHODS:
                errors.append(f"{prefix}: Solana method '{method}' invalid for ethereum")
            elif chain_type == "solana" and method in VALID_EVM_METHODS:
                errors.append(f"{prefix}: EVM method '{method}' invalid for solana")

        action = rule.get("action")
        if not action or not isinstance(action, str):
            errors.append(f"{prefix}: 'action' required (ALLOW or DENY)")
        else:
            action = action.upper()
            if action not in ("ALLOW", "DENY"):
                errors.append(f"{prefix}: 'action' must be ALLOW or DENY")

        conditions = rule.get("conditions") or []
        if not isinstance(conditions, list):
            errors.append(f"{prefix}: 'conditions' must be a list")
            conditions = []

        cleaned_conditions = []
        for j, cond in enumerate(conditions):
            cp = f"{prefix}, Cond #{j + 1}"
            if not isinstance(cond, dict):
                errors.append(f"{cp}: must be an object")
                continue

            fs = cond.get("field_source")
            if not fs:
                errors.append(f"{cp}: 'field_source' required")
            elif fs not in VALID_FIELD_SOURCES:
                errors.append(f"{cp}: unknown field_source '{fs}'")
            else:
                if chain_type == "ethereum" and fs in SOLANA_FIELD_SOURCES:
                    errors.append(f"{cp}: Solana field_source '{fs}' invalid for ethereum")
                elif chain_type == "solana" and fs in EVM_FIELD_SOURCES:
                    errors.append(f"{cp}: EVM field_source '{fs}' invalid for solana")
                if method and method in METHOD_FIELD_SOURCES:
                    allowed = METHOD_FIELD_SOURCES[method]
                    if fs not in allowed:
                        errors.append(f"{cp}: '{fs}' not valid for method '{method}'")

            field_val = cond.get("field")
            if not field_val:
                errors.append(f"{cp}: 'field' required")
            elif fs and fs in FIELD_SOURCE_VALID_FIELDS:
                if field_val not in FIELD_SOURCE_VALID_FIELDS[fs]:
                    errors.append(f"{cp}: field '{field_val}' invalid for '{fs}'")

            op = cond.get("operator")
            if not op:
                errors.append(f"{cp}: 'operator' required")
            elif op not in VALID_OPERATORS:
                errors.append(f"{cp}: unknown operator '{op}'")

            if "value" not in cond:
                errors.append(f"{cp}: 'value' required")
            else:
                value = cond["value"]
                if op == "in_condition_set":
                    errors.append(f"{cp}: 'in_condition_set' not supported, use 'in' with array")
                elif op == "in":
                    if not isinstance(value, list):
                        errors.append(f"{cp}: 'in' requires array")
                    elif len(value) > 100:
                        errors.append(f"{cp}: 'in' max 100 values")
                    else:
                        cond["value"] = [str(v) for v in value if v is not None]
                elif op and value is not None and not isinstance(value, str):
                    cond["value"] = str(value)

            cleaned_conditions.append({k: v for k, v in cond.items() if k in CONDITION_KEYS})

        if method and method in METHODS_REQUIRING_CONDITIONS and not cleaned_conditions:
            errors.append(f"{prefix}: method '{method}' requires at least one condition")

        cleaned_rule = {k: v for k, v in rule.items() if k in RULE_KEYS}
        cleaned_rule["conditions"] = cleaned_conditions
        if "action" in cleaned_rule and isinstance(cleaned_rule["action"], str):
            cleaned_rule["action"] = cleaned_rule["action"].upper()
        cleaned.append(cleaned_rule)

    return cleaned, errors


async def get_policy(chain_type: str = "ethereum") -> dict:
    """Get current wallet policy status."""
    return await wallet_request("GET", f"/agent/policy?chain_type={chain_type}")

"""
Transfer & signing tools — EVM and Solana.
"""

from .common import is_fly_machine, wallet_request, logger


async def evm_transfer(
    to: str, amount: str, chain_id: int = 1,
    data: str = "", gas_limit: str = "", gas_price: str = "",
    max_fee_per_gas: str = "", max_priority_fee_per_gas: str = "",
    nonce: str = "", tx_type: int = None,
) -> dict:
    """Sign and BROADCAST an EVM transaction. Gas is sponsored."""
    body = {"to": to, "amount": amount, "chain_id": chain_id}
    if data:
        body["data"] = data
    if gas_limit:
        body["gas_limit"] = gas_limit
    if gas_price:
        body["gas_price"] = gas_price
    if max_fee_per_gas:
        body["max_fee_per_gas"] = max_fee_per_gas
    if max_priority_fee_per_gas:
        body["max_priority_fee_per_gas"] = max_priority_fee_per_gas
    if nonce:
        body["nonce"] = nonce
    if tx_type is not None:
        body["tx_type"] = tx_type
    return await wallet_request("POST", "/agent/transfer", body)


async def evm_sign_transaction(
    to: str, amount: str, chain_id: int = 1,
    data: str = "", gas_limit: str = "", gas_price: str = "",
    max_fee_per_gas: str = "", max_priority_fee_per_gas: str = "",
    nonce: str = "", tx_type: int = None,
) -> dict:
    """Sign an EVM transaction WITHOUT broadcasting."""
    body = {"to": to, "amount": amount, "chain_id": chain_id}
    if data:
        body["data"] = data
    if gas_limit:
        body["gas_limit"] = gas_limit
    if gas_price:
        body["gas_price"] = gas_price
    if max_fee_per_gas:
        body["max_fee_per_gas"] = max_fee_per_gas
    if max_priority_fee_per_gas:
        body["max_priority_fee_per_gas"] = max_priority_fee_per_gas
    if nonce:
        body["nonce"] = nonce
    if tx_type is not None:
        body["tx_type"] = tx_type
    return await wallet_request("POST", "/agent/sign-transaction", body)


async def evm_sign_message(message: str) -> dict:
    """EIP-191 personal_sign."""
    return await wallet_request("POST", "/agent/sign", {"message": message})


async def evm_sign_typed_data(
    domain: dict, types: dict, primaryType: str, message: dict,
) -> dict:
    """Sign EIP-712 structured data."""
    return await wallet_request("POST", "/agent/sign-typed-data", {
        "domain": domain, "types": types,
        "primaryType": primaryType, "message": message,
    })


async def evm_transactions(chain: str = "ethereum", asset: str = "eth", limit: int = 20) -> dict:
    """Get EVM transaction history."""
    qs = f"?chain_type=ethereum&chain={chain}&asset={asset}&limit={limit}"
    return await wallet_request("GET", f"/agent/transactions{qs}")


# ── Solana ───────────────────────────────────────────────────────────────────

async def sol_transfer(transaction: str, caip2: str = "solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp") -> dict:
    """Sign and BROADCAST a Solana transaction."""
    return await wallet_request("POST", "/agent/sol/transfer", {
        "transaction": transaction, "caip2": caip2,
    })


async def sol_sign_transaction(transaction: str) -> dict:
    """Sign a Solana transaction WITHOUT broadcasting."""
    return await wallet_request("POST", "/agent/sol/sign-transaction", {
        "transaction": transaction,
    })


async def sol_sign_message(message: str) -> dict:
    """Sign a message with Solana wallet (base64)."""
    return await wallet_request("POST", "/agent/sol/sign", {"message": message})


async def sol_transactions(chain: str = "solana", asset: str = "sol", limit: int = 20) -> dict:
    """Get Solana transaction history."""
    qs = f"?chain_type=solana&chain={chain}&asset={asset}&limit={limit}"
    return await wallet_request("GET", f"/agent/transactions{qs}")

"""
Wallet Tool Wrappers — BaseTool classes for Agent framework.
EVM + Solana: balances, transfers, signing, policy.
"""

import asyncio
import logging
import os
import time
from core.tool import BaseTool, ToolContext, ToolResult

from .tools.common import is_fly_machine
from .tools.info import get_wallet_info
from .tools.balance import get_evm_balance, get_sol_balance, get_all_balances
from .tools.transfer import (
    evm_transfer, evm_sign_transaction, evm_sign_message, evm_sign_typed_data,
    evm_transactions, sol_transfer, sol_sign_transaction, sol_sign_message,
    sol_transactions,
)
from .tools.policy import get_policy, validate_and_clean_rules

logger = logging.getLogger(__name__)

EVM_CHAINS = ["ethereum", "base", "arbitrum", "optimism", "polygon", "linea"]

def _fly_check():
    if not is_fly_machine():
        return ToolResult(success=False, error="Not running on a Fly Machine — wallet unavailable")
    return None


# ── Info ─────────────────────────────────────────────────────────────────────

class WalletInfoTool(BaseTool):
    @property
    def name(self): return "wallet_info"
    @property
    def description(self): return "Get all on-chain wallet addresses for this agent (one per chain)."
    @property
    def parameters(self): return {"type": "object", "properties": {}}

    async def execute(self, ctx: ToolContext, **kw) -> ToolResult:
        if err := _fly_check(): return err
        try:
            return ToolResult(success=True, output=await get_wallet_info())
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── EVM Balance ──────────────────────────────────────────────────────────────

class WalletBalanceTool(BaseTool):
    @property
    def name(self): return "wallet_balance"
    @property
    def description(self): return """Get EVM wallet balance on a specific chain. Omit 'asset' to discover ALL tokens.
Chains: ethereum, base, arbitrum, optimism, polygon, linea.
Use wallet_get_all_balances for all chains at once."""
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "chain": {"type": "string", "enum": EVM_CHAINS, "description": "Required. Blockchain network."},
            "address": {"type": "string", "description": "EVM address (0x...). Omit for own wallet."},
            "asset": {"type": "string", "description": "Asset filter. Omit for ALL tokens."},
        },
        "required": ["chain"],
    }

    async def execute(self, ctx: ToolContext, chain="", address="", asset="", **kw) -> ToolResult:
        if not chain or chain not in EVM_CHAINS:
            return ToolResult(success=False, error=f"'chain' required. One of: {', '.join(EVM_CHAINS)}")
        try:
            data = await get_evm_balance(chain, address, asset)
            if "error" in data:
                return ToolResult(success=False, error=data["error"])
            return ToolResult(success=True, output=data)
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── Solana Balance ───────────────────────────────────────────────────────────

class WalletSolBalanceTool(BaseTool):
    @property
    def name(self): return "wallet_sol_balance"
    @property
    def description(self): return "Get Solana wallet balance. Omit 'asset' to discover ALL SPL tokens."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "address": {"type": "string", "description": "Solana address. Omit for own wallet."},
            "asset": {"type": "string", "description": "Asset filter. Omit for ALL tokens."},
        },
    }

    async def execute(self, ctx: ToolContext, address="", asset="", **kw) -> ToolResult:
        try:
            data = await get_sol_balance(address, asset)
            if "error" in data:
                return ToolResult(success=False, error=data["error"])
            return ToolResult(success=True, output=data)
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── All Balances ─────────────────────────────────────────────────────────────

class WalletGetAllBalancesTool(BaseTool):
    @property
    def name(self): return "wallet_get_all_balances"
    @property
    def description(self): return "Get complete balance snapshot across ALL chains (EVM + Solana) with USD values."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "evm_address": {"type": "string", "description": "EVM address (0x...). Omit for own."},
            "sol_address": {"type": "string", "description": "Solana address. Omit for own."},
        },
    }

    async def execute(self, ctx: ToolContext, evm_address="", sol_address="", **kw) -> ToolResult:
        try:
            data = await get_all_balances(evm_address, sol_address)
            if "error" in data:
                return ToolResult(success=False, error=data["error"])
            return ToolResult(success=True, output=data)
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── EVM Transfer ─────────────────────────────────────────────────────────────

class WalletTransferTool(BaseTool):
    @property
    def name(self): return "wallet_transfer"
    @property
    def description(self): return """Sign and BROADCAST an EVM transaction. Gas is sponsored.
Use '0' amount for contract calls. Policy-gated if enabled."""
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "to": {"type": "string", "description": "Target address (0x...)"},
            "amount": {"type": "string", "description": "Amount in wei"},
            "chain_id": {"type": "integer", "description": "Chain ID (default: 1)"},
            "data": {"type": "string", "description": "Hex calldata for contract calls"},
            "gas_limit": {"type": "string"}, "gas_price": {"type": "string"},
            "max_fee_per_gas": {"type": "string"}, "max_priority_fee_per_gas": {"type": "string"},
            "nonce": {"type": "string"}, "tx_type": {"type": "integer", "description": "0=legacy, 2=EIP-1559"},
        },
        "required": ["to", "amount"],
    }

    async def execute(self, ctx: ToolContext, to="", amount="", chain_id=1,
                      data="", gas_limit="", gas_price="",
                      max_fee_per_gas="", max_priority_fee_per_gas="",
                      nonce="", tx_type=None, **kw) -> ToolResult:
        if err := _fly_check(): return err
        if not to or not amount:
            return ToolResult(success=False, error="'to' and 'amount' required")
        try:
            resp = await evm_transfer(to, amount, chain_id, data, gas_limit, gas_price,
                                      max_fee_per_gas, max_priority_fee_per_gas, nonce, tx_type)
            return ToolResult(success=True, output=resp)
        except Exception as e:
            msg = str(e)
            if "policy" in msg.lower():
                return ToolResult(success=False, error=f"Policy violation: {msg}")
            return ToolResult(success=False, error=msg)


# ── EVM Sign Transaction ────────────────────────────────────────────────────

class WalletSignTransactionTool(BaseTool):
    @property
    def name(self): return "wallet_sign_transaction"
    @property
    def description(self): return "Sign an EVM transaction WITHOUT broadcasting. Returns signed tx data."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "to": {"type": "string"}, "amount": {"type": "string"},
            "chain_id": {"type": "integer"}, "data": {"type": "string"},
            "gas_limit": {"type": "string"}, "gas_price": {"type": "string"},
            "max_fee_per_gas": {"type": "string"}, "max_priority_fee_per_gas": {"type": "string"},
            "nonce": {"type": "string"}, "tx_type": {"type": "integer"},
        },
        "required": ["to", "amount"],
    }

    async def execute(self, ctx: ToolContext, to="", amount="", chain_id=1,
                      data="", gas_limit="", gas_price="",
                      max_fee_per_gas="", max_priority_fee_per_gas="",
                      nonce="", tx_type=None, **kw) -> ToolResult:
        if err := _fly_check(): return err
        if not to or not amount:
            return ToolResult(success=False, error="'to' and 'amount' required")
        try:
            resp = await evm_sign_transaction(to, amount, chain_id, data, gas_limit, gas_price,
                                              max_fee_per_gas, max_priority_fee_per_gas, nonce, tx_type)
            return ToolResult(success=True, output=resp)
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── EVM Sign Message ─────────────────────────────────────────────────────────

class WalletSignTool(BaseTool):
    @property
    def name(self): return "wallet_sign"
    @property
    def description(self): return "Sign a message (EIP-191 personal_sign). Proves wallet ownership."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {"message": {"type": "string", "description": "Message to sign"}},
        "required": ["message"],
    }

    async def execute(self, ctx: ToolContext, message="", **kw) -> ToolResult:
        if err := _fly_check(): return err
        if not message: return ToolResult(success=False, error="'message' required")
        try:
            return ToolResult(success=True, output=await evm_sign_message(message))
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── EVM Sign Typed Data ──────────────────────────────────────────────────────

class WalletSignTypedDataTool(BaseTool):
    @property
    def name(self): return "wallet_sign_typed_data"
    @property
    def description(self): return "Sign EIP-712 structured data (permits, orders, etc.)."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "domain": {"type": "object", "description": "EIP-712 domain separator"},
            "types": {"type": "object", "description": "Type definitions"},
            "primaryType": {"type": "string", "description": "Primary type name"},
            "message": {"type": "object", "description": "Data to sign"},
        },
        "required": ["domain", "types", "primaryType", "message"],
    }

    async def execute(self, ctx: ToolContext, domain=None, types=None,
                      primaryType="", message=None, **kw) -> ToolResult:
        if err := _fly_check(): return err
        if not all([domain, types, primaryType, message]):
            return ToolResult(success=False, error="All params required")
        try:
            return ToolResult(success=True, output=await evm_sign_typed_data(domain, types, primaryType, message))
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── EVM Transactions ─────────────────────────────────────────────────────────

class WalletTransactionsTool(BaseTool):
    @property
    def name(self): return "wallet_transactions"
    @property
    def description(self): return "Get recent EVM transaction history."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "chain": {"type": "string"}, "asset": {"type": "string"},
            "limit": {"type": "integer", "description": "Max 100"},
        },
    }

    async def execute(self, ctx: ToolContext, chain="ethereum", asset="eth", limit=20, **kw) -> ToolResult:
        if err := _fly_check(): return err
        try:
            return ToolResult(success=True, output=await evm_transactions(chain, asset, limit))
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── Solana Transfer ──────────────────────────────────────────────────────────

class WalletSolTransferTool(BaseTool):
    @property
    def name(self): return "wallet_sol_transfer"
    @property
    def description(self): return "Sign and BROADCAST a Solana transaction. Gas sponsored. Policy-gated if enabled."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "transaction": {"type": "string", "description": "Base64-encoded Solana tx"},
            "caip2": {"type": "string", "description": "CAIP-2 chain ID (default: mainnet)"},
        },
        "required": ["transaction"],
    }

    async def execute(self, ctx: ToolContext, transaction="", caip2="solana:5eykt4UsFv8P8NJdTREpY1vzqKqZKvdp", **kw) -> ToolResult:
        if err := _fly_check(): return err
        if not transaction: return ToolResult(success=False, error="'transaction' required")
        try:
            return ToolResult(success=True, output=await sol_transfer(transaction, caip2))
        except Exception as e:
            msg = str(e)
            if "policy" in msg.lower():
                return ToolResult(success=False, error=f"Policy violation: {msg}")
            return ToolResult(success=False, error=msg)


# ── Solana Sign Transaction ─────────────────────────────────────────────────

class WalletSolSignTransactionTool(BaseTool):
    @property
    def name(self): return "wallet_sol_sign_transaction"
    @property
    def description(self): return "Sign a Solana transaction WITHOUT broadcasting."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {"transaction": {"type": "string", "description": "Base64-encoded Solana tx"}},
        "required": ["transaction"],
    }

    async def execute(self, ctx: ToolContext, transaction="", **kw) -> ToolResult:
        if err := _fly_check(): return err
        if not transaction: return ToolResult(success=False, error="'transaction' required")
        try:
            return ToolResult(success=True, output=await sol_sign_transaction(transaction))
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── Solana Sign Message ──────────────────────────────────────────────────────

class WalletSolSignTool(BaseTool):
    @property
    def name(self): return "wallet_sol_sign"
    @property
    def description(self): return "Sign a message with Solana wallet (base64)."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {"message": {"type": "string", "description": "Base64-encoded message"}},
        "required": ["message"],
    }

    async def execute(self, ctx: ToolContext, message="", **kw) -> ToolResult:
        if err := _fly_check(): return err
        if not message: return ToolResult(success=False, error="'message' required")
        try:
            return ToolResult(success=True, output=await sol_sign_message(message))
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── Solana Transactions ──────────────────────────────────────────────────────

class WalletSolTransactionsTool(BaseTool):
    @property
    def name(self): return "wallet_sol_transactions"
    @property
    def description(self): return "Get recent Solana transaction history."
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "chain": {"type": "string"}, "asset": {"type": "string"},
            "limit": {"type": "integer"},
        },
    }

    async def execute(self, ctx: ToolContext, chain="solana", asset="sol", limit=20, **kw) -> ToolResult:
        if err := _fly_check(): return err
        try:
            return ToolResult(success=True, output=await sol_transactions(chain, asset, limit))
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── Get Policy ───────────────────────────────────────────────────────────────

class WalletGetPolicyTool(BaseTool):
    @property
    def name(self): return "wallet_get_policy"
    @property
    def description(self): return """Get wallet policy status.
- enabled=false → allow-all (default)
- enabled=true, rules=[] → deny-all
- enabled=true, rules=[...] → rules enforced"""
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "chain_type": {"type": "string", "enum": ["ethereum", "solana"], "default": "ethereum"},
        },
    }

    async def execute(self, ctx: ToolContext, chain_type="ethereum", **kw) -> ToolResult:
        if err := _fly_check(): return err
        try:
            return ToolResult(success=True, output=await get_policy(chain_type))
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# ── Propose Policy ───────────────────────────────────────────────────────────

class WalletProposePolicyTool(BaseTool):
    @property
    def name(self): return "wallet_propose_policy"
    @property
    def description(self): return """Propose a wallet policy update. Sends action_request to frontend for user confirmation.
For both EVM and Solana, call TWICE (once per chain_type)."""
    @property
    def parameters(self): return {
        "type": "object",
        "properties": {
            "chain_type": {"type": "string", "enum": ["ethereum", "solana"]},
            "rules": {"type": "array", "description": "Privy policy rule objects", "items": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"}, "method": {"type": "string"},
                    "conditions": {"type": "array", "items": {"type": "object"}},
                    "action": {"type": "string", "enum": ["ALLOW", "DENY"]},
                },
            }},
            "title": {"type": "string", "description": "Short title (shown in UI)"},
            "description": {"type": "string", "description": "What this policy does"},
        },
        "required": ["chain_type", "rules", "title", "description"],
    }

    async def execute(self, ctx: ToolContext, chain_type="", rules=None,
                      title="", description="", **kw) -> ToolResult:
        if not chain_type or chain_type not in ("ethereum", "solana"):
            return ToolResult(success=False, error="chain_type must be 'ethereum' or 'solana'")
        if rules is None:
            return ToolResult(success=False, error="'rules' required")
        if not title:
            return ToolResult(success=False, error="'title' required")

        cleaned_rules, validation_errors = validate_and_clean_rules(rules, chain_type)
        if validation_errors:
            return ToolResult(
                success=False,
                error="Rule validation failed:\n" + "\n".join(f"- {e}" for e in validation_errors),
            )

        # Truncate names to Privy limit
        for rule in cleaned_rules:
            if isinstance(rule, dict) and "name" in rule and len(rule["name"]) > 50:
                rule["name"] = rule["name"][:50]

        container_id = os.environ.get("FLY_MACHINE_ID", "") or os.environ.get("FLY_ALLOC_ID", "") or "local-dev"
        action_id = f"act_{int(time.time())}_{os.urandom(4).hex()}"

        payload = {
            "container_id": container_id,
            "chain_type": chain_type,
            "rules": cleaned_rules,
        }

        streaming = getattr(ctx, "streaming", None)
        if streaming:
            streaming.action_request(
                action_id=action_id,
                action="update_wallet_policy",
                title=title,
                description=description,
                payload=payload,
                require_signature=True,
            )
            return ToolResult(success=True, output={
                "status": "action_request_sent",
                "action_id": action_id,
                "message": f"Policy proposal sent. Chain: {chain_type}, Rules: {len(cleaned_rules)}.",
            })
        else:
            return ToolResult(
                success=False,
                error="Streaming context not available — cannot send action_request.",
            )