import asyncio
import concurrent.futures
from asyncio.exceptions import CancelledError
from collections import defaultdict
from typing import TYPE_CHECKING, Awaitable, Callable

import ccxt.pro as cxp
import numpy as np
import pandas as pd
from ccxt import ExchangeClosedByUser, ExchangeError, ExchangeNotAvailable, NetworkError

from qubx import logger
from qubx.connectors.registry import account_processor
from qubx.core.account import BasicAccountProcessor
from qubx.core.basics import (
    CtrlChannel,
    Deal,
    Instrument,
    ITimeProvider,
    Order,
    Position,
    RestoredState,
    Timestamped,
    TransactionCostsCalculator,
    dt_64,
)
from qubx.core.helpers import extract_price
from qubx.core.interfaces import IHealthMonitor, ISubscriptionManager
from qubx.utils.marketdata.ccxt import ccxt_symbol_to_instrument
from qubx.utils.misc import AsyncThreadLoop

from .exceptions import CcxtSymbolNotRecognized
from .exchange_manager import ExchangeManager
from .utils import (
    ccxt_convert_balance,
    ccxt_convert_deal_info,
    ccxt_convert_order_info,
    ccxt_convert_positions,
    ccxt_convert_ticker,
    ccxt_extract_deals_from_exec,
    ccxt_find_instrument,
    ccxt_restore_position_from_deals,
    instrument_to_ccxt_symbol,
)

if TYPE_CHECKING:
    from qubx.utils.runner.accounts import AccountConfigurationManager


