"""
BitMart 被动做市/高频刷单策略 (修复版 V2)
修复内容：
1. 修正 get_order_book 中解析深度数据的方式，由字典键名访问改为列表索引访问 (['price'] -> [0])
"""

import time
import requests
from typing import Optional, Dict, List, Tuple
from dataclasses import dataclass
from loguru import logger
from threading import Lock

from DrissionPage import ChromiumPage, ChromiumOptions
from bitmart.api_contract import APIContract


# ================================================================
# 📊 配置类
# ================================================================

@dataclass
class MarketMakingConfig:
    """做市策略配置"""
    # API配置（仅用于查询，不下单）
    api_key: str = "a0fb7b98464fd9bcce67e7c519d58ec10d0c38a8"
    secret_key: str = "4eaeba78e77aeaab1c2027f846a276d164f264a44c2c1bb1c5f3be50c8de1ca5"
    memo: str = "合约交易"
    contract_symbol: str = "ETHUSDT"

    # 浏览器配置
    tge_id: int = 196495  # TGE浏览器ID
    tge_url: str = "http://127.0.0.1:50326"
    tge_headers: Dict = None
    trading_url: str = "https://derivatives.bitmart.com/zh-CN/futures/ETHUSDT"

    # 做市基础参数
    spread_percent: float = 0.04  # 基础价差 (0.04% 约为 $1左右 on ETH)
    order_size_usdt: float = 10.0  # 每单金额（USDT）
    max_position_usdt: float = 100.0  # 最大持仓金额（USDT）

    # 🚀 高级策略参数
    # 库存倾斜：每持有100U，价格偏移多少。正数表示持有多单时价格下调（利于卖出，不利于买入）
    inventory_skew_factor: float = 0.0005
    # 价格容忍度：只有当(目标价 - 当前挂单价) / 目标价 > 0.05% 时才改单，避免频繁操作
    price_tolerance: float = 0.0005

    # 风险控制
    max_daily_loss: float = 50.0
    leverage: str = "30"
    open_type: str = "cross"

    def __post_init__(self):
        """初始化TGE headers"""
        if self.tge_headers is None:
            self.tge_headers = {
                "Authorization": "Bearer asp_174003986c9b0799677c5b2c1adb76e402735d753bc91a91",
                "Content-Type": "application/json"
            }


# ================================================================
# 📊 订单簿数据结构
# ================================================================

@dataclass
class OrderBook:
    """订单簿数据"""
    bids: List[Tuple[float, float]]  # [(价格, 数量), ...]
    asks: List[Tuple[float, float]]  # [(价格, 数量), ...]
    timestamp: float

    @property
    def mid_price(self) -> Optional[float]:
        """中间价"""
        if self.bids and self.asks:
            return (self.bids[0][0] + self.asks[0][0]) / 2
        return None


# ================================================================
# 📊 浏览器管理器
# ================================================================

class BrowserManager:
    """浏览器管理器：负责浏览器的启动、接管和标签页管理"""

    def __init__(self, config: MarketMakingConfig):
        self.config = config
        self.tge_port: Optional[int] = None
        self.page: Optional[ChromiumPage] = None

    def open_browser(self) -> bool:
        """打开浏览器并获取端口"""
        try:
            response = requests.post(
                f"{self.config.tge_url}/api/browser/start",
                json={"envId": self.config.tge_id},
                headers=self.config.tge_headers,
                timeout=10
            )
            data = response.json()
            if "data" in data and "port" in data["data"]:
                self.tge_port = data["data"]["port"]
                logger.success(f"成功打开浏览器，端口：{self.tge_port}")
                return True
            else:
                logger.error(f"打开浏览器响应异常: {data}")
                return False
        except Exception as e:
            logger.error(f"打开浏览器失败: {e}")
            return False

    def take_over_browser(self) -> bool:
        """接管浏览器"""
        if not self.tge_port:
            logger.error("浏览器端口未设置")
            return False

        try:
            co = ChromiumOptions()
            co.set_local_port(self.tge_port)
            self.page = ChromiumPage(addr_or_opts=co)
            logger.success("成功接管浏览器")
            return True
        except Exception as e:
            logger.error(f"接管浏览器失败: {e}")
            return False

    def close_extra_tabs(self) -> bool:
        """关闭多余的标签页，只保留第一个"""
        if not self.page:
            return False
        try:
            tabs = self.page.get_tabs()
            for idx, tab in enumerate(tabs):
                if idx == 0: continue
                tab.close()
            return True
        except Exception as e:
            logger.warning(f"关闭多余标签页失败: {e}")
            return False


# ================================================================
# 📊 浏览器交易执行器
# ================================================================

