wallet.py

"""
wallet.py  –  Core async Solana wallet for the Agentic Wallet SDK
Handles: keypair management, SOL balance, SPL token accounts, transfers, history
"""
import asyncio, json, os
from typing import Optional, Dict, Any, List

from solders.keypair import Keypair
from solders.pubkey import Pubkey
from solders.system_program import TransferParams, transfer
from solders.transaction import Transaction
from solders.message import Message
from solana.rpc.async_api import AsyncClient
from solana.rpc.types import TxOpts
from solana.rpc.commitment import Confirmed

DEVNET_RPC  = "https://api.devnet.solana.com"
MAINNET_RPC = "https://api.mainnet-beta.solana.com"
TOKEN_PROGRAM_ID = Pubkey.from_string("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA")


class SolanaWallet:
    """Async Solana wallet. Safe for AI agent use – all writes are explicit."""

    def __init__(self, keypair=None, keypair_path=None,
                 rpc_url=DEVNET_RPC, network="devnet"):
        self.rpc_url = rpc_url
        self.network = network
        self._client: Optional[AsyncClient] = None
        if keypair:
            self.keypair = keypair
        elif keypair_path and os.path.exists(keypair_path):
            self.keypair = self._load_keypair(keypair_path)
        else:
            self.keypair = Keypair()
            print(f"[Wallet] Generated new keypair: {self.keypair.pubkey()}")
        self.pubkey = self.keypair.pubkey()

    # ── keypair helpers ──────────────────────────────────────────────────
    def _load_keypair(self, path: str) -> Keypair:
        with open(path) as f:
            return Keypair.from_bytes(bytes(json.load(f)))

    def save_keypair(self, path: str) -> None:
        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        with open(path, "w") as f:
            json.dump(list(bytes(self.keypair)), f)
        print(f"[Wallet] Keypair saved → {path}")

    # ── RPC client ───────────────────────────────────────────────────────
    async def _client_(self) -> AsyncClient:
        if not self._client:
            self._client = AsyncClient(self.rpc_url)
        return self._client

    async def close(self):
        if self._client:
            await self._client.close()
            self._client = None

    # ── balance / portfolio ──────────────────────────────────────────────
    async def get_sol_balance(self, pubkey: Optional[str] = None) -> float:
        """Returns SOL balance (not lamports)."""
        c = await self._client_()
        pk = Pubkey.from_string(pubkey) if pubkey else self.pubkey
        r = await c.get_balance(pk, commitment=Confirmed)
        return r.value / 1_000_000_000

    async def get_token_accounts(self, pubkey: Optional[str] = None) -> List[Dict]:
        """All SPL token accounts with parsed amounts."""
        c = await self._client_()
        pk = Pubkey.from_string(pubkey) if pubkey else self.pubkey
        try:
            from solana.rpc.types import TokenAccountOpts
            r = await c.get_token_accounts_by_owner_json_parsed(
                pk, TokenAccountOpts(program_id=TOKEN_PROGRAM_ID), commitment=Confirmed)
            out = []
            for acct in (r.value or []):
                info = acct.account.data.parsed["info"]
                ta   = info.get("tokenAmount", {})
                out.append({"pubkey": str(acct.pubkey), "mint": info.get("mint"),
                             "amount": ta.get("uiAmount", 0), "decimals": ta.get("decimals", 0)})
            return out
        except Exception as e:
            print(f"[Wallet] token_accounts error: {e}")
            return []

    async def get_portfolio(self) -> Dict[str, Any]:
        sol    = await self.get_sol_balance()
        tokens = await self.get_token_accounts()
        return {"pubkey": str(self.pubkey), "network": self.network,
                "sol_balance": sol, "tokens": tokens, "token_count": len(tokens)}

    # ── transfers ────────────────────────────────────────────────────────
    async def transfer_sol(self, recipient: str, amount_sol: float) -> Dict[str, Any]:
        """Send SOL. Raises ValueError on insufficient balance."""
        c = await self._client_()
        lamports = int(amount_sol * 1_000_000_000)
        bal = await self.get_sol_balance()
        if bal < amount_sol + 0.001:
            raise ValueError(f"Insufficient: have {bal:.6f} SOL, need {amount_sol+0.001:.6f}")
        ix  = transfer(TransferParams(from_pubkey=self.pubkey,
                                      to_pubkey=Pubkey.from_string(recipient),
                                      lamports=lamports))
        bh  = (await c.get_latest_blockhash(commitment=Confirmed)).value.blockhash
        msg = Message.new_with_blockhash([ix], self.pubkey, bh)
        tx  = Transaction([self.keypair], msg, bh)
        sig = str((await c.send_transaction(
            tx, opts=TxOpts(skip_preflight=False, preflight_commitment=Confirmed))).value)
        return {"success": True, "signature": sig,
                "from": str(self.pubkey), "to": recipient,
                "amount_sol": amount_sol, "lamports": lamports,
                "network": self.network,
                "explorer": f"https://explorer.solana.com/tx/{sig}?cluster={self.network}"}

    # ── devnet utils ─────────────────────────────────────────────────────
    async def request_airdrop(self, amount_sol: float = 1.0) -> str:
        if self.network != "devnet":
            raise ValueError("Airdrops only available on devnet")
        c   = await self._client_()
        sig = str((await c.request_airdrop(self.pubkey,
                                           int(amount_sol * 1_000_000_000))).value)
        await asyncio.sleep(2)
        return sig

    # ── history ──────────────────────────────────────────────────────────
    async def get_transaction_history(self, limit: int = 10) -> List[Dict]:
        c = await self._client_()
        r = await c.get_signatures_for_address(self.pubkey, limit=limit, commitment=Confirmed)
        return [{"signature": str(s.signature), "slot": s.slot,
                 "err": s.err, "block_time": s.block_time} for s in r.value]

    # ── helpers ──────────────────────────────────────────────────────────
    def get_public_key(self) -> str: return str(self.pubkey)
    def get_explorer_url(self) -> str:
        return f"https://explorer.solana.com/address/{self.pubkey}?cluster={self.network}"
    def __repr__(self): return f"SolanaWallet({self.pubkey}, {self.network})"