import math
from collections import defaultdict, deque
from typing import Any, Iterator, TypeAlias

import pandas as pd

from qubx import logger
from qubx.core.basics import DataType, Instrument, Timestamped
from qubx.core.exceptions import SimulationError
from qubx.data.readers import (
    AsDict,
    AsQuotes,
    AsTrades,
    DataReader,
    DataTransformer,
    RestoredBarsFromOHLC,
    RestoreQuotesFromOHLC,
    RestoreTradesFromOHLC,
)

SlicerOutData: TypeAlias = tuple[str, int, Timestamped] | tuple


class IteratedDataStreamsSlicer(Iterator[SlicerOutData]):
    """
    This class manages seamless iteration over multiple time-series data streams,
    ensuring that events are processed in the correct chronological order regardless of their source.
    It supports adding / removing new data streams to the slicer on the fly (during the itration).
    """

    _iterators: dict[str, Iterator[list[Timestamped]]]
    _buffers: dict[str, list[Timestamped]]
    _keys: deque[str]
    _iterating: bool

    def __init__(self):
        self._buffers = defaultdict(list)
        self._iterators = {}
        self._keys = deque()
        self._iterating = False

    def put(self, data: dict[str, Iterator[list[Timestamped]]]):
        _rebuild = False
        for k, vi in data.items():
            if k not in self._keys:
                self._iterators[k] = vi
                self._buffers[k] = self._load_next_chunk_to_buffer(k)  # do initial chunk fetching
                self._keys.append(k)
                _rebuild = True

        # - rebuild strategy
        if _rebuild and self._iterating:
            self._build_initial_iteration_seq()

    def __add__(self, data: dict[str, Iterator]) -> "IteratedDataStreamsSlicer":
        self.put(data)
        return self

    def remove(self, keys: list[str] | str):
        """
        Remove data iterator and associated keys from the queue.
        If the key is not found, it does nothing.
        """
        _keys = keys if isinstance(keys, list) else [keys]
        _rebuild = False
        for i in _keys:
            if i in self._buffers:
                self._buffers.pop(i)
                self._iterators.pop(i)
                self._keys.remove(i)
                _rebuild = True

        # - rebuild strategy
        if _rebuild and self._iterating:
            self._build_initial_iteration_seq()

    def __iter__(self) -> Iterator:
        self._build_initial_iteration_seq()
        self._iterating = True
        return self

    def _build_initial_iteration_seq(self):
        _init_seq = {k: self._buffers[k][-1].time for k in self._keys}
        _init_seq = dict(sorted(_init_seq.items(), key=lambda item: item[1]))
        self._keys = deque(_init_seq.keys())

    def _load_next_chunk_to_buffer(self, index: str) -> list[Timestamped]:
        return list(reversed(next(self._iterators[index])))

    def _remove_iterator(self, key: str):
        self._buffers.pop(key)
        self._iterators.pop(key)
        self._keys.remove(key)

    def _pop_top(self, k: str) -> Timestamped:
        """
        Removes and returns the most recent timestamped data element from the buffer associated with the given key.
        If the buffer is empty after popping, it attempts to load the next chunk of data into the buffer.
        If no more data is available, the iterator associated with the key is removed.

        Parameters:
            k (str): The key identifying the data stream buffer to pop from.

        Returns:
            Timestamped: The most recent timestamped data element from the buffer.
        """
        v = (data := self._buffers[k]).pop()
        if not data:
            try:
                data.extend(self._load_next_chunk_to_buffer(k))  # - get next chunk of data
            except StopIteration:
                self._remove_iterator(k)  # - remove iterable data
        return v

    def fetch_before_time(self, key: str, time_ns: int) -> list[Timestamped]:
        """
        Fetches and returns all timestamped data elements from the buffer associated with the given key
        that have a timestamp earlier than the specified time.

        Parameters:
            - key (str): The key identifying the data stream buffer to fetch from.
            - time_ns (int): The timestamp in nanoseconds. All returned elements will have a timestamp less than this value.

        Returns:
            - list[Timestamped]: A list of timestamped data elements that occur before the specified time.
        """
        values = []
        data = self._buffers[key]
        if not data:
            try:
                data.extend(self._load_next_chunk_to_buffer(key))  # - get next chunk of data
            except StopIteration:
                self._remove_iterator(key)

        # pull most past elements
        v = data[-1]
        while v.time < time_ns:
            values.append(data.pop())
            if not data:
                try:
                    data.extend(self._load_next_chunk_to_buffer(key))  # - get next chunk of data
                except StopIteration:
                    self._remove_iterator(key)
                    break
            v = data[-1]

        return values

    def __next__(self) -> SlicerOutData:
        """
        Advances the iterator to the next available timestamped data element across all data streams.

        Returns:
            - SlicerOutData: A tuple containing the key of the data stream, the timestamp of the data element, and the data element itself.

        Raises:
            - StopIteration: If there are no more data elements to iterate over.
        """
        if not self._keys:
            self._iterating = False
            raise StopIteration

        _min_t = math.inf
        _min_k = self._keys[0]
        for i in self._keys:
            _x = self._buffers[i][-1]
            if _x.time < _min_t:
                _min_t = _x.time
                _min_k = i

        _v = self._pop_top(_min_k)
        return (_min_k, _v.time, _v)


