"""
BitMart 30分钟K线数据抓取脚本
从 BitMart API 获取30分钟K线数据并存储到数据库
"""

import time
from loguru import logger
from bitmart.api_contract import APIContract
from models.bitmart import BitMart30


class BitMartDataCollector:
    def __init__(self):
        self.api_key = "a0fb7b98464fd9bcce67e7c519d58ec10d0c38a8"
        self.secret_key = "4eaeba78e77aeaab1c2027f846a276d164f264a44c2c1bb1c5f3be50c8de1ca5"
        self.memo = "数据抓取"
        self.contract_symbol = "ETHUSDT"
        self.contractAPI = APIContract(self.api_key, self.secret_key, self.memo, timeout=(5, 15))

    def get_klines(self, start_time=None, end_time=None, limit=200):
        """
        获取K线数据
        :param start_time: 开始时间戳（秒级）
        :param end_time: 结束时间戳（秒级）
        :param limit: 获取数量限制
        :return: K线数据列表
        """
        try:
            if not end_time:
                end_time = int(time.time())
            if not start_time:
                start_time = end_time - 3600 * 24 * 7  # 默认获取最近7天

            response = self.contractAPI.get_kline(
                contract_symbol=self.contract_symbol,
                step=30,  # 30分钟
                start_time=start_time,
                end_time=end_time
            )[0]

            if response['code'] != 1000:
                logger.error(f"获取K线失败: {response}")
                return []

            klines = response.get('data', [])
            formatted = []
            for k in klines:
                # BitMart API 返回的时间戳是秒级，需要转换为毫秒级
                # 根据 bitmart/框架.py 中的使用方式，API返回的是秒级时间戳
                timestamp_ms = int(k["timestamp"]) * 1000
                
                formatted.append({
                    'id': timestamp_ms,
                    'open': float(k["open_price"]),
                    'high': float(k["high_price"]),
                    'low': float(k["low_price"]),
                    'close': float(k["close_price"])
                })
            
            # 按时间戳排序
            formatted.sort(key=lambda x: x['id'])
            return formatted
        except Exception as e:
            logger.error(f"获取K线异常: {e}")
            return []

    def save_klines(self, klines):
        """
        保存K线数据到数据库
        :param klines: K线数据列表
        :return: 保存的数量
        """
        saved_count = 0
        for kline in klines:
            try:
                BitMart30.get_or_create(
                    id=kline['id'],
                    defaults={
                        'open': kline['open'],
                        'high': kline['high'],
                        'low': kline['low'],
                        'close': kline['close'],
                    }
                )
                saved_count += 1
            except Exception as e:
                logger.error(f"保存K线数据失败 {kline['id']}: {e}")
        
        return saved_count

    def collect_historical_data(self, days=30):
        """
        抓取历史数据
        :param days: 抓取最近多少天的数据
        """
        logger.info(f"开始抓取 BitMart {self.contract_symbol} 最近 {days} 天的30分钟K线数据")
        
        end_time = int(time.time())
        start_time = end_time - 3600 * 24 * days
        
        # 分批获取，每次获取7天的数据
        batch_days = 7
        total_saved = 0
        
        current_start = start_time
        while current_start < end_time:
            current_end = min(current_start + 3600 * 24 * batch_days, end_time)
            
            logger.info(f"抓取时间段: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(current_start))} "
                       f"到 {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(current_end))}")
            
            klines = self.get_klines(start_time=current_start, end_time=current_end)
            if klines:
                saved = self.save_klines(klines)
                total_saved += saved
                logger.info(f"本批次保存 {saved} 条数据，累计 {total_saved} 条")
            else:
                logger.warning("本批次未获取到数据")
            
            current_start = current_end
            time.sleep(1)  # 避免请求过快
        
        logger.success(f"数据抓取完成，共保存 {total_saved} 条K线数据")

    def collect_realtime_data(self):
        """
        实时抓取最新数据（用于定时任务）
        """
        logger.info("开始抓取 BitMart 最新30分钟K线数据")
        
        # 获取最近1小时的数据（确保能获取到最新的K线）
        end_time = int(time.time())
        start_time = end_time - 3600 * 2  # 最近2小时
        
        klines = self.get_klines(start_time=start_time, end_time=end_time)
        if klines:
            saved = self.save_klines(klines)
            logger.success(f"保存 {saved} 条最新K线数据")
        else:
            logger.warning("未获取到最新数据")


if __name__ == '__main__':
    collector = BitMartDataCollector()
    
    # 抓取最近30天的历史数据
    collector.collect_historical_data(days=365)
    
    # 如果需要实时抓取，可以取消下面的注释
    # collector.collect_realtime_data()
