"""从Binance获取加密货币OHLCV数据"""

import os
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional, Literal

import pandas as pd
from binance.client import Client
import pytz
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()


class BinanceDataFetcher:
    """从Binance公开API获取和缓存OHLCV数据"""

    INTERVALS = {
        '1h': Client.KLINE_INTERVAL_1HOUR,
        '15m': Client.KLINE_INTERVAL_15MINUTE,
    }

    def __init__(self, cache_dir: str = "./data/cache"):
        """初始化Binance客户端和缓存目录

        Args:
            cache_dir: 缓存下载数据的目录
        """
        # 配置代理
        proxies = None
        use_proxy = os.getenv('USE_PROXY', 'false').lower() == 'true'

        if use_proxy:
            socks5_proxy = os.getenv('SOCKS5_PROXY', 'socks5://127.0.0.1:7890')
            proxies = {
                'http': socks5_proxy,
                'https': socks5_proxy
            }
            print(f"✓ 使用SOCKS5代理: {socks5_proxy}")

        try:
            # 初始化Binance客户端（带代理）
            if proxies:
                self.client = Client(requests_params={'proxies': proxies})
            else:
                self.client = Client()

            # 测试连接
            self.client.ping()
            print("✓ 已成功连接到Binance API")
        except Exception as e:
            print(f"⚠ Binance API连接失败: {e}")
            print("提示：")
            print("  1. 检查代理设置 (.env文件)")
            print("  2. 确保代理服务正在运行")
            print("  3. 或者使用已缓存的数据")
            self.client = None

        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def fetch_ohlcv(
        self,
        symbol: Literal['BTC', 'ETH', 'SOL'],
        timeframe: Literal['1h', '15m'],
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        use_cache: bool = True
    ) -> pd.DataFrame:
        """获取加密货币对的OHLCV数据

        Args:
            symbol: 加密货币符号 (BTC, ETH, SOL)
            timeframe: 时间框架间隔 (1h 或 15m)
            start_date: 开始日期，格式：'YYYY-MM-DD' (默认: 1年前)
            end_date: 结束日期，格式：'YYYY-MM-DD' (默认: 现在)
            use_cache: 是否使用缓存数据

        Returns:
            包含以下列的DataFrame: timestamp, open, high, low, close, volume
        """
        # 生成缓存文件名
        cache_file = self._get_cache_filename(symbol, timeframe, start_date, end_date)

        # 检查缓存
        if use_cache and cache_file.exists():
            print(f"从缓存加载数据：{cache_file}")
            return pd.read_csv(cache_file, parse_dates=['timestamp'])

        # 检查客户端是否可用
        if self.client is None:
            raise ConnectionError(
                "Binance客户端不可用。请使用VPN或使用已缓存的数据。"
            )

        # 设置默认日期
        if end_date is None:
            end_dt = datetime.now(pytz.UTC)
        else:
            end_dt = datetime.strptime(end_date, '%Y-%m-%d').replace(tzinfo=pytz.UTC)

        if start_date is None:
            start_dt = end_dt - timedelta(days=365)  # 1年前
        else:
            start_dt = datetime.strptime(start_date, '%Y-%m-%d').replace(tzinfo=pytz.UTC)

        # Binance使用USDT交易对
        pair = f"{symbol}USDT"
        interval = self.INTERVALS[timeframe]

        print(f"获取 {symbol} {timeframe} 数据，时间范围：{start_dt.date()} 至 {end_dt.date()}...")

        # 获取K线数据
        klines = self.client.get_historical_klines(
            symbol=pair,
            interval=interval,
            start_str=str(int(start_dt.timestamp() * 1000)),
            end_str=str(int(end_dt.timestamp() * 1000))
        )

        if not klines:
            raise ValueError(f"未获取到 {symbol} {timeframe} 的数据")

        # 转换为DataFrame
        df = pd.DataFrame(klines, columns=[
            'timestamp', 'open', 'high', 'low', 'close', 'volume',
            'close_time', 'quote_volume', 'trades', 'taker_buy_base',
            'taker_buy_quote', 'ignore'
        ])

        # 选择并转换列
        df = df[['timestamp', 'open', 'high', 'low', 'close', 'volume']].copy()
        df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')

        for col in ['open', 'high', 'low', 'close', 'volume']:
            df[col] = df[col].astype(float)

        # 按时间排序
        df = df.sort_values('timestamp').reset_index(drop=True)

        print(f"已获取 {len(df)} 行数据，时间范围：{df['timestamp'].min()} 至 {df['timestamp'].max()}")

        # 保存到缓存
        df.to_csv(cache_file, index=False)
        print(f"已缓存数据至：{cache_file}")

        return df

    def _get_cache_filename(
        self,
        symbol: str,
        timeframe: str,
        start_date: Optional[str],
        end_date: Optional[str]
    ) -> Path:
        """根据参数生成缓存文件名"""
        start = start_date or "1y"
        end = end_date or "now"
        filename = f"{symbol}_{timeframe}_{start}_{end}.csv"
        return self.cache_dir / filename

    def get_latest_price(self, symbol: Literal['BTC', 'ETH', 'SOL']) -> float:
        """获取符号的当前价格

        Args:
            symbol: 加密货币符号

        Returns:
            当前USDT价格
        """
        pair = f"{symbol}USDT"
        ticker = self.client.get_symbol_ticker(symbol=pair)
        return float(ticker['price'])


if __name__ == "__main__":
    # 测试数据获取器
    fetcher = BinanceDataFetcher()

    # 获取BTC 1小时数据（1年）
    df = fetcher.fetch_ohlcv('BTC', '1h')
    print("\nBTC 1小时数据:")
    print(df.head())
    print(f"\n形状：{df.shape}")
    print(f"日期范围：{df['timestamp'].min()} 至 {df['timestamp'].max()}")

    # 获取当前价格
    current_price = fetcher.get_latest_price('BTC')
    print(f"\n当前BTC价格：${current_price:,.2f}")