import io
from datetime import datetime
from typing import Optional, Union

import httpx
import pandas as pd
import polars as pl

from ._http import load_parquet, load_parquet_bytes, list_snapshots, delete_snapshot
from ._query import _to_timestamp


def _to_datetime(date: datetime | str | int) -> datetime:
    """Convert any date input to a timezone-aware datetime."""
    ts = _to_timestamp(date)  # returns 'YYYY-MM-DDTHH:MM:SSZ'
    return datetime.fromisoformat(ts.replace("Z", "+00:00"))

from .erc20 import ERC20Namespace, TRC20Namespace
from .binance import BinanceNamespace, HyperliquidNamespace
from .protocols import (
    AaveNamespace,
    UniswapNamespace,
    LidoNamespace,
    StaderNamespace,
    ThresholdNamespace,
    NativeNamespace,
    TronNativeNamespace,
    BitcoinNativeNamespace,
    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.trc20 = TRC20Namespace(self._session, self._url)
        self.tron_native = TronNativeNamespace(self._session, self._url)
        self.bitcoin_native = BitcoinNativeNamespace(self._session, self._url)
        self.hyperliquid = HyperliquidNamespace(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,
        since: Optional[Union[datetime, str, int]] = None,
        until: Optional[Union[datetime, str, int]] = None,
    ) -> pd.DataFrame:
        """Load a saved snapshot as a pandas DataFrame."""
        table = await load_parquet(self._session, self._url, key)
        df = table.to_pandas()
        if since is not None or until is not None:
            time_col = "timestamp" if "timestamp" in df.columns else "time"
            if time_col in df.columns:
                if since is not None:
                    df = df[df[time_col] >= pd.Timestamp(_to_datetime(since))]
                if until is not None:
                    df = df[df[time_col] <= pd.Timestamp(_to_datetime(until))]
                df = df.reset_index(drop=True)
        return df

    async def load_parquet_polars(
        self,
        key: str,
        since: Optional[Union[datetime, str, int]] = None,
        until: Optional[Union[datetime, str, int]] = None,
    ) -> pl.DataFrame:
        """Load a saved snapshot as a polars DataFrame."""
        raw = await load_parquet_bytes(self._session, self._url, key)
        df = pl.read_parquet(io.BytesIO(raw))
        if since is not None or until is not None:
            time_col = "timestamp" if "timestamp" in df.columns else "time"
            if time_col in df.columns:
                if since is not None:
                    df = df.filter(pl.col(time_col) >= _to_datetime(since))
                if until is not None:
                    df = df.filter(pl.col(time_col) <= _to_datetime(until))
        return df

    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()
