client.py

import logging
from typing import Dict, List, Optional

import requests

from models import SwapQuote

logger = logging.getLogger(__name__)

WELL_KNOWN_ASSETS = {
    "XOR": "0x0200000000000000000000000000000000000000000000000000000000000000",
    "VAL": "0x0200040000000000000000000000000000000000000000000000000000000000",
    "PSWAP": "0x0200050000000000000000000000000000000000000000000000000000000000",
    "DAI": "0x0200060000000000000000000000000000000000000000000000000000000000",
    "ETH": "0x0200070000000000000000000000000000000000000000000000000000000000",
    "XSTUSD": "0x0200080000000000000000000000000000000000000000000000000000000000",
}


class SoraClient:
    """Client for the SORA v2 decentralized economy platform API.

    Provides methods for querying account balances, obtaining swap
    quotes through the Polkaswap DEX, listing liquidity pools,
    and fetching network-level information.
    """

    def __init__(self, config: dict) -> None:
        self.api_url = config.get("api_url", "").rstrip("/")
        self.dex_id = config.get("default_dex_id", 0)
        self.timeout = config.get("timeout", 30)
        self._session = requests.Session()
        self._session.headers.update({
            "Accept": "application/json",
            "Content-Type": "application/json",
        })

    def _resolve_asset_id(self, asset: str) -> str:
        return WELL_KNOWN_ASSETS.get(asset.upper(), asset)

    def _get(self, path: str, params: dict = None) -> Optional[dict]:
        url = f"{self.api_url}{path}"
        try:
            resp = self._session.get(url, params=params, timeout=self.timeout)
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            logger.error("GET %s failed: %s", path, exc)
            return None

    def _post(self, path: str, payload: dict = None) -> Optional[dict]:
        url = f"{self.api_url}{path}"
        try:
            resp = self._session.post(url, json=payload or {}, timeout=self.timeout)
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            logger.error("POST %s failed: %s", path, exc)
            return None

    def get_balances(self, address: str) -> List[Dict[str, str]]:
        data = self._get(f"/accounts/{address}/balances")
        if data is None:
            return []

        results = []
        for item in data if isinstance(data, list) else data.get("balances", []):
            asset_id = item.get("assetId", "")
            symbol = item.get("symbol", "")
            name = item.get("name", symbol)
            balance = item.get("balance", item.get("free", "0"))

            if not symbol:
                for sym, aid in WELL_KNOWN_ASSETS.items():
                    if aid == asset_id:
                        symbol = sym
                        break

            results.append({
                "asset_id": asset_id,
                "symbol": symbol or asset_id[:10],
                "name": name or symbol,
                "balance": str(balance),
            })

        return results

    def get_swap_quote(
        self,
        input_asset: str,
        output_asset: str,
        amount: str,
        slippage: float = 0.5,
    ) -> Optional[SwapQuote]:
        input_id = self._resolve_asset_id(input_asset)
        output_id = self._resolve_asset_id(output_asset)

        data = self._get("/dex/quote", params={
            "dexId": self.dex_id,
            "inputAssetId": input_id,
            "outputAssetId": output_id,
            "amount": amount,
            "swapVariant": "WithDesiredInput",
            "selectedSourceTypes": "XYKPool,MulticollateralBondingCurvePool",
        })

        if data is None:
            return None

        return SwapQuote(
            input_asset=input_asset.upper(),
            output_asset=output_asset.upper(),
            amount_in=amount,
            amount_out=data.get("amount", "0"),
            price_impact=data.get("priceImpact", "0"),
            fee=data.get("fee", "0"),
            route=data.get("route", [input_asset.upper(), output_asset.upper()]),
        )

    def execute_swap(
        self,
        input_asset: str,
        output_asset: str,
        amount: str,
        slippage: float = 0.5,
    ) -> str:
        input_id = self._resolve_asset_id(input_asset)
        output_id = self._resolve_asset_id(output_asset)

        result = self._post("/dex/swap", payload={
            "dexId": self.dex_id,
            "inputAssetId": input_id,
            "outputAssetId": output_id,
            "amount": amount,
            "slippageTolerance": slippage,
            "swapVariant": "WithDesiredInput",
            "selectedSourceTypes": ["XYKPool", "MulticollateralBondingCurvePool"],
        })

        if result is None:
            raise RuntimeError("Swap request failed")

        return result.get("txHash", result.get("hash", str(result)))

    def get_pools(self) -> List[Dict[str, str]]:
        data = self._get("/dex/pools", params={"dexId": self.dex_id})
        if data is None:
            return []

        results = []
        for pool in data if isinstance(data, list) else data.get("pools", []):
            base_sym = pool.get("baseAssetSymbol", pool.get("baseAssetId", "?")[:10])
            target_sym = pool.get("targetAssetSymbol", pool.get("targetAssetId", "?")[:10])

            results.append({
                "pair": f"{base_sym}/{target_sym}",
                "base_reserve": str(pool.get("baseAssetReserves", "0")),
                "target_reserve": str(pool.get("targetAssetReserves", "0")),
                "total_liquidity": str(pool.get("totalLiquidity", "0")),
            })

        return results

    def get_network_info(self) -> Dict[str, str]:
        data = self._get("/network/info")
        if data is None:
            return {"status": "unavailable"}

        return {
            "Network": data.get("name", "SORA"),
            "Chain": data.get("chain", "N/A"),
            "Node Version": data.get("nodeVersion", "N/A"),
            "Spec Version": str(data.get("specVersion", "N/A")),
            "Block Height": str(data.get("blockHeight", "N/A")),
            "Finalized Block": str(data.get("finalizedBlock", "N/A")),
            "Peer Count": str(data.get("peerCount", "N/A")),
            "Total Issuance (XOR)": data.get("totalIssuance", "N/A"),
        }