class DataFetcher:
    _fetcher_id: str
    _reader: DataReader
    _requested_data_type: str
    _producing_data_type: str
    _params: dict[str, object]
    _specs: list[str]

    _transformer: DataTransformer
    _timeframe: str | None = None
    _warmup_period: pd.Timedelta | None = None
    _chunksize: int = 5000

    def __init__(
        self,
        fetcher_id: str,
        reader: DataReader,
        subtype: str,
        params: dict[str, Any],
        warmup_period: pd.Timedelta | None = None,
        chunksize: int = 5000,
        open_close_time_indent_secs=1.0,  # open/close shift may depends on simulation
    ) -> None:
        self._fetcher_id = fetcher_id
        self._params = params
        self._reader = reader

        match subtype:
            case DataType.OHLC_QUOTES:
                # - requested restore quotes from OHLC
                self._transformer = RestoreQuotesFromOHLC(open_close_time_shift_secs=open_close_time_indent_secs)
                self._requested_data_type = "ohlc"
                self._producing_data_type = "quote"
                if "timeframe" in params:
                    self._timeframe = params.get("timeframe", "1Min")

            case DataType.OHLC_TRADES:
                # - requested restore trades from OHLC
                self._transformer = RestoreTradesFromOHLC(open_close_time_shift_secs=open_close_time_indent_secs)
                self._requested_data_type = "ohlc"
                self._producing_data_type = "trade"
                if "timeframe" in params:
                    self._timeframe = params.get("timeframe", "1Min")

            case DataType.OHLC:
                # - requested restore bars from OHLC
                self._transformer = RestoredBarsFromOHLC(open_close_time_shift_secs=open_close_time_indent_secs)
                self._requested_data_type = "ohlc"
                self._producing_data_type = "ohlc"
                if "timeframe" in params:
                    self._timeframe = params.get("timeframe", "1Min")

            case DataType.TRADE:
                self._requested_data_type = "trade"
                self._producing_data_type = "trade"
                self._transformer = AsTrades()

            case DataType.QUOTE:
                self._requested_data_type = "orderbook"
                self._producing_data_type = "quote"  # ???
                self._transformer = AsQuotes()

            case _:
                self._requested_data_type = subtype
                self._producing_data_type = subtype
                self._transformer = AsDict()

        self._warmup_period = warmup_period
        self._warmed = {}
        self._specs = []
        self._chunksize = chunksize

    @staticmethod
    def _make_request_id(instrument: Instrument) -> str:
        return f"{instrument.exchange}:{instrument.symbol}"

    def attach_instrument(self, instrument: Instrument) -> str:
        _data_id = self._make_request_id(instrument)

        if _data_id not in self._specs:
            self._specs.append(_data_id)
            self._warmed[_data_id] = False

        return self._fetcher_id + "." + _data_id

    def remove_instrument(self, instrument: Instrument) -> str:
        _data_id = self._make_request_id(instrument)

        if _data_id in self._specs:
            self._specs.remove(_data_id)
            del self._warmed[_data_id]

        return self._fetcher_id + "." + _data_id

    def has_instrument(self, instrument: Instrument) -> bool:
        return self._make_request_id(instrument) in self._specs

    def get_instruments_indices(self) -> list[str]:
        return [self._fetcher_id + "." + i for i in self._specs]

    def get_instrument_index(self, instrument: Instrument) -> str:
        return self._fetcher_id + "." + self._make_request_id(instrument)

    def load(
        self, start: str | pd.Timestamp, end: str | pd.Timestamp, to_load: list[Instrument] | None
    ) -> dict[str, Iterator]:
        """
        Loads data for specified instruments within a given time range.

        Parameters:
            - start (str | pd.Timestamp): The start time for data loading, can be a string or a pandas Timestamp.
            - end (str | pd.Timestamp): The end time for data loading, can be a string or a pandas Timestamp.
            - to_load (list[Instrument] | None): A list of instruments to load data for. If None, data for all subscribed instruments is loaded.

        Returns:
            - dict[str, Iterator]: A dictionary where keys are instrument identifiers and values are iterators over the loaded data.
        """
        _requests = self._specs if not to_load else set(self._make_request_id(i) for i in to_load)
        _r_iters = {}

        for _r in _requests:  # - TODO: replace this loop with multi-instrument request after DataReader refactoring
            if _r in self._specs:
                _start = pd.Timestamp(start)
                if self._warmup_period and not self._warmed.get(_r):
                    _start -= self._warmup_period
                    self._warmed[_r] = True

                _args = dict(
                    data_id=_r,
                    start=_start,
                    stop=end,
                    transform=self._transformer,
                    data_type=self._requested_data_type,
                    chunksize=self._chunksize,
                )

                if self._timeframe:
                    _args["timeframe"] = self._timeframe

                try:
                    _r_iters[self._fetcher_id + "." + _r] = self._reader.read(**_args)  # type: ignore
                except Exception as e:
                    logger.error(f">>> (DataFetcher::load) - failed to load <g>'{self._fetcher_id}'</g> data: {e}")
            else:
                raise IndexError(
                    f"Instrument {_r} is not subscribed for this data {self._requested_data_type} in {self._fetcher_id} !"
                )

        return _r_iters

    def __repr__(self) -> str:
        return f"{self._requested_data_type}({self._params}) (-{self._warmup_period if self._warmup_period else '--'}) [{','.join(self._specs)}] :-> {self._transformer.__class__.__name__}"


