#!/usr/bin/env python3
"""
mt5_data_fetcher.py — MT5 历史数据获取器
替代 yfinance，使用 MetaTrader5 Python 包获取真实的 XAUUSD 历史K线数据。
输出与 yfinance 兼容的 pandas DataFrame 格式。

用法：
    from src.mt5_data_fetcher import fetch_mt5_history
    df = fetch_mt5_history(symbol="XAUUSD", timeframe="M5", days=30)
"""
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd

ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT))

DATA_DIR = ROOT / "data" / "offline"

TIMEFRAME_MAP = {
    "1m": "TIMEFRAME_M1",
    "5m": "TIMEFRAME_M5",
    "15m": "TIMEFRAME_M15",
    "30m": "TIMEFRAME_M30",
    "1h": "TIMEFRAME_H1",
    "4h": "TIMEFRAME_H4",
    "1d": "TIMEFRAME_D1",
}

# 近似每个时间框架每天K线数量（用于估算fetch数量）
BARS_PER_DAY = {
    "1m": 1440, "5m": 288, "15m": 96, "30m": 48,
    "1h": 24, "4h": 6, "1d": 1,
}


def _get_mt5_timeframe(tf_str: str):
    """将字符串时间框架转换为 MT5 常量"""
    try:
        import MetaTrader5 as mt5
    except ImportError:
        return None
    attr = TIMEFRAME_MAP.get(tf_str.lower(), "TIMEFRAME_M5")
    return getattr(mt5, attr, mt5.TIMEFRAME_M5)


def _mt5_available() -> bool:
    """检查 MT5 Python 包是否可用并尝试连接"""
    try:
        import MetaTrader5 as mt5
        if not mt5.initialize():
            mt5.shutdown()
            return False
        mt5.shutdown()
        return True
    except Exception:
        return False


def _fetch_from_mt5(symbol: str, timeframe: str, days: int) -> Optional[pd.DataFrame]:
    """从 MT5 获取历史数据"""
    try:
        import MetaTrader5 as mt5
        import pytz
    except ImportError as e:
        print(f"[mt5_data] MetaTrader5 或 pytz 包未安装: {e}")
        return None

    if not mt5.initialize():
        print("[mt5_data] MT5 initialize() 失败，确保 MT5 客户端正在运行")
        return None

    try:
        mt5_tf = _get_mt5_timeframe(timeframe)
        if mt5_tf is None:
            print(f"[mt5_data] 不支持的时间框架: {timeframe}")
            return None

        utc_tz = pytz.timezone("Etc/UTC")
        utc_to = datetime.now(utc_tz)
        bars_count = max(days * BARS_PER_DAY.get(timeframe, 288), 100)

        rates = mt5.copy_rates_from(symbol, mt5_tf, utc_to, bars_count)
        if rates is None or len(rates) == 0:
            error = mt5.last_error()
            print(f"[mt5_data] copy_rates_from 返回空: {error}")
            return None

        df = pd.DataFrame(rates)
        df["time"] = pd.to_datetime(df["time"], unit="s")
        df.set_index("time", inplace=True)
        df.rename(columns={
            "open": "Open", "high": "High", "low": "Low",
            "close": "Close", "tick_volume": "Volume",
        }, inplace=True)

        print(f"[mt5_data] MT5 获取成功: {symbol} {timeframe} {days}d → {len(df)} 根K线")
        return df

    finally:
        mt5.shutdown()


def _fetch_from_cache(symbol: str, timeframe: str, days: int) -> Optional[pd.DataFrame]:
    """从缓存的 CSV 文件读取"""
    csv_path = DATA_DIR / f"mt5_{symbol}_{days}d_{timeframe.lower()}.csv"
    if not csv_path.exists():
        return None
    try:
        df = pd.read_csv(csv_path, index_col=0, parse_dates=True)
        if len(df) > 0:
            print(f"[mt5_data] 缓存读取: {csv_path} → {len(df)} 根K线")
            return df
    except Exception as e:
        print(f"[mt5_data] 缓存读取失败: {e}")
    return None


