from __future__ import annotations

import enum
import logging
from pathlib import Path
from typing import Any, Literal, cast

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


class ValueTimestampResolution(enum.Enum):
    """Timestamp matching modes used by query helpers.

    Attributes
    ----------
    EXACT : ValueTimestampResolution
        Require an exact timestamp match.
    CLOSEST_BEFORE : ValueTimestampResolution
        Resolve to the nearest value at or before the timestamp.
    CLOSEST_AFTER : ValueTimestampResolution
        Resolve to the nearest value at or after the timestamp.
    """

    EXACT = "exact"
    CLOSEST_BEFORE = "closest_before"
    CLOSEST_AFTER = "closest_after"


class MarketData:
    """Container and query interface for tick and candle market data.

    Parameters
    ----------
    ticks : pd.DataFrame
        Tick dataframe indexed by timestamp, expected to include at least
        ``bid`` and ``ask`` columns.
    """

    def __init__(self, ticks: pd.DataFrame):
        """Initialize market data with lazy candle caches.

        Parameters
        ----------
        ticks : pd.DataFrame
            Tick dataframe indexed by timestamp.
        """
        self.ticks = ticks
        self._candles_1m: pd.DataFrame | None = None
        self._candles_5m: pd.DataFrame | None = None
        self._candles_15m: pd.DataFrame | None = None
        self._candles_1h: pd.DataFrame | None = None
        self._candles_4h: pd.DataFrame | None = None
        self._candles_1d: pd.DataFrame | None = None

    @property
    def candles_1m(self) -> pd.DataFrame:
        """Return cached 1-minute candles.

        Returns
        -------
        pd.DataFrame
            1-minute OHLC candle dataframe.
        """
        if self._candles_1m is None:
            self._candles_1m = self._compute_candles("1min")
        return self._candles_1m

    @property
    def candles_5m(self) -> pd.DataFrame:
        """Return cached 5-minute candles.

        Returns
        -------
        pd.DataFrame
            5-minute OHLC candle dataframe.
        """
        if self._candles_5m is None:
            self._candles_5m = self._compute_candles("5min")
        return self._candles_5m

    @property
    def candles_15m(self) -> pd.DataFrame:
        """Return cached 15-minute candles.

        Returns
        -------
        pd.DataFrame
            15-minute OHLC candle dataframe.
        """
        if self._candles_15m is None:
            self._candles_15m = self._compute_candles("15min")
        return self._candles_15m

    @property
    def candles_1h(self) -> pd.DataFrame:
        """Return cached 1-hour candles.

        Returns
        -------
        pd.DataFrame
            1-hour OHLC candle dataframe.
        """
        if self._candles_1h is None:
            self._candles_1h = self._compute_candles("1h")
        return self._candles_1h

    @property
    def candles_4h(self) -> pd.DataFrame:
        """Return cached 4-hour candles.

        Returns
        -------
        pd.DataFrame
            4-hour OHLC candle dataframe.
        """
        if self._candles_4h is None:
            self._candles_4h = self._compute_candles("4h")
        return self._candles_4h

    @property
    def candles_1d(self) -> pd.DataFrame:
        """Return cached 1-day candles.

        Returns
        -------
        pd.DataFrame
            1-day OHLC candle dataframe.
        """
        if self._candles_1d is None:
            self._candles_1d = self._compute_candles("1d")
        return self._candles_1d

    def filter(
        self, from_date: pd.Timestamp | None = None, to_date: pd.Timestamp | None = None
    ) -> MarketData:
        """Return a sliced :class:`MarketData` by timestamp boundaries.

        Parameters
        ----------
        from_date : pd.Timestamp | None, default=None
            Inclusive lower bound.
        to_date : pd.Timestamp | None, default=None
            Inclusive upper bound.

        Returns
        -------
        MarketData
            New market data instance with filtered ticks.
        """
        ticks = self.ticks
        if from_date is not None:
            ticks = ticks[ticks.index >= from_date]
        if to_date is not None:
            ticks = ticks[ticks.index <= to_date]
        # copy to avoid modifying original data
        ticks = ticks.copy()
        return MarketData(ticks)

    def get_candle(self, freq: str) -> pd.DataFrame:
        """Select a candle dataframe by frequency tag.

        Parameters
        ----------
        freq : str
            One of ``1m``, ``5m``, ``15m``, ``1h``, ``4h``, ``1d``.

        Returns
        -------
        pd.DataFrame
            Candle dataframe for the requested timeframe.

        Raises
        ------
        ValueError
            If ``freq`` is unknown.
        """
        if freq == "1m":
            return self.candles_1m
        elif freq == "5m":
            return self.candles_5m
        elif freq == "15m":
            return self.candles_15m
        elif freq == "1h":
            return self.candles_1h
        elif freq == "4h":
            return self.candles_4h
        elif freq == "1d":
            return self.candles_1d
        else:
            raise ValueError(f"Unknown frequency: {freq}")

    def query(
        self,
        data_type: str,
        from_date: pd.Timestamp | None = None,
        to_date: pd.Timestamp | None = None,
    ) -> pd.DataFrame:
        """Query ticks or candles optionally bounded by date range.

        Parameters
        ----------
        data_type : str
            ``ticks`` or a candle frequency tag supported by
            :meth:`get_candle`.
        from_date : pd.Timestamp | None, default=None
            Inclusive lower bound.
        to_date : pd.Timestamp | None, default=None
            Inclusive upper bound.

        Returns
        -------
        pd.DataFrame
            Filtered dataframe.
        """
        if data_type == "ticks":
            data = self.ticks
        else:
            data = self.get_candle(data_type)

        if from_date is not None:
            data = data[data.index >= from_date]
        if to_date is not None:
            data = data[data.index <= to_date]

        return data

    def query_value(
        self,
        data_type: str,
        col_name: str,
        at_date: pd.Timestamp,
        resolution: ValueTimestampResolution = ValueTimestampResolution.CLOSEST_BEFORE,
    ) -> Any:
        """Query a single scalar value at a timestamp.

        Parameters
        ----------
        data_type : str
            ``ticks`` or a candle frequency tag.
        col_name : str
            Column name to extract.
        at_date : pd.Timestamp
            Reference timestamp.
        resolution : ValueTimestampResolution, default=CLOSEST_BEFORE
            Timestamp resolution strategy.

        Returns
        -------
        Any
            Value from the selected column.

        Raises
        ------
        ValueError
            If no matching data or column exists.
        """

        iloc_idx = 0
        if resolution == ValueTimestampResolution.EXACT:
            data = self.query(data_type, from_date=at_date, to_date=at_date)
        elif resolution == ValueTimestampResolution.CLOSEST_BEFORE:
            data = self.query(data_type, to_date=at_date)
            iloc_idx = -1
        elif resolution == ValueTimestampResolution.CLOSEST_AFTER:
            data = self.query(data_type, from_date=at_date)
        else:
            raise ValueError(f"Unknown resolution: {resolution}")
        if data.empty:
            raise ValueError(
                f"No data found for type '{data_type}' at date {at_date} with resolution '{resolution}'."
            )
        if col_name not in data.columns:
            raise ValueError(
                f"Column '{col_name}' not found in data type '{data_type}'."
            )

        return data.iloc[iloc_idx][col_name]

    def query_values(
        self,
        data_type: str,
        col_name: str,
        timestamps: np.ndarray,
        resolution: ValueTimestampResolution = ValueTimestampResolution.CLOSEST_BEFORE,
        fill_nans: bool = True,
    ) -> np.ndarray:
        """Query multiple values aligned to many timestamps.

        Parameters
        ----------
        data_type : str
            ``ticks`` or a candle frequency tag.
        col_name : str
            Column name to extract.
        timestamps : np.ndarray
            Timestamp array to resolve.
        resolution : ValueTimestampResolution, default=CLOSEST_BEFORE
            Timestamp resolution strategy.
        fill_nans : bool, default=True
            If ``True``, unresolved values are backfilled with the opposite-side
            nearest value when possible.

        Returns
        -------
        np.ndarray
            Array of values aligned to ``timestamps``.

        Raises
        ------
        ValueError
            If ``resolution`` is invalid, timestamps are missing for EXACT
            resolution, or ``col_name`` is absent.
        """

        data = self.ticks if data_type == "ticks" else self.get_candle(data_type)

        if col_name not in data.columns:
            raise ValueError(
                f"Column '{col_name}' not found in data type '{data_type}'."
            )

        idx = data.index
        values = data[col_name].values
        ts = pd.DatetimeIndex(timestamps)

        if resolution == ValueTimestampResolution.EXACT:
            locs = idx.get_indexer(ts)
            if (locs == -1).any():
                raise ValueError("Some timestamps not found for EXACT resolution")
            return cast(np.ndarray, values[locs])

        out = np.full(len(ts), np.nan, dtype=np.float64)

        locs_before = idx.searchsorted(ts, side="right") - 1
        locs_after = idx.searchsorted(ts, side="left")

        if resolution == ValueTimestampResolution.CLOSEST_BEFORE:
            mask = locs_before >= 0
            out[mask] = cast(np.ndarray, values[locs_before[mask]])

            if fill_nans:
                mask_after = (locs_after < len(idx)) & ~mask
                out[mask_after] = cast(np.ndarray, values[locs_after[mask_after]])

        elif resolution == ValueTimestampResolution.CLOSEST_AFTER:
            mask = locs_after < len(idx)
            out[mask] = cast(np.ndarray, values[locs_after[mask]])

            if fill_nans:
                mask_before = (locs_before >= 0) & ~mask
                out[mask_before] = cast(np.ndarray, values[locs_before[mask_before]])

        else:
            raise ValueError(f"Unknown resolution: {resolution}")

        return out

    def _compute_candles(self, freq: str) -> pd.DataFrame:
        """
        Build OHLC candles from tick data.

        ticks columns:
            bid
            ask
            volume (optional)
            spread (optional)

        index:
            datetime64[ns]
        """

        ticks = self.ticks
        bid = ticks["bid"]
        ask = ticks["ask"]

        # --- OHLC ---
        bid_ohlc = bid.resample(freq, label="right", closed="left").ohlc()
        ask_ohlc = ask.resample(freq, label="right", closed="left").ohlc()

        bid_ohlc.rename(
            columns={
                "open": "bid_open",
                "high": "bid_high",
                "low": "bid_low",
                "close": "bid_close",
            },
            inplace=True,
        )
        ask_ohlc.rename(
            columns={
                "open": "ask_open",
                "high": "ask_high",
                "low": "ask_low",
                "close": "ask_close",
            },
            inplace=True,
        )

        candles = pd.concat([bid_ohlc, ask_ohlc], axis=1)

        if "volume" in ticks:
            candles["volume"] = ticks["volume"].resample(freq).sum()

        if "spread" in ticks:
            candles["spread"] = ticks["spread"].resample(freq).mean()

        candles = candles.dropna(subset=["bid_open", "ask_open"])

        return candles

    @staticmethod
    def _load_from_csv(
        file: str | Path,
        bid_col: str = "bid",
        ask_col: str = "ask",
        timestamp_col: str = "timestamp",
        ask_volume_col: str | None = None,
        bid_volume_col: str | None = None,
        volume_col: str | None = None,
        real_volume_col: str | None = None,
        epoch_unit: Literal["D", "s", "ms", "us", "ns"] | None = None,
    ) -> pd.DataFrame:
        """Load and normalize tick data from a CSV file.

        Parameters
        ----------
        file : str | Path
            Input CSV path.
        bid_col : str, default="bid"
            Bid column name in source file.
        ask_col : str, default="ask"
            Ask column name in source file.
        timestamp_col : str, default="timestamp"
            Timestamp column name in source file.
        ask_volume_col : str | None, default=None
            Optional ask volume column name.
        bid_volume_col : str | None, default=None
            Optional bid volume column name.
        volume_col : str | None, default=None
            Optional total volume column name.
        real_volume_col : str | None, default=None
            Optional real volume column name.
        epoch_unit : Literal["D", "s", "ms", "us", "ns"] | None, default=None
            Datetime epoch unit used when parsing timestamps.

        Returns
        -------
        pd.DataFrame
            Normalized tick dataframe indexed by timestamp.
        """
        df = pd.read_csv(file)
        df.rename(
            columns={
                bid_col: "bid",
                ask_col: "ask",
                **({ask_volume_col: "volume_ask"} if ask_volume_col else {}),
                **({bid_volume_col: "volume_bid"} if bid_volume_col else {}),
                **({volume_col: "volume"} if volume_col else {}),
                **({real_volume_col: "volume_real"} if real_volume_col else {}),
            },
            inplace=True,
        )
        if "volume" not in df.columns:
            if "bidVolume" in df.columns and "askVolume" in df.columns:
                df["volume"] = df["bidVolume"] + df["askVolume"]
            elif "volume_real" in df.columns:
                df["volume"] = df["volume_real"]
        df["timestamp"] = pd.to_datetime(df[timestamp_col], unit=epoch_unit)
        df.set_index("timestamp", inplace=True)
        if timestamp_col != "timestamp":
            df.drop(columns=[timestamp_col], inplace=True)
        df["spread"] = df["ask"] - df["bid"]
        return df

    @staticmethod
    def load_from_ticks(ticks: str | Path | list[str | Path]) -> MarketData:
        """Load and concatenate tick CSV files into one dataset.

        Parameters
        ----------
        ticks : str | Path | list[str | Path]
            One or more CSV file paths containing tick data.

        Returns
        -------
        MarketData
            Combined market data sorted by timestamp.
        """
        normalized_ticks = [ticks] if isinstance(ticks, (str, Path)) else ticks
        global_df = pd.DataFrame()
        for file in normalized_ticks:
            logger.log(logging.INFO, "Loading tick data from file: %s", file)
            df = MarketData._load_from_csv(file)
            global_df = pd.concat([global_df, df])

        global_df.sort_index(inplace=True)
        global_df.drop_duplicates(inplace=True)
        return MarketData(global_df)
