from pydmm_connectors.core.logger import MarketLogger
from pydmm_strategy.volume.config import VolumeMode
from pydmm_connectors.models import Ticker
from pydmm_connectors.models.connector_responses import Depth, Balance, BestBook
from pydmm_strategy.volume.market_factory import VolumeMarketFactory, MarketDataType, MarketSnapshot
from enum import auto, Enum
from dataclasses import dataclass

import numpy as np
import time
from typing import Tuple, Optional, List, Any, Callable
from concurrent.futures import Future
from concurrent.futures import ThreadPoolExecutor
from threading import Thread

class ExecutionStatus(Enum):
    SUCCESS = auto()
    FAILURE = auto()  # просто ошибка, не критичная
    STOPPED = auto()  # Бот нужно остановить
    _SUCCESS_REBALANCED = auto()  # (внутренний статус): Сделали ребаланс -> основной трейд в этой итерации не делаем

@dataclass
class ExecutorInput:
    price: float
    quantity: float
    target_price: float
    book: BestBook
    side: Optional[bool] = None  # True=buy, False=sell. Used by pam_cex path; None for spot self-trades.
    is_5m_open: bool = False     # True when this iter fires right after a 5m candle boundary (set by CandleCloser)


class VolumeExecutor:
    def __init__(
        self,
        market_factory: VolumeMarketFactory,
        logger: MarketLogger,
        ticker: Ticker,
        mode: VolumeMode,
        n_markets: int,
        force_rebalance: bool,
        bbo_spread_size: float,
        buy_ratio: float,
        tight_mode_enabled: bool,
        tight_min_base: float,
        tight_min_quote: float,
        min_quote_size: float = 0.0,
    ):
        self._logger = logger
        self._market_factory = market_factory
        self._thread_pool = ThreadPoolExecutor(max_workers=8) # TODO: отдельный тред пул для экзекутора
        self._mode = mode
        self._n_markets = n_markets
        self.ticker = ticker
        self._min_quote_size = min_quote_size if min_quote_size > 0 else ticker.min_notional
        self._step_size = 0.1 ** ticker.price_prec
        self._i1 = 0
        self._i2 = 0
        self._last_trade_price: Optional[float] = None
        self._bbo_spread_size = bbo_spread_size
        self._buy_ratio = buy_ratio
        self._tight_mode_enabled = tight_mode_enabled
        self._tight_min_base = tight_min_base
        self._tight_min_quote = tight_min_quote
        self._maker_taker_delay = 0.01
        # Anti-overflow self-match side selection (see _pick_self_trade_side): when the quoted
        # price sits in the outer band of the spread, choose which leg is the maker so the TAKER
        # leg points toward the FAR edge. Band is a fraction of the spread, floored to a tick
        # count so it stays meaningful on narrow spreads.
        self._edge_band_frac = 0.10
        self._edge_band_min_ticks = 1
        if force_rebalance:
            book = self._market_factory.get_market_data_and_log(MarketDataType.BEST_BOOK.value).best_book
            self._logger.event('force_rebalance_book', bid=book.bid_price, ask=book.ask_price)
            self._last_trade_price = (book.ask_price + book.bid_price) / 2

    def _generate_random_indexes(self) -> Tuple[int, int]:
        if self._mode == VolumeMode.SMART and self._n_markets > 1:
            i1 = np.random.randint(0, self._n_markets)
            i2 = np.random.randint(0, self._n_markets - 1)
            i2 += (i2 >= i1)
        else:
            i1 = 0
            i2 = 0

        return i1, i2

    def _check_last_price_and_log(self) -> bool:
        """Helper: проверяет last_price и логирует ошибку если None"""
        if self._last_trade_price is None:
            self._logger.error(
                'there will be no trades! '
                'turn on "Force rebalance" flag '
                '(be careful if the price has changed significantly)'
            )
            return True
        return False

    # TODO: не использовать book
    def _execute_bbo_trade(self, order_futures: List[Future[Any]], qty: float, book: BestBook, depth: Depth) -> float:
        ask_qty, bid_qty = 0, 0
        if len(depth.asks) > 0:
            ask_qty = depth.asks[0].quantity
        if len(depth.bids) > 0:
            bid_qty = depth.bids[0].quantity
        qty = max(
            min(qty, min(ask_qty, bid_qty)),
            self._min_quote_size / book.bid_price
        )
        self._logger.event(
            'bbo trade',
            qty=qty,
            amt_bid=qty * book.bid_price,
            amt_ask=qty * book.ask_price,
        )
        if book.bid_price <= 0 or (book.ask_price - book.bid_price) / book.bid_price >= self._bbo_spread_size:  # TODO: config
            self._logger.error(f"Can't do BBO trade: spread is too wide")
            return 0
        else:
            order_futures.append(self._thread_pool.submit(self._market_factory.new_limit, 0, book.bid_price, qty, False))
            order_futures.append(self._thread_pool.submit(self._market_factory.new_limit, 0, book.ask_price, qty, True))
            return qty * (book.bid_price + book.ask_price)

    def _execute_smart_trade(self, order_futures: List[Future[Any]], price: float, qty: float) -> float:
        if np.random.random() <= self._buy_ratio: # TODO: config
            order_futures.append(self._thread_pool.submit(self._market_factory.new_limit_maker, self._i1, price, qty, False))
            time.sleep(self._maker_taker_delay)
            order_futures.append(self._thread_pool.submit(self._market_factory.new_limit, self._i2, price, qty, True))
            return qty * price
        else:
            order_futures.append(self._thread_pool.submit(self._market_factory.new_limit_maker, self._i1, price, qty, True))
            time.sleep(self._maker_taker_delay)
            order_futures.append(self._thread_pool.submit(self._market_factory.new_limit, self._i2, price, qty, False))
            return qty * price

    def _execute_tight_trade(self, request: ExecutorInput, balance: Balance) -> float:
        book = request.book
        qty = request.quantity
        target_rounded = round(request.target_price, self.ticker.price_prec)
        self._logger.event(
            'execution in tight',
            bid=book.bid_price,
            ask=book.ask_price,
            target_rounded=target_rounded,
        )

        if book.bid_price <= target_rounded <= book.ask_price:
            if request.is_5m_open and (request.price == book.ask_price or request.price == book.bid_price):
                is_buy = (request.price == book.ask_price)
                self._logger.event('tight 5m opener', side='buy' if is_buy else 'sell')
            else:
                is_buy = np.random.randint(0, 1000) / 1000 < self._buy_ratio
        else:
            is_buy = (target_rounded >= book.ask_price)
        if is_buy:
            qty = max(qty, self._min_quote_size / book.ask_price)
            if balance.quote - qty * book.ask_price >= self._tight_min_quote:
                self._logger.event('tight buy ok', qty=qty)
                self._market_factory.new_limit(self._i1, book.ask_price, qty, True) # убрал поток
                self._last_trade_price = book.ask_price
                return book.ask_price * qty
            else:
                self._logger.event(
                    'tight buy skip',
                    quote_bal=balance.quote,
                    need=qty * book.ask_price,
                    min_quote=self._tight_min_quote,
                )
                return 0
        else:
            qty = max(qty, self._min_quote_size / book.bid_price)
            if balance.base - qty >= self._tight_min_base:
                self._logger.event('tight sell ok', qty=qty)
                self._market_factory.new_limit(self._i1, book.bid_price, qty, False) # убрал поток
                self._last_trade_price = book.bid_price
                return book.bid_price * qty
            else:
                self._logger.event(
                    'tight sell skip',
                    base_bal=balance.base,
                    qty=qty,
                    min_base=self._tight_min_base,
                )
                return 0

    def _execute_open_5m_tight_trade(self, request: ExecutorInput, balance: Balance) -> float:
        """Срабатывает когда узкий спред и нужно открыть пятиминутку. Нужна только для улучшения читаемости."""
        self._logger.event('try to do tight 5m opener')
        return self._execute_tight_trade(request, balance)

    def _pick_self_trade_side(self, price: float, book: BestBook, default_buy: bool) -> bool:
        """Pick which leg of the self-match is the maker, to keep the taker leg from crossing
        into a resting (e.g. MM-bot) order at the near edge of the spread.

        Near the LOWER edge the danger is the bid jumping up into a SELL taker, so we make the
        maker a SELL => taker BUY, which is immune to bid-side moves (a BUY only matches asks
        <= price; the nearest ask is a whole spread away). Near the UPPER edge we mirror it:
        maker BUY => taker SELL. In the middle of the spread the directional risk is low, so we
        keep the caller's default (buy_ratio) side untouched.

        Returns buy_trade: True => taker BUY / maker SELL.

        Applies only within the outer band of the spread. The band is a fraction of the spread
        (_edge_band_frac), floored to _edge_band_min_ticks so it stays meaningful when the spread
        is only a few ticks wide.
        """
        spread = book.ask_price - book.bid_price
        if spread <= 0:
            return default_buy

        eps = self._step_size * 1e-3
        edge_band = max(self._edge_band_frac * spread, self._edge_band_min_ticks * self._step_size) + eps
        dist_from_bid = price - book.bid_price
        dist_from_ask = book.ask_price - price

        if dist_from_bid <= edge_band and dist_from_bid <= dist_from_ask:
            buy_trade, region = True, 'lower_edge'    # maker SELL, taker BUY
        elif dist_from_ask <= edge_band and dist_from_ask < dist_from_bid:
            buy_trade, region = False, 'upper_edge'   # maker BUY, taker SELL
        else:
            return default_buy  # middle of the spread -> keep buy_ratio, no log (common no-op)

        # Only log when the edge rule actually engages.
        self._logger.event(
            'self trade side select',
            region=region,
            maker_side='sell' if buy_trade else 'buy',
            price=price, bid=book.bid_price, ask=book.ask_price,
        )
        return buy_trade

    def _execute_simple_self_trade(
        self,
        order_futures: List[Future[Any]],
        buy_trade: bool,
        price: float,
        qty: float,
        use_limit_only: bool = False,
    ) -> float:
        self._logger.event(
            'self trade preview',
            side='buy' if buy_trade else 'sell',
            p=price,
            amt=price * qty,
        )

        if use_limit_only:
            order_futures.append(
                self._thread_pool.submit(
                    self._market_factory.new_limit, 0, price, qty, not buy_trade
                )
            )
        else: # primary
            order_futures.append(
                self._thread_pool.submit(
                    self._market_factory.new_limit_maker, 0, price, qty, not buy_trade
                )
            )
        time.sleep(self._maker_taker_delay)
        order_futures.append(
            self._thread_pool.submit(
                self._market_factory.new_limit, 0, price, qty, buy_trade
            )
        )
        return price * qty

    def _execute_open_5m_self_trade(
        self,
        order_futures: List[Future[Any]],
        buy_trade: bool,
        price: float,
        qty: float,
    ) -> float:
        self._logger.event('open 5m candle self trade')

        return self._execute_simple_self_trade(
            order_futures, buy_trade, price, qty, use_limit_only=True
        )

    def _execute_rebalance_trade(
        self,
        balances: Tuple[Balance, Balance],
        book: BestBook,
        qty: float,
        price: float,
    ) -> Tuple[ExecutionStatus, float]:

        exec_amt = 0.0

        if balances[0].base < qty:
            self._logger.warning(f"base balance ({balances[0].base}) lower than trade quantity ({qty})")
            if self._check_last_price_and_log():
                return (ExecutionStatus.STOPPED, 0)
            if abs(book.ask_price / self._last_trade_price - 1) <= 0.01:
                self._market_factory.new_limit(self._i1, book.ask_price, qty, True)
                self._last_trade_price = book.ask_price
                exec_amt = qty * book.ask_price

        elif balances[0].quote < qty * price: # TODO
            self._logger.warning(f"base balance ({balances[0].quote}) lower than trade quantity ({qty * price})")
            if self._check_last_price_and_log():
                return (ExecutionStatus.STOPPED, 0)
            if abs(book.bid_price / self._last_trade_price - 1) <= 0.01:
                self._market_factory.new_limit(self._i1, book.bid_price, qty, False)
                self._last_trade_price = book.bid_price
                exec_amt = qty * book.bid_price

        elif balances[1].base < qty: # TODO: два if это х2 объём, но при этом хочется это сделать отдельным if
            self._logger.warning(f"base balance ({balances[0].base}) lower than trade quantity ({qty})")
            if self._check_last_price_and_log():
                return (ExecutionStatus.STOPPED, 0)
            if abs(book.ask_price / self._last_trade_price - 1) <= 0.01: # TODO: если else, то это неудача по сути
                self._market_factory.new_limit(self._i2, book.ask_price, qty, True)
                self._last_trade_price = book.ask_price
                exec_amt = qty * book.ask_price

        elif balances[1].quote < qty * price: # TODO
            self._logger.warning(f"base balance ({balances[0].quote}) lower than trade quantity ({qty * price})")
            if self._check_last_price_and_log():
                return (ExecutionStatus.STOPPED, 0)
            if abs(book.bid_price / self._last_trade_price - 1) <= 0.01:
                self._market_factory.new_limit(self._i2, book.bid_price, qty, False)
                self._last_trade_price = book.bid_price
                exec_amt = book.bid_price * qty

        else:
            return (ExecutionStatus.SUCCESS, 0)

        return (ExecutionStatus._SUCCESS_REBALANCED, exec_amt)

    def execute_volume_trade(self, request: ExecutorInput, snapshot: MarketSnapshot) -> Tuple[ExecutionStatus, float] | None:
        self._i1, self._i2 = self._generate_random_indexes()
        balances: Tuple[Balance, Balance] = (snapshot.balance[self._i1], snapshot.balance[self._i2])
        depth: Depth = snapshot.depth
        book: BestBook = request.book
        price = request.price

        execution_status, exec_amt = self._execute_rebalance_trade(balances, book, request.quantity, request.price)
        if execution_status == ExecutionStatus.STOPPED:
            return (ExecutionStatus.STOPPED, exec_amt)

        if execution_status == ExecutionStatus._SUCCESS_REBALANCED:
            self._logger.event('rebalance ok')
            return (ExecutionStatus.SUCCESS, exec_amt)

        if execution_status == ExecutionStatus.SUCCESS:
            order_futures: List[Future[Any]] = []

            if request.price <= book.ask_price - self._step_size / 2 and request.price >= book.bid_price + self._step_size / 2:

                buy_trade = (np.random.random() <= self._buy_ratio)

                if self._mode == VolumeMode.BBO:
                    exec_amt = self._execute_bbo_trade(order_futures, request.quantity, book, depth)

                elif self._mode == VolumeMode.SMART:
                    exec_amt = self._execute_smart_trade(order_futures, request.price, request.quantity)

                elif request.is_5m_open:
                    exec_amt = self._execute_open_5m_self_trade(
                        order_futures, buy_trade=buy_trade, price=request.price, qty=request.quantity
                    )

                else:
                    buy_trade = self._pick_self_trade_side(request.price, book, default_buy=buy_trade)
                    exec_amt = self._execute_simple_self_trade(
                        order_futures, buy_trade=buy_trade, price=request.price, qty=request.quantity
                    )

            elif self._tight_mode_enabled:
                if request.is_5m_open:
                    exec_amt = self._execute_open_5m_tight_trade(request, balances[0])
                else:
                    exec_amt = self._execute_tight_trade(request, balances[0])
                price = self._last_trade_price

            elif self._mode == VolumeMode.BBO:
                self._logger.event('BBO trade in tight...')
                exec_amt = self._execute_bbo_trade(order_futures, request.quantity, book, depth)

            else:
                self._logger.event(
                    'trade_skip_out_of_bounds',
                    price=request.price,
                    bid=book.bid_price,
                    ask=book.ask_price,
                )
                return (ExecutionStatus.FAILURE, exec_amt)

            self._last_trade_price = price

        return (ExecutionStatus.SUCCESS, exec_amt)

    def execute_pam_trade(self, request: ExecutorInput, snapshot: MarketSnapshot) -> Tuple[ExecutionStatus, float] | None:
        """Place a single limit order for periodic-auction matching (binance.pam_cex).

        Side is dictated by the quoter (strict alternation). No maker+taker pair: in PAM
        the auction matches crossed levels, so we just submit our side and wait.
        """
        if request.side is None:
            self._logger.error('execute_pam_trade: side is None')
            return (ExecutionStatus.FAILURE, 0.0)
        is_buy = bool(request.side)

        balances = snapshot.balance or []
        if balances:
            bal = balances[0]
            need_quote = request.price * request.quantity
            if is_buy and bal.quote < need_quote:
                self._logger.warning(
                    f'pam_trade buy: insufficient quote {bal.quote} < {need_quote}'
                )
                return (ExecutionStatus.FAILURE, 0.0)
            if not is_buy and bal.base < request.quantity:
                self._logger.warning(
                    f'pam_trade sell: insufficient base {bal.base} < {request.quantity}'
                )
                return (ExecutionStatus.FAILURE, 0.0)

        self._logger.event(
            'pam_trade',
            side='buy' if is_buy else 'sell',
            p=request.price,
            q=request.quantity,
            amt=request.price * request.quantity,
        )
        try:
            if self._mode == VolumeMode.SMART:
                id = np.random.randint(0, 2)
                self._market_factory.new_limit(id, request.price, request.quantity, is_buy)
                self._market_factory.new_limit(id ^ 1, request.price, request.quantity, not is_buy)
            else:
                self._market_factory.new_limit(id, request.price, request.quantity, is_buy)

        except Exception as e:
            self._logger.error(f'pam_trade place error: {e}')
            return (ExecutionStatus.FAILURE, 0.0)

        self._last_trade_price = request.price
        return (ExecutionStatus.SUCCESS, request.price * request.quantity)

    def execute_cancel_orders(self):
        def _safe_call(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
            try:
                fn(*args, **kwargs)
            except Exception as e:
                self._logger.error(f"VolumeExecutor.execute_cancel_orders exception in thread: {e}")

        Thread(target=_safe_call, args=(self._market_factory.cancel_open_orders, self._i1)).start()
        if self._i1 != self._i2:
            Thread(target=_safe_call, args=(self._market_factory.cancel_open_orders, self._i2)).start()

    def verify_last_trade_matched(self) -> Tuple[bool, float]:
        """Re-query private trades on {i1, i2} to confirm the self-trade matched.

        Returns (matched, matched_quote_volume); the volume double-counts (maker + taker legs).
        {i1, i2} covers all paths: DEFAULT/BBO use index 0, SMART uses i1 and i2."""
        matched_volume = 0.0
        trade_count = 0
        for idx in {self._i1, self._i2}:
            try:
                trades = self._market_factory.get_new_private_trades(idx)
            except Exception as e:
                self._logger.error(f'verify_last_trade_matched: get_new_private_trades({idx}) error: {e}')
                continue
            for t in (trades or []):
                trade_count += 1
                matched_volume += float(t.price) * float(t.quantity)
        return (trade_count > 0, matched_volume)

    def get_last_trade_price(self) -> Optional[float]:
        return self._last_trade_price