from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

import httpx
import pyarrow as pa

from ._http import fetch_table
from ._query import CacheableQuery

if TYPE_CHECKING:
    from typing import Self


class ERC20TransfersQuery(CacheableQuery):
    def __init__(self, session: httpx.AsyncClient, base_url: str, tokens: list[str]):
        super().__init__(session, base_url, {"tokens": tokens})
        self._min_amount = None

    def min_amount(self, amount: float) -> Self:
        self._min_amount = amount
        return self

    def max_amount(self, amount: float) -> Self:
        self._body["max_amount"] = amount
        return self

    def sender(self, address: str) -> Self:
        self._body["sender"] = address
        return self

    def receiver(self, address: str) -> Self:
        self._body["receiver"] = address
        return self

    def sender_label(self, label: str) -> Self:
        self._body["sender_label"] = label
        return self

    def sender_category(self, category: str) -> Self:
        self._body["sender_category"] = category
        return self

    def receiver_label(self, label: str) -> Self:
        self._body["receiver_label"] = label
        return self

    def receiver_category(self, category: str) -> Self:
        self._body["receiver_category"] = category
        return self

    def exclude_sender(self, address: str) -> Self:
        self._body["exclude_sender"] = address
        return self

    def exclude_sender_label(self, label: str) -> Self:
        self._body["exclude_sender_label"] = label
        return self

    def exclude_sender_category(self, category: str) -> Self:
        self._body["exclude_sender_category"] = category
        return self

    def exclude_receiver(self, address: str) -> Self:
        self._body["exclude_receiver"] = address
        return self

    def exclude_receiver_label(self, label: str) -> Self:
        self._body["exclude_receiver_label"] = label
        return self

    def exclude_receiver_category(self, category: str) -> Self:
        self._body["exclude_receiver_category"] = category
        return self

    def _resolve_path(self) -> str:
        if self._body.get("aggregate"):
            return "/erc20_transfers/aggregate"
        elif self._min_amount is not None:
            self._body["min_amount"] = self._min_amount
            return "/erc20_transfers/read/min"
        return "/erc20_transfers/read"

    async def _fetch_single(self, network: str) -> pa.Table:
        body = {**self._body, "network": network}
        body.pop("networks", None)
        path = self._resolve_path()
        return await fetch_table(self._session, self._base_url + path, body)

    async def _fetch_table(self) -> pa.Table:
        networks = self._body.get("networks")
        if networks:
            tables = await asyncio.gather(*[self._fetch_single(n) for n in networks])
            non_empty = [t for t in tables if t is not None and len(t) > 0]
            if not non_empty:
                return tables[0] if tables else pa.table({})
            return pa.concat_tables(non_empty)
        path = self._resolve_path()
        return await fetch_table(self._session, self._base_url + path, self._body)