def _save_cache(df: pd.DataFrame, symbol: str, timeframe: str, days: int):
    """保存到缓存"""
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    csv_path = DATA_DIR / f"mt5_{symbol}_{days}d_{timeframe.lower()}.csv"
    df.to_csv(csv_path)
    print(f"[mt5_data] 缓存已保存: {csv_path}")


def _fetch_from_yfinance(symbol: str, timeframe: str, days: int) -> Optional[pd.DataFrame]:
    """yfinance 兜底（仅 MT5 不可用时）"""
    try:
        import yfinance as yf
    except ImportError:
        return None

    # yfinance 使用不同的 symbol 和 interval 格式
    yf_interval_map = {
        "1m": "1m", "5m": "5m", "15m": "15m", "30m": "30m",
        "1h": "60m", "4h": "1h", "1d": "1d",
    }
    yf_symbol = "GC=F" if symbol == "XAUUSD" else symbol
    yf_interval = yf_interval_map.get(timeframe, "5m")

    for attempt in range(3):
        try:
            tk = yf.Ticker(yf_symbol)
            df = tk.history(period=f"{days}d", interval=yf_interval)
            if df is not None and len(df) > 0:
                print(f"[mt5_data] yfinance 兜底: {yf_symbol} {days}d → {len(df)} 根K线")
                return df
        except Exception as e:
            print(f"[mt5_data] yfinance 尝试 {attempt+1}/3 失败: {e}")
            if attempt < 2:
                import time
                time.sleep(3)
    return None


def fetch_history(
    symbol: str = "XAUUSD",
    timeframe: str = "M5",
    days: int = 30,
    prefer_mt5: bool = True,
    use_cache: bool = True,
) -> pd.DataFrame:
    """
    获取历史K线数据。

    优先级：MT5 > 缓存 > yfinance（兜底）

    Args:
        symbol: 品种代码，默认 XAUUSD
        timeframe: 时间框架 (M1, M5, M15, M30, H1, H4, D1)
        days: 拉取天数
        prefer_mt5: 优先使用 MT5（False 时直接用 yfinance）
        use_cache: 使用缓存

    Returns:
        pandas DataFrame with columns: Open, High, Low, Close, Volume
        Index: datetime
    """
    timeframe = timeframe.lower()

    # 1. 尝试缓存
    if use_cache:
        df = _fetch_from_cache(symbol, timeframe, days)
        if df is not None and len(df) >= 50:
            return df

    # 2. 尝试 MT5
    if prefer_mt5:
        df = _fetch_from_mt5(symbol, timeframe, days)
        if df is not None and len(df) >= 50:
            if use_cache:
                _save_cache(df, symbol, timeframe, days)
            return df
        print("[mt5_data] MT5 不可用，降级到 yfinance...")

    # 3. yfinance 兜底
    df = _fetch_from_yfinance(symbol, timeframe, days)
    if df is not None and len(df) >= 50:
        if use_cache:
            _save_cache(df, symbol, timeframe, days)
        return df

    raise RuntimeError(
        f"无法获取数据: {symbol} {timeframe} {days}d。"
        f"请确保 MT5 客户端正在运行，或网络可访问 Yahoo Finance。"
    )


def fetch_history_multi_timeframe(
    symbol: str = "XAUUSD",
    timeframes: list[str] = None,
    days: int = 30,
) -> dict[str, pd.DataFrame]:
    """
    批量获取多个时间框架的数据。

    Returns:
        {"M5": DataFrame, "M15": DataFrame, "H1": DataFrame}
    """
    if timeframes is None:
        timeframes = ["M5", "M15", "H1"]
    results = {}
    for tf in timeframes:
        try:
            results[tf] = fetch_history(symbol, tf, days)
        except Exception as e:
            print(f"[mt5_data] {tf} 获取失败: {e}")
            results[tf] = pd.DataFrame()
    return results