class IterableSimulationData(Iterator):
    """
    This class is a crucial component for backtesting system.
    It provides a flexible and efficient way to simulate market data feeds for strategy testing.

    Key Features:
        - Supports multiple data types (OHLC, trades, quotes) and instruments.
        - Allows for dynamic addition and removal of instruments during simulation.
        - Handles warmup periods for data preloading.
        - Manages historical and current data distinction during iteration.
        - Utilizes a data slicer (IteratedDataStreamsSlicer) to merge and order data from multiple sources.

    TODO:
        1. think how to provide initial "market quote" for each instrument
        2. optimization for historical data (return bunch of history instead of each bar in next(...))
    """

    _readers: dict[str, DataReader]
    _subtyped_fetchers: dict[str, DataFetcher]
    _warmups: dict[str, pd.Timedelta]
    _instruments: dict[str, tuple[Instrument, DataFetcher, DataType]]
    _open_close_time_indent_secs: int | float

    _slicer_ctrl: IteratedDataStreamsSlicer | None = None
    _slicing_iterator: Iterator | None = None
    _start: pd.Timestamp | None = None
    _stop: pd.Timestamp | None = None
    _current_time: int | None = None

    def __init__(
        self,
        readers: dict[str, DataReader],
        open_close_time_indent_secs=1,  # open/close ticks shift
    ):
        self._readers = dict(readers)
        self._instruments = {}
        self._subtyped_fetchers = {}
        self._warmups = {}
        self._open_close_time_indent_secs = open_close_time_indent_secs

    def set_typed_reader(self, data_type: str, reader: DataReader):
        self._readers[data_type] = reader
        if _fetcher := self._subtyped_fetchers.get(data_type):
            _fetcher._reader = reader

    def set_warmup_period(self, subscription: str, warmup_period: str | None = None):
        if warmup_period:
            _access_key, _, _ = self._parse_subscription_spec(subscription)
            self._warmups[_access_key] = pd.Timedelta(warmup_period)

    def _parse_subscription_spec(self, subscription: str) -> tuple[str, str, dict[str, object]]:
        _subtype, _params = DataType.from_str(subscription)
        match _subtype:
            case DataType.OHLC | DataType.OHLC_QUOTES:
                _timeframe = _params.get("timeframe", "1Min")
                _access_key = f"{_subtype}.{_timeframe}"
            case DataType.TRADE | DataType.QUOTE:
                _access_key = f"{_subtype}"
            case _:
                # - any arbitrary data type is passed as is
                _params = {}
                _subtype = subscription
                _access_key = f"{_subtype}"
        return _access_key, _subtype, _params

    def add_instruments_for_subscription(self, subscription: str, instruments: list[Instrument] | Instrument):
        instruments = instruments if isinstance(instruments, list) else [instruments]
        _subt_key, _data_type, _params = self._parse_subscription_spec(subscription)

        fetcher = self._subtyped_fetchers.get(_subt_key)
        if not fetcher:
            _reader = self._readers.get(_data_type)

            if _reader is None:
                raise SimulationError(f"No reader configured for data type: {_data_type}")

            self._subtyped_fetchers[_subt_key] = (
                fetcher := DataFetcher(
                    _subt_key,
                    _reader,
                    _data_type,
                    _params,
                    warmup_period=self._warmups.get(_subt_key),
                    open_close_time_indent_secs=self._open_close_time_indent_secs,
                )
            )

        _instrs_to_preload = []
        for i in instruments:
            if not fetcher.has_instrument(i):
                idx = fetcher.attach_instrument(i)
                self._instruments[idx] = (i, fetcher, subscription)  # type: ignore
                _instrs_to_preload.append(i)

        if self.is_running and _instrs_to_preload:
            self._slicer_ctrl += fetcher.load(
                pd.Timestamp(self._current_time, unit="ns"),  # type: ignore
                self._stop,  # type: ignore
                _instrs_to_preload,
            )

    def peek_historical_data(self, instrument: Instrument, subscription: str) -> list[Timestamped]:
        """
        Retrieves historical data for a specified instrument and subscription type up to the current simulation time.

        Parameters:
            - instrument (Instrument): instrument for which historical data is requested.
            - subscription (str): type of data subscription (e.g., OHLC, trades, quotes) for the instrument.

        Returns:
            - list[Timestamped]: A list of historical data elements for the specified instrument and subscription type
            that occurred before the current simulation time. If the simulation is not running, returns an empty list.

        Raises:
            SimulationError: If the instrument does not have the specified subscription in the simulation data provider.
        """
        if not self.has_subscription(instrument, subscription):
            raise SimulationError(
                f"Instrument: {instrument} has no subscription: {subscription} in this simulation data provider"
            )

        if not self.is_running:
            return []

        _subt_key, _, _ = self._parse_subscription_spec(subscription)
        _i_key = self._subtyped_fetchers[_subt_key].get_instrument_index(instrument)

        assert self._slicer_ctrl is not None and self._current_time is not None

        # fetch historical data for current time
        return self._slicer_ctrl.fetch_before_time(_i_key, self._current_time)

    def get_instruments_for_subscription(self, subscription: str) -> list[Instrument]:
        if subscription == DataType.ALL:
            return list((i for i, *_ in self._instruments.values()))

        _subt_key, _, _ = self._parse_subscription_spec(subscription)
        if (fetcher := self._subtyped_fetchers.get(_subt_key)) is not None:
            return [self._instruments[k][0] for k in fetcher.get_instruments_indices()]

        return []

    def get_subscriptions_for_instrument(self, instrument: Instrument | None) -> list[str]:
        r = []
        for i, f, s in self._instruments.values():
            if instrument is not None:
                if i == instrument:
                    r.append(s)
            else:
                r.append(s)
        return list(set(r))

    def has_subscription(self, instrument: Instrument, subscription_type: str) -> bool:
        for i, f, s in self._instruments.values():
            if i == instrument and s == subscription_type:
                return True
        return False

    def remove_instruments_from_subscription(self, subscription: str, instruments: list[Instrument] | Instrument):
        def _remove_from_fetcher(_subt_key: str, instruments: list[Instrument]):
            fetcher = self._subtyped_fetchers.get(_subt_key)
            if not fetcher:
                logger.warning(f"No configured data fetcher for '{_subt_key}' subscription !")
                return

            _keys_to_remove = []
            for i in instruments:
                # - try to remove from data fetcher
                if idx := fetcher.remove_instrument(i):
                    if idx in self._instruments:
                        self._instruments.pop(idx)
                        _keys_to_remove.append(idx)

            # print("REMOVING FROM:", _keys_to_remove)
            if self.is_running and _keys_to_remove:
                self._slicer_ctrl.remove(_keys_to_remove)  # type: ignore

        instruments = instruments if isinstance(instruments, list) else [instruments]

        # - if we want to remove instruments from all subscriptions
        if subscription == DataType.ALL:
            _f_keys = list(self._subtyped_fetchers.keys())
            for s in _f_keys:
                _remove_from_fetcher(s, instruments)
            return

        _subt_key, _, _ = self._parse_subscription_spec(subscription)
        _remove_from_fetcher(_subt_key, instruments)

    @property
    def is_running(self) -> bool:
        return self._current_time is not None

    def create_iterable(self, start: str | pd.Timestamp, stop: str | pd.Timestamp) -> Iterator:
        self._start = pd.Timestamp(start)
        self._stop = pd.Timestamp(stop)
        self._current_time = None
        self._slicer_ctrl = IteratedDataStreamsSlicer()
        return self

    def __iter__(self) -> Iterator:
        assert self._start is not None
        self._current_time = int(pd.Timestamp(self._start).timestamp() * 1e9)
        _ct_timestap = pd.Timestamp(self._current_time, unit="ns")

        for f in self._subtyped_fetchers.values():
            logger.debug(
                f"  [<c>IteratedDataStreamsSlicer</c>] :: Preloading initial data for {f._fetcher_id} {self._start} : {self._stop} ..."
            )
            self._slicer_ctrl += f.load(_ct_timestap, self._stop, None)  # type: ignore

        self._slicing_iterator = iter(self._slicer_ctrl)
        return self

    def __next__(self) -> tuple[Instrument, str, Timestamped, bool]:  # type: ignore
        try:
            while data := next(self._slicing_iterator):  # type: ignore
                k, t, v = data
                instr, fetcher, subt = self._instruments[k]
                data_type = fetcher._producing_data_type
                _is_historical = False
                if t < self._current_time:  # type: ignore
                    _is_historical = True
                else:
                    # only update the current time if the event is not historical
                    self._current_time = t

                return instr, data_type, v, _is_historical
        except StopIteration as e:  # noqa: F841
            raise StopIteration
