from __future__ import annotations

from datetime import datetime, timezone
from typing import TYPE_CHECKING

import httpx
import pandas as pd
import polars as pl
import pyarrow as pa

from ._http import fetch_table

if TYPE_CHECKING:
    from typing import Self

_TIME_COL = "timestamp"


def _to_timestamp(date: datetime | str | int) -> str:
    """Normalise any DateType value to 'YYYY-MM-DDTHH:MM:SSZ' for the server."""
    if isinstance(date, int):
        dt = datetime.fromtimestamp(date / 1000, tz=timezone.utc)
    elif isinstance(date, str):
        if "T" in date:
            dt = datetime.fromisoformat(date.replace("Z", "+00:00"))
        elif len(date) == 10:
            dt = datetime.strptime(date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
        else:
            dt = datetime.strptime(date, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc)
    elif isinstance(date, datetime):
        dt = date
    else:
        raise ValueError(f"Unsupported date type: {type(date)}")
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt.strftime("%Y-%m-%dT%H:%M:%SZ")


class BaseQuery:
    def __init__(self, session: httpx.AsyncClient, base_url: str, body: dict):
        self._session = session
        self._base_url = base_url
        self._body = body

    def network(self, n: str | list[str]) -> Self:
        if isinstance(n, list):
            self._body["networks"] = n
        else:
            self._body["network"] = n
        return self

    def with_network(self) -> Self:
        self._body["with_network"] = True
        return self

    def time_range(self, since: datetime | str | int, until: datetime | str | int) -> Self:
        self._body["since"] = _to_timestamp(since)
        self._body["until"] = _to_timestamp(until)
        return self

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

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

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

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

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

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

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

    def with_value(self) -> Self:
        self._body["with_value"] = True
        return self

    def verbose(self) -> Self:
        self._body["verbose"] = True
        return self

    def aggregate(self, group_by: str = "time", period: str = "1h") -> Self:
        self._body["aggregate"] = True
        self._body["group_by"] = group_by
        self._body["period"] = period
        return self

    async def as_pandas(self) -> pd.DataFrame:
        table = await self._fetch_table()
        df = table.to_pandas()
        if _TIME_COL in df.columns:
            df = df.sort_values(_TIME_COL, ignore_index=True)
        return df

    async def as_polars(self) -> pl.DataFrame:
        table = await self._fetch_table()
        df = pl.from_arrow(table)
        if _TIME_COL in df.columns:
            df = df.sort(_TIME_COL)
        return df

    async def as_parquet(self, key: str) -> None:
        """Save the query result as a named parquet snapshot on the server."""
        self._body["save_key"] = key
        try:
            await self._fetch_table()
        finally:
            del self._body["save_key"]


class CacheableQuery(BaseQuery):
    def cache(self, cache_type: str = "append") -> Self:
        self._body["cache"] = True
        self._body["cache_type"] = cache_type
        return self

    def parallel(self) -> Self:
        self._body["parallel"] = True
        return self


class EventQuery(CacheableQuery):
    def __init__(self, session: httpx.AsyncClient, base_url: str, path: str, body: dict):
        super().__init__(session, base_url, body)
        self._path = path

    def _resolve_path(self) -> str:
        if self._body.get("aggregate"):
            return self._path.rsplit("/read", 1)[0] + "/aggregate"
        return self._path

    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:
        import asyncio

        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)
