"""Parquet-based storage backend (optional — requires pyarrow)."""

from __future__ import annotations

from pathlib import Path
from typing import List, Optional

from sd_api.exceptions import StorageError, ImportError as _ImportError

try:
    import pandas as pd
    _HAS_PANDAS = True
except ImportError:
    _HAS_PANDAS = False

from sd_api.logging import get_logger
from sd_api.models.kline import KlineBar
from sd_api.models.quote import Quote
from sd_api.storage.base import BaseStore

logger = get_logger("storage.parquet")


class ParquetStore(BaseStore):
    """Columnar Parquet storage — efficient for analytics workloads.

    Requires pandas and pyarrow (install via ``pip install sd-api[storage]``).

    Usage::

        store = ParquetStore("./data/parquet")
        store.save_quotes(quotes)
        quotes = store.load_quotes("XAUUSD")
    """

    def __init__(self, base_dir: str = "./sd_api_data/parquet") -> None:
        if not _HAS_PANDAS:
            raise _ImportError(
                "pandas is required for ParquetStore. "
                "Install with: pip install sd-api[storage]"
            )
        self.base_dir = Path(base_dir)
        self.quotes_dir = self.base_dir / "quotes"
        self.kline_dir = self.base_dir / "kline"
        self.quotes_dir.mkdir(parents=True, exist_ok=True)
        self.kline_dir.mkdir(parents=True, exist_ok=True)

    def _quote_path(self, symbol: str) -> Path:
        return self.quotes_dir / f"{symbol.upper().strip()}.parquet"

    def _kline_path(self, symbol: str, tf: str) -> Path:
        sym_dir = self.kline_dir / symbol.upper().strip()
        sym_dir.mkdir(exist_ok=True)
        return sym_dir / f"{tf}.parquet"

    # ── quote I/O ──────────────────────────────────────────────────────────

    def save_quote(self, quote: Quote) -> None:
        self.save_quotes([quote])

    def save_quotes(self, quotes: List[Quote]) -> None:
        if not quotes:
            return
        by_symbol: dict = {}
        for q in quotes:
            by_symbol.setdefault(q.symbol, []).append(q)
        for symbol, qs in by_symbol.items():
            df = pd.DataFrame([q.to_dict() for q in qs])
            path = self._quote_path(symbol)
            mode = "ab" if path.exists() else "wb"
            existing = pd.read_parquet(path) if path.exists() else pd.DataFrame()
            combined = pd.concat([existing, df], ignore_index=True).drop_duplicates(
                subset=["symbol", "quote_ts_ms"], keep="last"
            )
            combined.to_parquet(path, index=False)
            logger.debug(f"Saved {len(qs)} quotes to Parquet {path}")

    def load_quotes(self, symbol: str, *, limit: int = 1000) -> List[Quote]:
        path = self._quote_path(symbol)
        if not path.exists():
            return []
        df = pd.read_parquet(path).tail(limit)
        return [Quote.from_api(row.to_dict()) for _, row in df.iterrows()]

    # ── kline I/O ─────────────────────────────────────────────────────────

    def save_kline(self, bars: List[KlineBar], symbol: str = "", tf: str = "") -> None:
        if not bars:
            return
        if not symbol or not tf:
            raise StorageError("symbol and tf are required for save_kline")
        path = self._kline_path(symbol, tf)
        df = pd.DataFrame([b.to_dict() for b in bars])
        df["symbol"] = symbol.upper()
        df["tf"] = tf
        existing = pd.read_parquet(path) if path.exists() else pd.DataFrame()
        combined = pd.concat([existing, df], ignore_index=True).drop_duplicates(
            subset=["symbol", "tf", "ts"], keep="last"
        )
        combined.to_parquet(path, index=False)
        logger.debug(f"Saved {len(bars)} bars to Parquet {path}")

    def load_kline(
        self,
        symbol: str,
        *,
        tf: str = "5min",
        limit: int = 500,
    ) -> List[KlineBar]:
        path = self._kline_path(symbol, tf)
        if not path.exists():
            return []
        df = pd.read_parquet(path).tail(limit)
        return [KlineBar.from_api(row.to_dict()) for _, row in df.iterrows()]
