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):
    """
    Enum for specifying the resolution of timestamp values in MarketData.
    """

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


class MarketData:
    def __init__(self, ticks: pd.DataFrame):
        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:
        if self._candles_1m is None:
            self._candles_1m = self._compute_candles("1min")
        return self._candles_1m

    @property
    def candles_5m(self) -> pd.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:
        if self._candles_15m is None:
            self._candles_15m = self._compute_candles("15min")
        return self._candles_15m

    @property
    def candles_1h(self) -> pd.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:
        if self._candles_4h is None:
            self._candles_4h = self._compute_candles("4h")
        return self._candles_4h

    @property
    def candles_1d(self) -> pd.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:
        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:
        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:
        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:

        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:

        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:
        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: list[str | Path]) -> MarketData:
        global_df = pd.DataFrame()
        for file in 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)
