from __future__ import annotations

import httpx
import pandas as pd
import polars as pl

from ._http import load_parquet, list_snapshots, delete_snapshot
from .erc20 import ERC20Namespace
from .binance import BinanceNamespace
from .protocols import (
    AaveNamespace,
    UniswapNamespace,
    LidoNamespace,
    StaderNamespace,
    ThresholdNamespace,
    NativeNamespace,
    CacheNamespace,
)


class DataProviderClient:
    def __init__(self, url: str):
        self._url = url.rstrip("/")
        self._session = httpx.AsyncClient(timeout=3600)
        self.erc20 = ERC20Namespace(self._session, self._url)
        self.binance = BinanceNamespace(self._session, self._url)
        self.aave = AaveNamespace(self._session, self._url)
        self.uniswap = UniswapNamespace(self._session, self._url)
        self.lido = LidoNamespace(self._session, self._url)
        self.stader = StaderNamespace(self._session, self._url)
        self.threshold = ThresholdNamespace(self._session, self._url)
        self.native = NativeNamespace(self._session, self._url)
        self.cache = CacheNamespace(self._session, self._url)

    async def health(self) -> bool:
        response = await self._session.get(self._url + "/health")
        response.raise_for_status()
        return True

    async def load_parquet(self, key: str) -> pd.DataFrame:
        """Load a saved snapshot as a pandas DataFrame."""
        table = await load_parquet(self._session, self._url, key)
        return table.to_pandas()

    async def load_parquet_polars(self, key: str) -> pl.DataFrame:
        """Load a saved snapshot as a polars DataFrame."""
        table = await load_parquet(self._session, self._url, key)
        return pl.from_arrow(table)

    async def list_snapshots(self) -> list[str]:
        """List all saved snapshot keys."""
        return await list_snapshots(self._session, self._url)

    async def delete_snapshot(self, key: str) -> None:
        """Delete a saved snapshot."""
        await delete_snapshot(self._session, self._url, key)

    async def close(self) -> None:
        await self._session.aclose()

    async def __aenter__(self) -> DataProviderClient:
        return self

    async def __aexit__(self, *_) -> None:
        await self.close()