class BrowserTradingExecutor:
    """浏览器交易执行器：通过浏览器自动化下单（获取高返佣）"""

    def __init__(self, page: ChromiumPage):
        self.page = page

    def click_safe(self, xpath: str, sleep: float = 0.5) -> bool:
        """安全点击"""
        try:
            ele = self.page.ele(xpath)
            if not ele:
                return False
            ele.scroll.to_see(center=True)
            time.sleep(sleep)
            ele.click()
            return True
        except Exception as e:
            logger.error(f"点击失败 {xpath}: {e}")
            return False

    def 开单(self, marketPriceLongOrder: int = 0, limitPriceShortOrder: int = 0,
             size: Optional[float] = None, price: Optional[float] = None) -> bool:
        """开单操作"""
        size = 0.1

        try:
            # 市价单 (代码略)
            if marketPriceLongOrder == -1: pass
            elif marketPriceLongOrder == 1: pass

            # 限价单
            if limitPriceShortOrder == -1:
                # 限价做空
                if not self.click_safe('x://button[normalize-space(text()) ="限价"]'): return False
                self.page.ele('x://*[ @id="price_0"]').input(vals=price, clear=True)
                time.sleep(0.2)
                self.page.ele('x://*[ @id="size_0"]').input(vals=size, clear=True)
                if not self.click_safe('x://span[normalize-space(text()) ="卖出/做空"]'): return False
                logger.success(f"浏览器下单: 限价做空 {size} @ {price}")
                return True

            elif limitPriceShortOrder == 1:
                # 限价做多
                if not self.click_safe('x://button[normalize-space(text()) ="限价"]'): return False
                self.page.ele('x://*[ @id="price_0"]').input(vals=price, clear=True)
                time.sleep(0.2)
                self.page.ele('x://*[ @id="size_0"]').input(vals=size, clear=True)
                if not self.click_safe('x://span[normalize-space(text()) ="买入/做多"]'): return False
                logger.success(f"浏览器下单: 限价做多 {size} @ {price}")
                return True

            return False
        except Exception as e:
            logger.error(f"开单异常: {e}")
            return False

    def place_limit_order(self, side: str, price: float, size: float) -> bool:
        """统一接口"""
        if side == "buy":
            return self.开单(limitPriceShortOrder=1, size=size, price=price)
        else:
            return self.开单(limitPriceShortOrder=-1, size=size, price=price)


# ================================================================
# 📊 BitMart API 封装 (修复 get_order_book)
# ================================================================

class BitMartMarketMakerAPI:
    """BitMart做市API封装（仅用于查询，不下单）"""

    def __init__(self, config: MarketMakingConfig):
        self.config = config
        self.contractAPI = APIContract(
            config.api_key,
            config.secret_key,
            config.memo,
            timeout=(5, 15)
        )

    def get_order_book(self) -> Optional[OrderBook]:
        try:
            # 移除不支持的 limit 参数
            response = self.contractAPI.get_depth(contract_symbol=self.config.contract_symbol)[0]

            if response.get('code') == 1000:
                data = response.get('data', {})
                bids = []
                asks = []
                # 解析数据
                if isinstance(data, dict):
                    bids_raw = data.get('bids', [])
                    asks_raw = data.get('asks', [])

                    # 修复：b 是列表 [price, size]，不是字典
                    for b in bids_raw[:10]:
                        # b[0] 是价格, b[1] 是数量
                        bids.append((float(b[0]), float(b[1])))

                    for a in asks_raw[:10]:
                        # a[0] 是价格, a[1] 是数量
                        asks.append((float(a[0]), float(a[1])))

                if bids and asks:
                    return OrderBook(bids=bids, asks=asks, timestamp=time.time())
            else:
                logger.warning(f"获取深度失败: {response}")
            return None
        except Exception as e:
            logger.error(f"获取订单簿异常: {e}")
            return None

    def get_position_net(self) -> float:
        """获取净持仓 (多为正，空为负)"""
        try:
            response = self.contractAPI.get_position(contract_symbol=self.config.contract_symbol)[0]
            if response.get('code') == 1000:
                data = response.get('data', [])
                if data:
                    pos = data[0]
                    current_amount = float(pos.get('current_amount', 0))
                    position_type = int(pos.get('position_type', 0)) # 1多 2空
                    if position_type == 1: return current_amount
                    if position_type == 2: return -current_amount
            return 0.0
        except Exception as e:
            logger.error(f"持仓查询异常: {e}")
            return 0.0

    def get_open_orders(self) -> List[Dict]:
        """获取当前挂单"""
        try:
            resp = self.contractAPI.get_open_order(contract_symbol=self.config.contract_symbol)[0]
            if resp.get("code") == 1000:
                return resp.get("data", [])
            return []
        except Exception as e:
            logger.error(f"查询挂单异常: {e}")
            return []

    def cancel_order(self, order_id: str) -> bool:
        """API撤单"""
        try:
            resp = self.contractAPI.post_cancel_order(contract_symbol=self.config.contract_symbol, order_id=order_id)[0]
            return resp.get("code") == 1000
        except Exception as e:
            logger.error(f"API撤单异常: {e}")
            return False

    def set_leverage(self):
        try:
            self.contractAPI.post_submit_leverage(contract_symbol=self.config.contract_symbol, leverage=self.config.leverage, open_type=self.config.open_type)
        except:
            pass


# ================================================================
# 🧠 策略核心
# ================================================================