@account_processor("ccxt")
class CcxtAccountProcessor(BasicAccountProcessor):
    """
    Subscribes to account information from the exchange.
    """

    exchange_manager: ExchangeManager
    channel: CtrlChannel
    base_currency: str
    balance_interval: str
    position_interval: str
    subscription_interval: str
    max_position_restore_days: int
    max_retries: int

    _loop: AsyncThreadLoop
    _polling_tasks: dict[str, concurrent.futures.Future]
    _subscription_manager: ISubscriptionManager | None
    _polling_to_init: dict[str, bool]
    _required_instruments: set[Instrument]
    _latest_instruments: set[Instrument]

    _free_capital: float = np.nan
    _total_capital: float = np.nan
    _instrument_to_last_price: dict[Instrument, tuple[dt_64, float]]

    def __new__(cls, exchange_name: str, **kwargs):
        if cls is CcxtAccountProcessor:
            from .exchanges import CUSTOM_ACCOUNTS

            custom_cls = CUSTOM_ACCOUNTS.get(exchange_name.lower())
            if custom_cls is not None:
                return custom_cls.__new__(custom_cls, exchange_name=exchange_name, **kwargs)
        return super().__new__(cls)

    def __init__(
        self,
        exchange_name: str,
        channel: CtrlChannel,
        time_provider: ITimeProvider,
        account_manager: "AccountConfigurationManager",
        tcc: TransactionCostsCalculator,
        health_monitor: IHealthMonitor,
        data_provider=None,
        restored_state: RestoredState | None = None,
        read_only: bool = False,
        loop: asyncio.AbstractEventLoop | None = None,
        balance_interval: str = "30Sec",
        position_interval: str = "30Sec",
        subscription_interval: str = "10Sec",
        open_order_interval: str = "1Min",
        open_order_backoff: str = "1Min",
        max_position_restore_days: int = 5,
        max_retries: int = 10,
        connection_timeout: int = 30,
        **kwargs,
    ):
        from qubx.connectors.ccxt.factory import get_ccxt_exchange_manager

        creds = account_manager.get_exchange_credentials(exchange_name)

        exchange_manager = get_ccxt_exchange_manager(
            exchange=exchange_name,
            use_testnet=creds.testnet,
            api_key=creds.api_key,
            secret=creds.secret,
            health_monitor=health_monitor,
            time_provider=time_provider,
            loop=loop,
            **(creds.model_extra or {}),
        )

        super().__init__(
            account_id=exchange_name,
            time_provider=time_provider,
            base_currency=creds.base_currency,
            exchange=exchange_name,
            tcc=tcc,
            health_monitor=health_monitor,
            initial_capital=0,
            restored_state=restored_state,
        )
        self.exchange_manager = exchange_manager
        self.channel = channel
        self.max_retries = max_retries
        self.balance_interval = balance_interval
        self.position_interval = position_interval
        self.subscription_interval = subscription_interval
        self.open_order_interval = open_order_interval
        self.open_order_backoff = open_order_backoff
        self.max_position_restore_days = max_position_restore_days
        self._loop = AsyncThreadLoop(exchange_manager.exchange.asyncio_loop)
        self._is_running = False
        self._polling_tasks = {}
        self._polling_to_init = defaultdict(bool)
        self._instrument_to_last_price = {}
        self._required_instruments = set()
        self._latest_instruments = set()
        self._subscription_manager = None
        self._read_only = read_only
        self._connection_timeout = connection_timeout

    def set_subscription_manager(self, manager: ISubscriptionManager) -> None:
        self._subscription_manager = manager

    def start(self):
        """Start the balance and position polling tasks"""
        channel = self.channel
        if channel is None or not channel.control.is_set():
            return
        if self._subscription_manager is None:
            return
        if self._is_running:
            logger.debug("Account polling is already running")
            return

        self._is_running = True

        if not self.exchange_manager.exchange.isSandboxModeEnabled:
            # - start polling tasks
            self._loop.submit(self.exchange_manager.exchange.load_markets()).result()
            self._polling_tasks["balance"] = self._loop.submit(
                self._poller("balance", self._update_balance, self.balance_interval)
            )
            self._polling_tasks["position"] = self._loop.submit(
                self._poller("position", self._update_positions, self.position_interval)
            )

            # - start initialization tasks
            _init_tasks = [
                self._loop.submit(self._init_spot_positions()),  # restore spot positions
                self._loop.submit(self._init_open_orders()),  # fetch open orders
            ]

            logger.info("Waiting for account polling tasks to be initialized")
            _waiter = self._loop.submit(self._wait_for_init(*_init_tasks))
            _waiter.result()
            logger.info("Account polling tasks have been initialized")

        # - start subscription polling task
        self._polling_tasks["subscription"] = self._loop.submit(
            self._poller("subscription", self._update_subscriptions, self.subscription_interval)
        )
        # - subscribe to order executions
        self._polling_tasks["executions"] = self._loop.submit(self._subscribe_executions("executions", channel))
        # - sync open orders
        self._polling_tasks["open_orders"] = self._loop.submit(
            self._poller(
                "open_orders", self._sync_open_orders, self.open_order_interval, backoff=self.open_order_backoff
            )
        )

    def stop(self):
        """Stop all polling tasks"""
        for task in self._polling_tasks.values():
            if not task.done():
                task.cancel()
        self._polling_tasks.clear()
        self._is_running = False

    def update_position_price(self, time: dt_64, instrument: Instrument, update: float | Timestamped) -> None:
        self._instrument_to_last_price[instrument] = (time, extract_price(update))
        super().update_position_price(time, instrument, update)

    def get_total_capital(self, exchange: str | None = None) -> float:
        # If polling is not running yet, we need to fetch balance data directly
        if not self._is_running and self.exchange_manager.exchange:
            try:
                future = self._loop.submit(self._update_balance())
                future.result(timeout=self._connection_timeout)
            except Exception as e:
                logger.warning(f"Failed to fetch balance data before polling started: {e}")

        # sum of balances + market value of all positions on non spot/margin
        _currency_to_value = {c: self._get_currency_value(b.total, c) for c, b in self._balances.items()}
        _positions_value = sum([p.market_value_funds for p in self._positions.values() if p.instrument.is_futures()])
        return sum(_currency_to_value.values()) + _positions_value

    def _get_instrument_for_currency(self, currency: str) -> Instrument:
        symbol = f"{currency}/{self.base_currency}"
        market = self.exchange_manager.exchange.market(symbol)
        exchange_name = self.exchange_manager.exchange.name
        assert exchange_name is not None
        return ccxt_symbol_to_instrument(exchange_name, market)

    def _get_currency_value(self, amount: float, currency: str) -> float:
        if not amount:
            return 0.0
        if currency == self.base_currency:
            return amount
        instr = self._get_instrument_for_currency(currency)
        _dt, _price = self._instrument_to_last_price.get(instr, (None, None))
        if not _dt or not _price:
            # logger.warning(f"Price for {instr} not available. Using 0.")
            return 0.0
        return amount * _price

    async def _poller(
        self,
        name: str,
        coroutine: Callable[[], Awaitable],
        interval: str,
        backoff: str | None = None,
    ):
        sleep_time = pd.Timedelta(interval).total_seconds()
        retries = 0

        if backoff is not None:
            sleep_time = pd.Timedelta(backoff).total_seconds()
            await asyncio.sleep(sleep_time)

        while self.channel.control.is_set():
            try:
                await coroutine()

                if not self._polling_to_init[name]:
                    logger.info(f"{name} polling task has been initialized")
                    self._polling_to_init[name] = True

                retries = 0  # Reset retry counter on success
            except CancelledError:
                logger.info(f"{name} listening has been cancelled")
                break
            except ExchangeClosedByUser:
                logger.info(f"{name} listening has been stopped")
                break
            except (NetworkError, ExchangeError, ExchangeNotAvailable) as e:
                logger.error(f"Error polling account data: {e}")
                retries += 1
                if retries >= self.max_retries:
                    logger.error(f"Max retries ({self.max_retries}) reached. Stopping poller.")
                    break
            except Exception as e:
                if not self.channel.control.is_set():
                    # If the channel is closed, then ignore all exceptions and exit
                    break
                logger.error(f"Unexpected error during account polling: {e}")
                logger.exception(e)
                retries += 1
                if retries >= self.max_retries:
                    logger.error(f"Max retries ({self.max_retries}) reached. Stopping poller.")
                    break
            finally:
                if not self.channel.control.is_set():
                    break
                await asyncio.sleep(min(sleep_time * (2 ** (retries)), 60))  # Exponential backoff capped at 60s

        logger.debug(f"{name} polling task has been stopped")

    async def _wait(self, condition: Callable[[], bool], sleep: float = 0.1) -> None:
        while not condition():
            await asyncio.sleep(sleep)

    async def _wait_for_init(self, *futures: concurrent.futures.Future) -> None:
        await self._wait(lambda: all(self._polling_to_init.values()))
        await self._wait(lambda: all([f.done() for f in futures]))

    async def _update_subscriptions(self) -> None:
        """Subscribe to required instruments"""
        assert self._subscription_manager is not None
        await asyncio.sleep(pd.Timedelta(self.subscription_interval).total_seconds())

        # if required instruments have changed, subscribe to them
        if not self._latest_instruments.issuperset(self._required_instruments):
            await self._subscribe_instruments(list(self._required_instruments))
            self._latest_instruments.update(self._required_instruments)

    async def _update_balance(self) -> None:
        """Fetch and update balances from exchange"""
        balances_raw = await self.exchange_manager.exchange.fetch_balance()
        balances = ccxt_convert_balance(balances_raw, self.exchange)
        current_balances = self.get_balances()

        # remove balances that are not there anymore
        current_currencies = {b.currency for b in current_balances}
        new_currencies = {b.currency for b in balances}
        _removed_currencies = current_currencies - new_currencies
        for currency in _removed_currencies:
            self.update_balance(currency, 0, 0)

        # update current balances
        for balance in balances:
            self.update_balance(currency=balance.currency, total=balance.total, locked=balance.locked)

        # update required instruments that we need to subscribe to
        currencies = [b.currency for b in self.get_balances()]
        instruments = [
            self._get_instrument_for_currency(c) for c in currencies if c.upper() != self.base_currency.upper()
        ]
        self._required_instruments.update(instruments)

        # fetch tickers for instruments that don't have recent price updates
        await self._fetch_missing_tickers(instruments)

    async def _update_positions(self) -> None:
        # fetch and update positions from exchange
        ccxt_positions = await self.exchange_manager.exchange.fetch_positions()
        positions = ccxt_convert_positions(
            ccxt_positions, self.exchange_manager.exchange.name, self.exchange_manager.exchange.markets
        )  # type: ignore
        # update required instruments that we need to subscribe to
        self._required_instruments.update([p.instrument for p in positions])
        # update positions
        _instrument_to_position = {p.instrument: p for p in positions}
        _current_instruments = set(self._positions.keys())
        _new_instruments = set([p.instrument for p in positions])
        # - spot positions should not be updated here, because exchanges don't provide spot positions
        # - so we have to trust deal updates to update spot positions
        _to_remove = {instr for instr in _current_instruments - _new_instruments if instr.is_futures()}
        _to_add = _new_instruments - _current_instruments
        _to_modify = _current_instruments.intersection(_new_instruments)
        _update_positions = [Position(i) for i in _to_remove] + [_instrument_to_position[i] for i in _to_modify]
        # - add new positions
        for i in _to_add:
            self._positions[i] = _instrument_to_position[i]
        # - modify existing positions
        _time = self.time_provider.time()
        for pos in _update_positions:
            self._update_instrument_position(_time, self._positions[pos.instrument], pos)

    def _update_instrument_position(self, timestamp: dt_64, current_pos: Position, new_pos: Position) -> None:
        instrument = current_pos.instrument
        quantity_diff = new_pos.quantity - current_pos.quantity
        if abs(quantity_diff) >= instrument.lot_size:
            _current_price = current_pos.last_update_price
            current_pos.change_position_by(timestamp, quantity_diff, _current_price)

        # Update maintenance margin from exchange if available
        if new_pos._maint_margin_external:
            current_pos.set_external_maint_margin(new_pos.maint_margin)

    def _get_start_time_in_ms(self, days_before: int) -> int:
        return (self.time_provider.time() - days_before * pd.Timedelta("1d")).asm8.item() // 1000000  # type: ignore

    def _is_our_order(self, order: Order) -> bool:
        if order.client_id is None:
            return False
        return order.client_id.startswith("qubx_")

    def _is_base_currency(self, currency: str) -> bool:
        return currency.upper() == self.base_currency

    async def _subscribe_instruments(self, instruments: list[Instrument]) -> None:
        assert self._subscription_manager is not None

        # find missing subscriptions
        _base_sub = self._subscription_manager.get_base_subscription()
        _subscribed_instruments = self._subscription_manager.get_subscribed_instruments(_base_sub)
        _add_instruments = list(set(instruments) - set(_subscribed_instruments))

        if _add_instruments:
            # subscribe to instruments
            self._subscription_manager.subscribe(_base_sub, _add_instruments)
            self._subscription_manager.commit()

    async def _fetch_missing_tickers(self, instruments: list[Instrument]) -> None:
        _current_time = self.time_provider.time()
        _fetch_instruments: list[Instrument] = []
        for instr in instruments:
            _dt, _ = self._instrument_to_last_price.get(instr, (None, None))
            if _dt is None or pd.Timedelta(_current_time - _dt) > pd.Timedelta(self.balance_interval):  # type: ignore
                _fetch_instruments.append(instr)

        _symbol_to_instrument = {instr.symbol: instr for instr in instruments}
        if _fetch_instruments:
            logger.debug(f"Fetching missing tickers for {_fetch_instruments}")
            _fetch_symbols = [instrument_to_ccxt_symbol(instr) for instr in _fetch_instruments]
            tickers: dict[str, dict] = await self.exchange_manager.exchange.fetch_tickers(_fetch_symbols)
            for symbol, ticker in tickers.items():
                instr = _symbol_to_instrument.get(symbol)
                if instr is not None:
                    quote = ccxt_convert_ticker(ticker)
                    self.update_position_price(_current_time, instr, quote)

    async def _init_spot_positions(self) -> None:
        # - wait for balance to be initialized
        await self._wait(lambda: self._polling_to_init["balance"])
        logger.debug("Restoring spot positions ...")

        # - get nonzero balances
        _nonzero_balances = {
            c: b.total for c, b in self._balances.items() if b.total > 0 and not self._is_base_currency(c)
        }
        _positions = []

        async def _restore_pos(currency: str, balance: float) -> None:
            try:
                _instrument = self._get_instrument_for_currency(currency)
                # - get latest order for instrument and check client id
                _latest_orders = await self._fetch_orders(_instrument, limit=1)
                if not _latest_orders:
                    return
                _latest_order = list(_latest_orders.values())[-1]
                if self._is_our_order(_latest_order):
                    # - if it's our order, then we fetch the deals and restore position
                    _deals = await self._fetch_deals(_instrument, self.max_position_restore_days)
                    _position = ccxt_restore_position_from_deals(Position(_instrument), balance, _deals)
                    _positions.append(_position)
            except Exception as e:
                logger.warning(f"Error restoring position for {currency}: {e}")

        # - restore positions
        await asyncio.gather(*[_restore_pos(c, b) for c, b in _nonzero_balances.items()])

        # - attach positions
        if _positions:
            self.attach_positions(*_positions)
            logger.debug("Restored positions ->")
            for p in _positions:
                logger.debug(f"  ::  {p}")

    async def _init_open_orders(self) -> None:
        # wait for balances and positions to be initialized
        await self._wait(lambda: all([self._polling_to_init[task] for task in ["balance", "position"]]))
        await self._sync_open_orders(initial_call=True)

    async def _sync_open_orders(self, initial_call: bool = False) -> None:
        logger.debug("[SYNC] Fetching open orders ...")

        # in order to minimize order requests we only fetch open orders for instruments that we have positions in
        _nonzero_balances = {
            c: b.total for c, b in self._balances.items() if b.total > 0 and not self._is_base_currency(c)
        }
        _balance_instruments = [self._get_instrument_for_currency(c) for c in _nonzero_balances.keys()]
        _position_instruments = list(self._positions.keys())
        _instruments = list(set(_balance_instruments + _position_instruments))

        _open_orders: dict[str, Order] = {}

        async def _add_open_orders(instrument: Instrument) -> None:
            try:
                _orders = await self._fetch_orders(instrument, is_open=True)
                _open_orders.update(_orders)
            except Exception as e:
                logger.warning(f"[SYNC] Error fetching open orders for {instrument}: {e}")

        await asyncio.gather(*[_add_open_orders(i) for i in _instruments])

        if initial_call:
            # - when it's the initial call, we add the open orders to the account
            self.add_active_orders(_open_orders)
            logger.debug(f"[SYNC] Found {len(_open_orders)} open orders ->")
            _instr_to_open_orders: dict[Instrument, list[Order]] = defaultdict(list)
            for od in _open_orders.values():
                _instr_to_open_orders[od.instrument].append(od)
            for instr, orders in _instr_to_open_orders.items():
                logger.debug(f"  :: [SYNC] {instr} ->")
                for order in orders:
                    logger.debug(f"    :: [SYNC] {order.side} {order.quantity} @ {order.price} ({order.status})")
        else:
            # - we need to cancel the unexpected orders
            if not self._read_only:
                await self._cancel_unexpected_orders(_open_orders)

    async def _cancel_unexpected_orders(self, open_orders: dict[str, Order]) -> None:
        _expected_orders = set(self._active_orders.keys())
        _unexpected_orders = set(open_orders.keys()) - _expected_orders
        if _unexpected_orders:
            logger.info(f"[SYNC] Canceling {len(_unexpected_orders)} unexpected open orders ...")
            _instr_to_orders = defaultdict(list)
            for _id in _unexpected_orders:
                _order = open_orders[_id]
                _instr_to_orders[_order.instrument].append(_order)

            async def _cancel_order(order: Order) -> None:
                try:
                    await self.exchange_manager.exchange.cancel_order(
                        order.id, symbol=instrument_to_ccxt_symbol(order.instrument)
                    )
                    logger.debug(
                        f"  :: [SYNC] Canceled {order.id} {order.instrument.symbol} {order.side} {order.quantity} @ {order.price} ({order.status})"
                    )
                except Exception as e:
                    logger.warning(f"[SYNC] Error canceling order {order.id}: {e}")

            for instr, orders in _instr_to_orders.items():
                logger.debug(
                    f"[SYNC] Canceling {len(orders)} (out of {len(open_orders)}) unexpected open orders for {instr}"
                )
                await asyncio.gather(*[_cancel_order(order) for order in orders])

    async def _fetch_orders(
        self, instrument: Instrument, days_before: int = 30, limit: int | None = None, is_open: bool = False
    ) -> dict[str, Order]:
        _start_ms = self._get_start_time_in_ms(days_before) if limit is None else None
        _ccxt_symbol = instrument_to_ccxt_symbol(instrument)
        _fetcher = (
            self.exchange_manager.exchange.fetch_open_orders if is_open else self.exchange_manager.exchange.fetch_orders
        )
        _raw_orders = await _fetcher(_ccxt_symbol, since=_start_ms, limit=limit)
        _orders = [ccxt_convert_order_info(instrument, o) for o in _raw_orders]
        _id_to_order = {o.id: o for o in _orders}
        return dict(sorted(_id_to_order.items(), key=lambda x: x[1].time, reverse=False))

    async def _fetch_deals(self, instrument: Instrument, days_before: int = 30) -> list[Deal]:
        _start_ms = self._get_start_time_in_ms(days_before)
        _ccxt_symbol = instrument_to_ccxt_symbol(instrument)
        deals_data = await self.exchange_manager.exchange.fetch_my_trades(_ccxt_symbol, since=_start_ms)
        deals: list[Deal] = [ccxt_convert_deal_info(o) for o in deals_data]
        return sorted(deals, key=lambda x: x.time) if deals else []

    # TODO: this should take the exchange manager instead of cxp.Exchange
    async def _listen_to_stream(
        self,
        subscriber: Callable[[], Awaitable[None]],
        exchange: cxp.Exchange,
        channel: CtrlChannel,
        name: str,
    ):
        logger.info(f"Listening to {name}")
        n_retry = 0
        while channel.control.is_set():
            try:
                await subscriber()
                n_retry = 0
            except CcxtSymbolNotRecognized:
                continue
            except CancelledError:
                break
            except ExchangeClosedByUser:
                # - we closed connection so just stop it
                logger.info(f"{name} listening has been stopped")
                break
            except (NetworkError, ExchangeError, ExchangeNotAvailable) as e:
                logger.error(f"Error in {name} : {e}")
                await asyncio.sleep(1)
                continue
            except Exception as e:
                if not channel.control.is_set():
                    # If the channel is closed, then ignore all exceptions and exit
                    break
                logger.error(f"exception in {name} : {e}")
                logger.exception(e)
                n_retry += 1
                if n_retry >= self.max_retries:
                    logger.error(f"Max retries reached for {name}. Closing connection.")
                    del exchange
                    break
                await asyncio.sleep(min(2**n_retry, 60))  # Exponential backoff with a cap at 60 seconds

    async def _subscribe_executions(self, name: str, channel: CtrlChannel):
        _symbol_to_instrument = {}

        async def _watch_executions():
            exec = await self.exchange_manager.exchange.watch_orders()
            for report in exec:
                instrument = ccxt_find_instrument(
                    report["symbol"], self.exchange_manager.exchange, _symbol_to_instrument
                )
                order = ccxt_convert_order_info(instrument, report)
                deals = ccxt_extract_deals_from_exec(report)
                channel.send((instrument, "order", order, False))
                if deals:
                    channel.send((instrument, "deals", deals, False))

        await self._listen_to_stream(
            subscriber=_watch_executions,
            exchange=self.exchange_manager.exchange,
            channel=channel,
            name=name,
        )
