wallet_manager.py

"""
錢包管理模組
處理錢包私鑰載入、餘額查詢等功能
"""
from typing import List, Optional

import base58
from solana.rpc.async_api import AsyncClient
from solders.keypair import Keypair
from solders.pubkey import Pubkey
from spl.token.async_client import AsyncToken
from spl.token.constants import TOKEN_PROGRAM_ID

from logger import get_logger


class WalletError(Exception):
    """錢包相關錯誤"""

    pass


class Wallet:
    """錢包類"""

    def __init__(self, keypair: Keypair):
        """
        初始化錢包

        Args:
            keypair: Solana 密鑰對
        """
        self.keypair = keypair
        self.public_key = keypair.pubkey()
        self.address = str(self.public_key)

    @classmethod
    def from_private_key(cls, private_key: str) -> "Wallet":
        """
        從私鑰創建錢包

        Args:
            private_key: Base58 編碼的私鑰

        Returns:
            Wallet 實例

        Raises:
            WalletError: 私鑰格式錯誤時拋出
        """
        try:
            # 嘗試解碼 Base58 私鑰
            private_key_bytes = base58.b58decode(private_key)

            # 創建密鑰對
            keypair = Keypair.from_bytes(private_key_bytes)

            return cls(keypair)
        except Exception as e:
            raise WalletError(f"無效的私鑰格式: {e}")

    async def get_sol_balance(self, client: AsyncClient) -> float:
        """
        獲取 SOL 餘額

        Args:
            client: Solana RPC 客戶端

        Returns:
            SOL 餘額（單位：SOL）
        """
        try:
            response = await client.get_balance(self.public_key)
            if response.value is None:
                return 0.0
            # 轉換從 lamports 到 SOL (1 SOL = 10^9 lamports)
            return response.value / 1_000_000_000
        except Exception as e:
            get_logger().error(f"獲取 SOL 餘額失敗: {e}")
            return 0.0

    async def get_token_balance(
        self, client: AsyncClient, token_mint: str
    ) -> float:
        """
        獲取 SPL Token 餘額

        Args:
            client: Solana RPC 客戶端
            token_mint: Token mint 地址

        Returns:
            Token 餘額
        """
        try:
            from spl.token.instructions import get_associated_token_address

            mint_pubkey = Pubkey.from_string(token_mint)

            # 獲取關聯代幣帳戶地址 (Associated Token Account)
            ata = get_associated_token_address(self.public_key, mint_pubkey)

            # 獲取帳戶餘額
            response = await client.get_token_account_balance(ata)

            if response.value:
                return float(response.value.ui_amount or 0)

            return 0.0
        except Exception as e:
            # 如果帳戶不存在或其他錯誤，返回 0
            get_logger().debug(f"獲取 Token 餘額失敗（可能帳戶不存在）: {e}")
            return 0.0

    async def get_usdc_balance(self, client: AsyncClient, usdc_mint: str) -> float:
        """
        獲取 USDC 餘額

        Args:
            client: Solana RPC 客戶端
            usdc_mint: USDC mint 地址

        Returns:
            USDC 餘額
        """
        return await self.get_token_balance(client, usdc_mint)

    def __repr__(self) -> str:
        """字符串表示"""
        return f"Wallet({self.address[:4]}...{self.address[-4:]})"


class WalletManager:
    """錢包管理器"""

    def __init__(self, private_keys: List[str]):
        """
        初始化錢包管理器

        Args:
            private_keys: 私鑰列表（Base58 格式）

        Raises:
            WalletError: 私鑰格式錯誤時拋出
        """
        self.wallets: List[Wallet] = []
        self.logger = get_logger()

        # 載入所有錢包
        for i, private_key in enumerate(private_keys, 1):
            try:
                wallet = Wallet.from_private_key(private_key)
                self.wallets.append(wallet)
                self.logger.info(
                    f"✅ 載入錢包 {i}/{len(private_keys)}: "
                    f"{wallet.address[:4]}...{wallet.address[-4:]}"
                )
            except WalletError as e:
                self.logger.error(f"❌ 載入錢包 {i} 失敗: {e}")
                raise

        if not self.wallets:
            raise WalletError("沒有成功載入任何錢包")

        self.logger.info(f"✅ 成功載入 {len(self.wallets)} 個錢包")

    def get_wallets(self) -> List[Wallet]:
        """獲取所有錢包"""
        return self.wallets

    def get_wallet(self, index: int) -> Optional[Wallet]:
        """
        根據索引獲取錢包

        Args:
            index: 錢包索引（從 0 開始）

        Returns:
            Wallet 實例，如果索引無效則返回 None
        """
        if 0 <= index < len(self.wallets):
            return self.wallets[index]
        return None

    def __len__(self) -> int:
        """返回錢包數量"""
        return len(self.wallets)

    def __iter__(self):
        """使 WalletManager 可迭代"""
        return iter(self.wallets)