class MarketMakingStrategy:
    """优化版被动做市策略"""

    def __init__(self, config: MarketMakingConfig):
        self.config = config
        self.api = BitMartMarketMakerAPI(config)
        self.browser_manager = BrowserManager(config)
        self.trading_executor: Optional[BrowserTradingExecutor] = None
        self.running = False

        # 初始化流程
        if not self._initialize_browser():
            raise Exception("浏览器初始化失败")
        self.api.set_leverage()

    def _initialize_browser(self) -> bool:
        try:
            if not self.browser_manager.open_browser(): return False
            if not self.browser_manager.take_over_browser(): return False
            self.browser_manager.close_extra_tabs()

            # 访问交易页
            logger.info(f"正在访问交易页: {self.config.trading_url}")
            self.browser_manager.page.get(self.config.trading_url)
            time.sleep(3)

            self.trading_executor = BrowserTradingExecutor(self.browser_manager.page)
            logger.success("浏览器环境就绪")
            return True
        except Exception as e:
            logger.error(f"浏览器初始化异常: {e}")
            return False

    def calculate_target_prices(self, mid_price: float, net_position: float) -> Tuple[float, float]:
        """核心算法：计算考虑了库存倾斜的目标买卖价"""
        # 1. 基础价差的一半
        half_spread = mid_price * (self.config.spread_percent / 100) / 2

        # 2. 库存倾斜调整
        skew_adjust = net_position * self.config.inventory_skew_factor * mid_price

        quote_mid = mid_price - skew_adjust

        target_bid = quote_mid - half_spread
        target_ask = quote_mid + half_spread

        # 3. 价格修正 (防止穿仓)
        if target_ask <= target_bid:
            target_ask = target_bid + mid_price * 0.0001

        return round(target_bid, 2), round(target_ask, 2)

    def reconcile_orders(self, target_bid: float, target_ask: float):
        """调节逻辑：对比API实际挂单 vs 目标价格"""
        open_orders = self.api.get_open_orders()

        current_bids = []
        current_asks = []

        for o in open_orders:
            side = o.get('side')
            # 兼容API返回
            side_str = str(side).lower()
            if side_str == '1' or 'buy' in side_str:
                current_bids.append(o)
            elif side_str == '2' or 'sell' in side_str:
                current_asks.append(o)

        # --- 调节买单 ---
        valid_bid_exists = False
        for order in current_bids:
            price = float(order.get('price', 0))
            diff_pct = abs(price - target_bid) / target_bid

            if diff_pct < self.config.price_tolerance:
                valid_bid_exists = True
            else:
                logger.info(f"♻️ 买单价格偏离 (现{price} vs 标{target_bid})，撤单")
                self.api.cancel_order(order.get('order_id') or order.get('id'))

        if not valid_bid_exists:
            # 计算张数
            size_contract = self.config.order_size_usdt / target_bid / 0.01
            size_contract = max(1, int(size_contract))
            logger.info(f"➕ 补买单: {target_bid} (数量:{size_contract})")
            self.trading_executor.place_limit_order("buy", target_bid, size_contract)

        # --- 调节卖单 ---
        valid_ask_exists = False
        for order in current_asks:
            price = float(order.get('price', 0))
            diff_pct = abs(price - target_ask) / target_ask

            if diff_pct < self.config.price_tolerance:
                valid_ask_exists = True
            else:
                logger.info(f"♻️ 卖单价格偏离 (现{price} vs 标{target_ask})，撤单")
                self.api.cancel_order(order.get('order_id') or order.get('id'))

        if not valid_ask_exists:
            size_contract = self.config.order_size_usdt / target_ask / 0.01
            size_contract = max(1, int(size_contract))
            logger.info(f"➕ 补卖单: {target_ask} (数量:{size_contract})")
            self.trading_executor.place_limit_order("sell", target_ask, size_contract)

    def run(self):
        self.running = True
        logger.info("🚀 策略已启动")

        while self.running:
            try:
                # 1. 获取市场数据
                ob = self.api.get_order_book()
                if not ob:
                    time.sleep(1)
                    continue
                mid_price = ob.mid_price

                # 2. 获取持仓
                net_position = self.api.get_position_net()

                # 3. 计算目标价
                t_bid, t_ask = self.calculate_target_prices(mid_price, net_position)
                logger.info(f"Mid:{mid_price:.2f} | Pos:{net_position} | Target Bid:{t_bid} Ask:{t_ask}")

                # 4. 调节挂单
                self.reconcile_orders(t_bid, t_ask)

                # 5. 循环间隔
                time.sleep(3)

            except KeyboardInterrupt:
                logger.warning("停止策略")
                break
            except Exception as e:
                logger.error(f"Loop Exception: {e}")
                time.sleep(2)

if __name__ == '__main__':
    config = MarketMakingConfig(
        contract_symbol="ETHUSDT",
        spread_percent=0.04,
        order_size_usdt=0.1,
        max_position_usdt=50.0,
        inventory_skew_factor=0.0005,
        price_tolerance=0.0005
    )

    strategy = MarketMakingStrategy(config)
    strategy.run()
