"""
Margin Grid Bot — Grid engine.
Builds a grid of BUY/SELL orders around a center price.
Supports both real trading and demo simulation.
"""

import math
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional

@dataclass
class GridLevel:
    level_id: int           # negative=below center (BUY), positive=above center (SELL)
    price: float
    side: str               # "BUY" or "SELL"
    qty: float
    step_num: int           # 1..N from center
    filled: bool = False

    def to_dict(self):
        return {
            "level_id": self.level_id,
            "price": round(self.price, 8),
            "side": self.side,
            "qty": self.qty,
            "step_num": self.step_num,
            "filled": self.filled,
        }


@dataclass
class GridState:
    center_price: float = 0.0
    step_percent: float = 0.005
    step_price: float = 0.0
    levels: list = field(default_factory=list)
    total_levels: int = 0
    take_profit_percent: float = 0.5
    symbol: str = "USDC_BTC"
    mode: str = "demo"        # "demo" or "live"
    last_update: float = 0.0


class GridEngine:
    """
    Grid trading engine.
    
    Buy orders below center:   center - n×step  (level_id: -1, -2, ... -N)
    Sell orders above center:  center + n×step  (level_id: +1, +2, ... +N)
    
    Each round-trip = profit = step × qty.
    When a level is filled → check if total PnL >= take_profit → close position.
    """

    def __init__(
        self,
        levels: int = 10,
        step_percent: float = 0.005,
        take_profit_percent: float = 0.5,
        symbol: str = "USDC_BTC",
        mode: str = "demo",
        grid_mode: str = "arithmetic",
    ):
        self.levels_count = levels
        self.step_percent = step_percent
        self.take_profit_percent = take_profit_percent
        self.symbol = symbol
        self.mode = mode
        # 'arithmetic' — равные $шаги; 'geometric' — равные %шаги.
        # Для BTC на широком диапазоне geometric точнее (см. QuantDinger levels.py).
        self.grid_mode = str(grid_mode or "arithmetic").lower().strip()
        if self.grid_mode not in ("arithmetic", "geometric"):
            self.grid_mode = "arithmetic" 

    def calc_step(self, center_price: float) -> float:
        """Absolute step size at the center price.

        arithmetic: constant $ step = center * step_percent
        geometric:  multiplicative — step at level n is roughly center * step_percent * ratio^n.
                   The constant returned is the *level-1* step in $.
        """
        if self.grid_mode == "geometric" and center_price > 0 and self.step_percent > 0:
            # Return nominal $ step at center for compat with old code paths
            # (rebuild threshold, etc.). Real price spacing is computed in
            # get_grid_levels() via multiplicative ratio.
            return center_price * self.step_percent
        return center_price * self.step_percent

    def get_grid_levels(
        self,
        center_price: float,
        base_qty: float = 0.0001,
        buy_center: float | None = None,
        sell_center: float | None = None,
    ) -> list[GridLevel]:
        """Build grid levels.

        buy_center  - anchor for BUY levels  (default: center_price)
        sell_center - anchor for SELL levels (default: center_price)
        If they equal center_price, behaviour matches the legacy symmetric grid.
        If they differ, BUY is built around buy_center and SELL around sell_center.
        Used by Kronos logic: at bias=up the BUY anchor stays near current_price (aggressive),
        the SELL anchor is shifted toward the forecasted price; at bias=down - mirrored.

        Modes:
          - arithmetic (default): levels at center ± n*step$
          - geometric:             levels at center / ratio^n  and  center * ratio^n
                                   ratio = 1 + step_percent (constant % step).
        """
        if buy_center is None:
            buy_center = center_price
        if sell_center is None:
            sell_center = center_price
        levels = []

        if self.grid_mode == "geometric" and self.step_percent > 0:
            ratio = 1.0 + self.step_percent
            for n in range(1, self.levels_count + 1):
                factor = ratio ** n
                buy_price = buy_center / factor
                if buy_price <= 0:
                    continue
                levels.append(GridLevel(
                    level_id=-n,
                    price=round(buy_price, 8),
                    side="BUY",
                    qty=base_qty,
                    step_num=n,
                ))
                sell_price = sell_center * factor
                levels.append(GridLevel(
                    level_id=+n,
                    price=round(sell_price, 8),
                    side="SELL",
                    qty=base_qty,
                    step_num=n,
                ))
        else:
            step = self.calc_step(center_price)
            for n in range(1, self.levels_count + 1):
                offset = n * step
                buy_price = buy_center - offset
                if buy_price <= 0:
                    continue
                levels.append(GridLevel(
                    level_id=-n,
                    price=round(buy_price, 8),
                    side="BUY",
                    qty=base_qty,
                    step_num=n,
                ))
                sell_price = sell_center + offset
                levels.append(GridLevel(
                    level_id=+n,
                    price=round(sell_price, 8),
                    side="SELL",
                    qty=base_qty,
                    step_num=n,
                ))

        # Sort: BUY first (negatives), then SELL (positives)
        levels.sort(key=lambda x: x.level_id)
        return levels

    def build_state(self, current_price: float, base_qty: float = 0.001) -> GridState:
        levels = self.get_grid_levels(current_price, base_qty)
        step = self.calc_step(current_price)

        return GridState(
            center_price=current_price,
            step_percent=self.step_percent,
            step_price=step,
            levels=[l.to_dict() for l in levels],
            total_levels=len(levels),
            take_profit_percent=self.take_profit_percent,
            symbol=self.symbol,
            mode=self.mode,
            last_update=0.0,
        )

    def should_rebuild(self, current_price: float, center_price: float, threshold: float = 3) -> bool:
        """Rebuild grid if price moved by threshold levels."""
        step = self.calc_step(center_price)
        if step == 0:
            return False
        levels_moved = abs(current_price - center_price) / step
        return levels_moved >= threshold