class TRC20TransfersQuery(CacheableQuery):
    def __init__(self, session: httpx.AsyncClient, base_url: str, tokens: list[str]):
        super().__init__(session, base_url, {"tokens": tokens})
        self._min_amount = None

    def min_amount(self, amount: float) -> Self:
        self._min_amount = amount
        return self

    def max_amount(self, amount: float) -> Self:
        self._body["max_amount"] = amount
        return self

    def sender(self, address: str) -> Self:
        self._body["sender"] = address
        return self

    def receiver(self, address: str) -> Self:
        self._body["receiver"] = address
        return self

    def sender_label(self, label: str) -> Self:
        self._body["sender_label"] = label
        return self

    def sender_category(self, category: str) -> Self:
        self._body["sender_category"] = category
        return self

    def receiver_label(self, label: str) -> Self:
        self._body["receiver_label"] = label
        return self

    def receiver_category(self, category: str) -> Self:
        self._body["receiver_category"] = category
        return self

    def exclude_sender(self, address: str) -> Self:
        self._body["exclude_sender"] = address
        return self

    def exclude_sender_label(self, label: str) -> Self:
        self._body["exclude_sender_label"] = label
        return self

    def exclude_sender_category(self, category: str) -> Self:
        self._body["exclude_sender_category"] = category
        return self

    def exclude_receiver(self, address: str) -> Self:
        self._body["exclude_receiver"] = address
        return self

    def exclude_receiver_label(self, label: str) -> Self:
        self._body["exclude_receiver_label"] = label
        return self

    def exclude_receiver_category(self, category: str) -> Self:
        self._body["exclude_receiver_category"] = category
        return self

    def _resolve_path(self) -> str:
        if self._body.get("aggregate"):
            return "/trc20_transfers/aggregate"
        elif self._min_amount is not None:
            self._body["min_amount"] = self._min_amount
            return "/trc20_transfers/read/min"
        return "/trc20_transfers/read"

    async def _fetch_single(self, network: str) -> pa.Table:
        body = {**self._body, "network": network}
        body.pop("networks", None)
        path = self._resolve_path()
        return await fetch_table(self._session, self._base_url + path, body)

    async def _fetch_table(self) -> pa.Table:
        networks = self._body.get("networks")
        if networks:
            tables = await asyncio.gather(*[self._fetch_single(n) for n in networks])
            non_empty = [t for t in tables if t is not None and len(t) > 0]
            if not non_empty:
                return tables[0] if tables else pa.table({})
            return pa.concat_tables(non_empty)
        path = self._resolve_path()
        return await fetch_table(self._session, self._base_url + path, self._body)


class ERC20Namespace:
    def __init__(self, session: httpx.AsyncClient, base_url: str):
        self._session = session
        self._base_url = base_url

    def transfers(self, tokens: list[str]) -> ERC20TransfersQuery:
        return ERC20TransfersQuery(self._session, self._base_url, tokens)

    async def flush(self, network: str | None = None, token: str | None = None) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        await self._session.post(self._base_url + "/erc20_transfers/flush", json=body)

    async def flush_aggregate(
        self,
        network: str | None = None,
        token: str | None = None,
        group_by: str | None = None,
        period: str | None = None,
    ) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        if group_by is not None:
            body["group_by"] = group_by
        if period is not None:
            body["period"] = period
        await self._session.post(self._base_url + "/erc20_transfers/aggregate/flush", json=body)

    async def compact(self, network: str | None = None, token: str | None = None) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        await self._session.post(self._base_url + "/erc20_transfers/compact", json=body)

    async def compact_aggregate(
        self,
        network: str | None = None,
        token: str | None = None,
        group_by: str | None = None,
        period: str | None = None,
    ) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        if group_by is not None:
            body["group_by"] = group_by
        if period is not None:
            body["period"] = period
        await self._session.post(self._base_url + "/erc20_transfers/aggregate/compact", json=body)


class TRC20Namespace:
    def __init__(self, session: httpx.AsyncClient, base_url: str):
        self._session = session
        self._base_url = base_url

    def transfers(self, tokens: list[str]) -> TRC20TransfersQuery:
        return TRC20TransfersQuery(self._session, self._base_url, tokens)

    async def flush(self, network: str | None = None, token: str | None = None) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        await self._session.post(self._base_url + "/trc20_transfers/flush", json=body)

    async def flush_aggregate(
        self,
        network: str | None = None,
        token: str | None = None,
        group_by: str | None = None,
        period: str | None = None,
    ) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        if group_by is not None:
            body["group_by"] = group_by
        if period is not None:
            body["period"] = period
        await self._session.post(self._base_url + "/trc20_transfers/aggregate/flush", json=body)

    async def compact(self, network: str | None = None, token: str | None = None) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        await self._session.post(self._base_url + "/trc20_transfers/compact", json=body)

    async def compact_aggregate(
        self,
        network: str | None = None,
        token: str | None = None,
        group_by: str | None = None,
        period: str | None = None,
    ) -> None:
        body = {}
        if network is not None:
            body["network"] = network
        if token is not None:
            body["token"] = token
        if group_by is not None:
            body["group_by"] = group_by
        if period is not None:
            body["period"] = period
        await self._session.post(self._base_url + "/trc20_transfers/aggregate/compact", json=body)