@dataclass
class DemoPosition:
    entry_price: float
    qty: float
    side: str               # "BUY" or "SELL"
    opened_at: float = 0.0

    def current_value(self, current_price: float) -> float:
        if self.side == "BUY":
            return self.qty * current_price
        else:
            return self.qty * (2 * self.entry_price - current_price)

    def pnl_percent(self, current_price: float) -> float:
        if self.side == "BUY":
            return (current_price - self.entry_price) / self.entry_price * 100
        else:
            return (self.entry_price - current_price) / self.entry_price * 100


class DemoSimulator:
    """
    Simulates trading on a virtual wallet of DEMO_START_BALANCE USDC.
    Uses real market prices from Tradernet.
    """

    def __init__(self, start_balance: float = 100.0):
        self.start_balance = start_balance
        self.balance = start_balance       # USDC
        self.position: Optional[DemoPosition] = None
        self.orders: list = []             # simulated open orders
        self.trade_log: list = []

    def reset(self):
        self.balance = self.start_balance
        self.position = None
        self.orders = []
        self.trade_log = []

    def place_buy_order(self, price: float, qty: float, level_id: int):
        """Simulate placing a BUY limit order.
        Auto-shrinks qty if balance is insufficient for full size.
        Min qty = 0.0001 BTC (Tradernet minimum lot for BTC-USDT).
        """
        min_qty = 0.0001
        cost = price * qty
        if self.balance >= cost:
            # Full size fits
            self.balance -= cost
            self.orders.append({"side": "BUY", "price": price, "qty": qty, "level_id": level_id, "filled": False})
            return True
        # Try to fit a smaller qty into available balance
        affordable_qty = round(self.balance / price, 8) if price > 0 else 0
        if affordable_qty >= min_qty:
            cost = price * affordable_qty
            self.balance -= cost
            self.orders.append({"side": "BUY", "price": price, "qty": affordable_qty, "level_id": level_id, "filled": False})
            return True
        # Insufficient balance for even min lot — warn loudly (was silent before 15.06 patch)
        import logging
        logging.getLogger(__name__).warning(
            f"[DEMO] skip BUY @ ${price:,.2f} (qty={qty} BTC): "
            f"balance=${self.balance:.2f} < min_cost=${min_qty*price:.2f} "
            f"(short by ${min_qty*price - self.balance:.2f}); level_id={level_id}"
        )
        return False

    def place_sell_order(self, price: float, qty: float, level_id: int):
        """Simulate placing a SELL limit order."""
        if self.position and self.position.side == "BUY" and self.position.qty >= qty:
            self.orders.append({"side": "SELL", "price": price, "qty": qty, "level_id": level_id, "filled": False})
            return True
        return False

    def check_fill_buy(self, current_price: float) -> list[dict]:
        """Check if any BUY orders should be filled at current_price."""
        filled = []
        for order in self.orders:
            if not order["filled"] and order["side"] == "BUY" and current_price <= order["price"]:
                order["filled"] = True
                self.position = DemoPosition(
                    entry_price=order["price"],
                    qty=order["qty"],
                    side="BUY",
                    opened_at=time.time(),
                )
                self.balance += order["price"] * order["qty"]   # rest of balance
                filled.append(order)
        self.orders = [o for o in self.orders if not o["filled"]]
        return filled

    def _cancel_open_sells(self, reason: str) -> int:
        """Отменить все висящие (не-filled) SELL-ордера. Возвращает кол-во отменённых.
        Вызывается при ЗАКРЫТИИ позиции — иначе SELL-ордер висит и при fill
        без открытой позиции даёт фантомный pnl=0 и теряет proceeds.
        """
        before = len(self.orders)
        self.orders = [o for o in self.orders if o["filled"] or o["side"] != "SELL"]
        cancelled = before - len(self.orders)
        if cancelled > 0:
            # Логгируем на уровне INFO через глобальный logger, если есть
            import logging
            logging.getLogger("grid").info(
                f"[DEMO] cancelled {cancelled} open SELL order(s): {reason}"
            )
        return cancelled

    def cancel_open_buys(self, reason: str = "") -> int:
        """Cancel all unfilled BUY orders AND refund reserved balance.

        For BUY orders the cost (price * qty) was debited at
        place_buy_order — cancelling must credit it back, otherwise
        every grid rebuild silently drains the wallet. Returns the
        count of cancelled orders. SELL orders are not touched
        (no balance was reserved for them).
        """
        new_orders = []
        refunded = 0.0
        cancelled = 0
        for o in self.orders:
            if o["side"] == "BUY" and not o.get("filled"):
                refunded += o["price"] * o["qty"]
                cancelled += 1
            else:
                new_orders.append(o)
        self.orders = new_orders
        if cancelled > 0:
            self.balance += refunded
            import logging
            logging.getLogger("grid").info(
                f"[DEMO] cancelled {cancelled} open BUY order(s) and refunded "
                f"${refunded:.2f}; reason={reason}"
            )
        return cancelled

    def check_fill_sell(self, current_price: float) -> list[dict]:
        """Check if any SELL orders should be filled at current_price."""
        filled = []
        for order in self.orders:
            if not order["filled"] and order["side"] == "SELL" and current_price >= order["price"]:
                order["filled"] = True
                # Close BUY position
                if self.position and self.position.side == "BUY":
                    proceeds = order["qty"] * order["price"]
                    self.balance += proceeds
                    pnl = proceeds - (self.position.qty * self.position.entry_price)
                    exit_ts = time.time()
                    self.trade_log.append({
                        "side": "BUY→SELL",
                        "entry": self.position.entry_price,
                        "exit": order["price"],
                        "qty": order["qty"],
                        "pnl_usdc": round(pnl, 8),
                        "pnl_pct": round(pnl / (self.position.qty * self.position.entry_price) * 100, 2),
                        "entry_time": datetime.fromtimestamp(self.position.opened_at, tz=timezone.utc).isoformat() if self.position.opened_at else None,
                        "exit_time": datetime.fromtimestamp(exit_ts, tz=timezone.utc).isoformat(),
                    })
                    self.position = None
                    # BUGFIX: SELL закрыл позицию → остальные SELL ордера-зомби
                    # больше не нужны, отменяем.
                    self._cancel_open_sells("position closed via SELL fill")
                else:
                    # SELL-zombie: ордер заполнился, но позиции уже нет.
                    # Не начисляем proceeds (нет BTC), логгируем warning.
                    import logging
                    logging.getLogger("grid").warning(
                        f"[DEMO] SELL-zombie fill @ ${order['price']:,.2f} qty={order['qty']} BTC "
                        f"— no position to close, no proceeds credited"
                    )
                filled.append(order)
        self.orders = [o for o in self.orders if not o["filled"]]
        return filled

    def check_take_profit(self, current_price: float, take_profit_pct: float) -> bool:
        """If position PnL >= take_profit_pct, close it."""
        if not self.position:
            return False
        pnl_pct = self.position.pnl_percent(current_price)
        if pnl_pct >= take_profit_pct:
            # Force close at current price
            if self.position.side == "BUY":
                proceeds = self.position.qty * current_price
                self.balance += proceeds
                pnl = proceeds - (self.position.qty * self.position.entry_price)
                exit_ts = time.time()
                self.trade_log.append({
                    "side": "TAKE_PROFIT",
                    "entry": self.position.entry_price,
                    "exit": current_price,
                    "qty": self.position.qty,
                    "pnl_usdc": round(pnl, 8),
                    "pnl_pct": round(pnl_pct, 2),
                    "entry_time": datetime.fromtimestamp(self.position.opened_at, tz=timezone.utc).isoformat() if self.position.opened_at else None,
                    "exit_time": datetime.fromtimestamp(exit_ts, tz=timezone.utc).isoformat(),
                })
                self.position = None
                # BUGFIX: TP закрыл позицию → SELL-ордера-зомби отменяем
                self._cancel_open_sells("position closed via TAKE_PROFIT")
                return True
        return False

    def get_status(self, current_price: float) -> dict:
        pnl_usdc = 0.0
        pnl_pct = 0.0
        unrealized = 0.0

        if self.position:
            if self.position.side == "BUY":
                unrealized = (current_price - self.position.entry_price) * self.position.qty
                pnl_pct = self.position.pnl_percent(current_price)
                pnl_usdc = unrealized

        return {
            "balance": round(self.balance, 8),
            "position": {
                "side": self.position.side if self.position else None,
                "entry_price": self.position.entry_price if self.position else None,
                "qty": self.position.qty if self.position else None,
                "current_value": round(self.position.current_value(current_price) if self.position else 0.0, 8),
                "unrealized_pnl": round(unrealized, 8),
                "pnl_pct": round(pnl_pct, 2),
            } if self.position else None,
            "open_orders": len(set(o["level_id"] for o in self.orders if not o["filled"])),
            "total_trades": len(self.trade_log),
            "total_pnl": round(sum(t["pnl_usdc"] for t in self.trade_log), 8),
            "last_trade": self.trade_log[-1] if self.trade_log else None,
        }