buffer_platform/app/storage_manager.py

import datetime as _datetime
import json
import logging
import time
from datetime import datetime

from sqlalchemy import text
from sqlalchemy.orm import sessionmaker

from app.models import MarketData, SymbolTimestamp

logger = logging.getLogger(__name__)

# Redis 缓存 TTL：30 天
CACHE_TTL_SECONDS = 30 * 24 * 60 * 60


class StorageManager:
    """存储管理器骨架，封装 Redis / MySQL 可用性检测与三级降级状态。"""

    def __init__(self, redis_client=None, mysql_engine=None, check_interval=30):
        self.redis_client = redis_client
        self.mysql_engine = mysql_engine
        self.check_interval = check_interval

        self.redis_available = False
        self.mysql_available = False
        self.last_redis_check = 0
        self.last_mysql_check = 0

    def check_redis(self):
        """惰性检测 Redis 可用性，30 秒间隔内返回缓存结果。"""
        now = time.time()
        if now - self.last_redis_check < self.check_interval:
            return self.redis_available

        if self.redis_client is None:
            self.redis_available = False
        else:
            try:
                self.redis_client.ping()
                self.redis_available = True
            except Exception as e:
                self.redis_available = False

        self.last_redis_check = now
        return self.redis_available

    def check_mysql(self):
        """惰性检测 MySQL 可用性，30 秒间隔内返回缓存结果。"""
        now = time.time()
        if now - self.last_mysql_check < self.check_interval:
            return self.mysql_available

        if self.mysql_engine is None:
            self.mysql_available = False
        else:
            try:
                with self.mysql_engine.connect() as conn:
                    conn.execute(text("SELECT 1"))
                self.mysql_available = True
            except Exception as e:
                self.mysql_available = False

        self.last_mysql_check = now
        return self.mysql_available

    def _parse_fetched_at(self, timestamp_raw):
        """优先使用输入 timestamp，解析失败时回退到当前时间。"""
        if not timestamp_raw:
            return _datetime.datetime.now()

        if isinstance(timestamp_raw, str):
            try:
                return _datetime.datetime.fromisoformat(timestamp_raw)
            except Exception:
                return _datetime.datetime.now()

        if hasattr(timestamp_raw, "isoformat"):
            return timestamp_raw

        return _datetime.datetime.now()

    def is_available(self):
        """判断 Redis 或 MySQL 至少一个可用。"""
        return self.check_redis() or self.check_mysql()

    def initialize(self, redis_client=None, mysql_engine=None):
        """启动时调用，可选注入 Redis/MySQL 客户端并检测各后端可用性。"""
        if redis_client is not None:
            self.redis_client = redis_client
        if mysql_engine is not None:
            self.mysql_engine = mysql_engine

        # 重置检测缓存，确保注入新客户端后重新探测后端可用性
        self.last_redis_check = 0
        self.last_mysql_check = 0

        redis_ok = self.check_redis()
        mysql_ok = self.check_mysql()

        if redis_ok:
            logger.info("Redis 可用")
        else:
            logger.warning("Redis 不可用")

        if mysql_ok:
            logger.info("MySQL 可用")
        else:
            logger.warning("MySQL 不可用")

    def _discover_periods(self, symbol, data_type):
        """从 Redis 或 MySQL 发现该品种已缓存的所有周期。"""
        if self.check_redis():
            try:
                pattern = f"market_data:{symbol}:*"
                keys = self.redis_client.keys(pattern)
                periods = set()
                for key in keys:
                    key_str = key.decode("utf-8") if isinstance(key, bytes) else key
                    parts = key_str.split(":")
                    if len(parts) == 3:
                        periods.add(parts[2])
                if periods:
                    return sorted(periods)
            except Exception as e:
                logger.warning(f"Redis 发现周期失败 [{symbol}]: {e}")

        if not self.check_mysql():
            return []

        try:
            SessionLocal = sessionmaker(bind=self.mysql_engine)
            with SessionLocal() as db:
                rows = (
                    db.query(MarketData.period)
                    .filter_by(symbol=symbol, data_type=data_type)
                    .distinct()
                    .all()
                )
                return sorted([r.period for r in rows if r.period])
        except Exception as e:
            logger.warning(f"MySQL 发现周期失败 [{symbol}]: {e}")
            return []

    def _get_market_data_single_period(self, symbol, data_type, period):
        """读取单个周期的行情数据：Redis → MySQL → None。"""
        cache_key = f"market_data:{symbol}:{period}"

        # 1. 检查 Redis 缓存
        if self.check_redis():
            try:
                raw = self.redis_client.get(cache_key)
                if raw:
                    return json.loads(raw)
            except Exception as e:
                logger.warning(f"Redis 读取失败 [{cache_key}]: {e}")

        # 2. Redis 未命中，检查 MySQL 可用性
        if not self.check_mysql():
            return None

        # 3. 从 MySQL 读取
        try:
            SessionLocal = sessionmaker(bind=self.mysql_engine)
            with SessionLocal() as db:
                record = (
                    db.query(MarketData)
                    .filter_by(symbol=symbol, data_type=data_type, period=period)
                    .first()
                )
                if not record:
                    return None

                result = {
                    "current_price": record.current_price,
                    "timestamp": (
                        record.fetched_at.isoformat() if record.fetched_at else None
                    ),
                    "candles": (
                        json.loads(record.candles_json) if record.candles_json else []
                    ),
                }

                # 4. 回填 Redis（TTL 30 天）
                if self.check_redis():
                    try:
                        self.redis_client.setex(
                            cache_key,
                            CACHE_TTL_SECONDS,
                            json.dumps(result, ensure_ascii=False),
                        )
                    except Exception as e:
                        logger.warning(f"Redis 回填失败 [{cache_key}]: {e}")

                return result
        except Exception as e:
            logger.warning(f"MySQL 读取失败 [{cache_key}]: {e}")
            return None

    @staticmethod
    def _parse_candle_time(candle):
        """解析单根 K 线的时间字段，返回 naive datetime 或 None。"""
        candle_time = candle.get("datetime") or candle.get("time")
        if not candle_time:
            return None

        if isinstance(candle_time, str):
            try:
                dt = datetime.fromisoformat(candle_time.replace("Z", "+00:00"))
                if dt.tzinfo is not None:
                    dt = dt.replace(tzinfo=None)
                return dt
            except Exception:
                return None

        if hasattr(candle_time, "tzinfo") and candle_time.tzinfo is not None:
            return candle_time.replace(tzinfo=None)
        return candle_time

    def _filter_and_truncate(self, merged, end_time, max_candles):
        """在合并结果上按 end_time 与 max_candles 过滤/截断每个周期。"""
        if merged is None:
            return None

        filter_end_time = end_time if end_time else datetime.now()
        if filter_end_time.tzinfo is not None:
            filter_end_time = filter_end_time.replace(tzinfo=None)

        timeframes = {}
        for period, candles in merged.get("timeframes", {}).items():
            filtered = []
            for candle in candles:
                candle_dt = self._parse_candle_time(candle)
                if candle_dt is None or candle_dt <= filter_end_time:
                    filtered.append(candle)

            if len(filtered) > max_candles:
                filtered = filtered[-max_candles:]

            timeframes[period] = filtered

        merged["timeframes"] = timeframes
        return merged

    def _merge_period_results(self, symbol, data_type, period_results):
        """将按周期返回的结果合并为统一格式。"""
        timeframes = {}
        current_price = None
        newest = None

        for period, result in period_results.items():
            if not result:
                continue
            candles = result.get("candles", [])
            timeframes[period] = candles
            if current_price is None:
                current_price = result.get("current_price")
            timestamp_raw = result.get("timestamp")
            if timestamp_raw:
                try:
                    timestamp_dt = datetime.fromisoformat(timestamp_raw)
                    if newest is None or timestamp_dt > newest:
                        newest = timestamp_dt
                except Exception:
                    pass

        if not timeframes:
            return None

        if newest is None:
            newest = datetime.now()

        now = datetime.now()
        is_fresh = (now - newest).total_seconds() < CACHE_TTL_SECONDS

        return {
            "symbol": symbol,
            "type": data_type,
            "current_price": current_price,
            "timestamp": newest.isoformat(),
            "timeframes": timeframes,
            "is_fresh": is_fresh,
            "fetched_at": newest.isoformat(),
        }

    def get_market_data(self, symbol, data_type, periods=None, end_time=None, max_candles=100):
        """
        读取行情数据：Redis → MySQL → None（由调用方降级到 SQLite）。

        Args:
            symbol: 品种合约代码
            data_type: 数据类型，如 futures
            periods: 周期或周期列表，如 "5min" 或 ["5min", "15min"];
                     为 None 时自动发现全部周期
            end_time: 结束时间（可选），过滤该时间之前的 K 线
            max_candles: 每个周期最大 K 线数量，默认 100

        Returns:
            periods 为字符串时返回旧版单周期 dict;
            periods 为列表/None 时返回合并后的 dict（包含 timeframes）或 None
        """
        single_period = isinstance(periods, str)
        query_periods = [periods] if single_period else periods

        if query_periods is None:
            query_periods = self._discover_periods(symbol, data_type)
            if not query_periods:
                return None

        period_results = {}
        for period in query_periods:
            period_results[period] = self._get_market_data_single_period(
                symbol, data_type, period
            )

        merged = self._merge_period_results(symbol, data_type, period_results)
        if merged is None:
            return None

        merged = self._filter_and_truncate(merged, end_time, max_candles)

        if single_period:
            period = query_periods[0]
            return {
                "current_price": merged.get("current_price"),
                "timestamp": merged.get("timestamp"),
                "candles": merged["timeframes"].get(period, []),
            }

        return merged

    def get_symbol_timestamp(self, symbol, data_type):
        """
        读取合约时间戳：Redis → MySQL → None（由调用方降级到 SQLite）。

        Args:
            symbol: 品种合约代码
            data_type: 数据类型，如 futures

        Returns:
            dict 或 None
        """
        cache_key = f"symbol_timestamps:{symbol}"

        # 1. 检查 Redis 缓存
        if self.check_redis():
            try:
                raw = self.redis_client.get(cache_key)
                if raw:
                    return json.loads(raw)
            except Exception as e:
                logger.warning(f"Redis 读取失败 [{cache_key}]: {e}")

        # 2. Redis 未命中，检查 MySQL 可用性
        if not self.check_mysql():
            return None

        # 3. 从 MySQL 读取
        try:
            SessionLocal = sessionmaker(bind=self.mysql_engine)
            with SessionLocal() as db:
                record = (
                    db.query(SymbolTimestamp)
                    .filter_by(symbol=symbol, data_type=data_type)
                    .first()
                )
                if not record:
                    return None

                result = {
                    "last_refresh_at": (
                        record.last_refresh_at.isoformat()
                        if record.last_refresh_at
                        else None
                    ),
                    "refresh_count": record.refresh_count,
                }

                # 4. 回填 Redis（TTL 30 天）
                if self.check_redis():
                    try:
                        self.redis_client.setex(
                            cache_key,
                            CACHE_TTL_SECONDS,
                            json.dumps(result, ensure_ascii=False),
                        )
                    except Exception as e:
                        logger.warning(f"Redis 回填失败 [{cache_key}]: {e}")

                return result
        except Exception as e:
            logger.warning(f"MySQL 读取失败 [{cache_key}]: {e}")
            return None

    def save_market_data(self, symbol, data):
        """
        保存行情数据：先删 Redis 缓存，再写 MySQL，MySQL 成功后回填 Redis。

        Args:
            symbol: 品种合约代码
            data: 采集脚本返回的数据，包含 timeframes, current_price, timestamp, type

        Returns:
            bool: MySQL 写入是否成功
        """
        data_type = data.get("type", "futures")
        timeframes = data.get("timeframes", {})
        current_price = data.get("current_price")
        fetched_at = self._parse_fetched_at(data.get("timestamp"))

        cache_keys = [f"market_data:{symbol}:{period}" for period in timeframes.keys()]

        # 1. 删除 Redis 缓存
        if self.check_redis() and cache_keys:
            try:
                self.redis_client.delete(*cache_keys)
            except Exception as e:
                logger.warning(f"Redis 删除缓存失败 [{symbol}]: {e}")

        # 2. 写入 MySQL（事务）
        if not self.check_mysql():
            logger.warning(f"MySQL 不可用，无法写入行情数据 [{symbol}]")
            return False

        try:
            SessionLocal = sessionmaker(bind=self.mysql_engine)
            with SessionLocal() as db:
                try:
                    for period, candles in timeframes.items():
                        candles_json = json.dumps(candles, ensure_ascii=False)
                        record = (
                            db.query(MarketData)
                            .filter_by(symbol=symbol, data_type=data_type, period=period)
                            .first()
                        )
                        if record:
                            record.candles_json = candles_json
                            record.current_price = current_price
                            record.fetched_at = fetched_at
                            record.candle_count = len(candles)
                        else:
                            record = MarketData(
                                symbol=symbol,
                                data_type=data_type,
                                period=period,
                                candles_json=candles_json,
                                current_price=current_price,
                                fetched_at=fetched_at,
                                candle_count=len(candles),
                            )
                            db.add(record)
                    db.commit()
                except Exception:
                    db.rollback()
                    raise
        except Exception as e:
            logger.error(f"MySQL 写入行情数据失败 [{symbol}]: {e}")
            return False

        # 3. MySQL 成功，回填 Redis（失败仅记录日志）
        if self.check_redis():
            for period, candles in timeframes.items():
                cache_key = f"market_data:{symbol}:{period}"
                cache_value = {
                    "current_price": current_price,
                    "timestamp": fetched_at.isoformat(),
                    "candles": candles,
                }
                try:
                    self.redis_client.setex(
                        cache_key,
                        CACHE_TTL_SECONDS,
                        json.dumps(cache_value, ensure_ascii=False),
                    )
                except Exception as e:
                    logger.warning(f"Redis 更新失败 [{cache_key}]: {e}")

        return True

    def save_symbol_timestamp(self, symbol, data_type, refresh_time):
        """
        保存合约时间戳：先删 Redis 缓存，再写 MySQL，MySQL 成功后回填 Redis。

        Args:
            symbol: 品种合约代码
            data_type: 数据类型，如 futures
            refresh_time: 刷新时间（datetime）

        Returns:
            bool: MySQL 写入是否成功
        """
        cache_key = f"symbol_timestamps:{symbol}"

        # 1. 删除 Redis 缓存
        if self.check_redis():
            try:
                self.redis_client.delete(cache_key)
            except Exception as e:
                logger.warning(f"Redis 删除缓存失败 [{cache_key}]: {e}")

        # 2. 写入 MySQL（事务）
        if not self.check_mysql():
            logger.warning(f"MySQL 不可用，无法写入合约时间戳 [{symbol}]")
            return False

        try:
            SessionLocal = sessionmaker(bind=self.mysql_engine)
            with SessionLocal() as db:
                try:
                    record = (
                        db.query(SymbolTimestamp)
                        .filter_by(symbol=symbol, data_type=data_type)
                        .first()
                    )
                    if record:
                        record.last_refresh_at = refresh_time
                        record.refresh_count += 1
                    else:
                        record = SymbolTimestamp(
                            symbol=symbol,
                            data_type=data_type,
                            last_refresh_at=refresh_time,
                            refresh_count=1,
                        )
                        db.add(record)
                    db.commit()
                except Exception:
                    db.rollback()
                    raise
        except Exception as e:
            logger.error(f"MySQL 写入合约时间戳失败 [{symbol}]: {e}")
            return False

        # 3. MySQL 成功，回填 Redis（失败仅记录日志）
        if self.check_redis():
            cache_value = {
                "last_refresh_at": refresh_time.isoformat(),
                "refresh_count": record.refresh_count,
            }
            try:
                self.redis_client.setex(
                    cache_key,
                    CACHE_TTL_SECONDS,
                    json.dumps(cache_value, ensure_ascii=False),
                )
            except Exception as e:
                logger.warning(f"Redis 更新失败 [{cache_key}]: {e}")

        return True

    def delete_cache(self, symbol, periods):
        """
        批量删除 Redis 行情数据缓存键。

        Args:
            symbol: 品种合约代码
            periods: 周期列表，如 ["5min", "15min"]
        """
        if not self.check_redis():
            return

        cache_keys = [f"market_data:{symbol}:{period}" for period in periods]
        if not cache_keys:
            return

        try:
            self.redis_client.delete(*cache_keys)
        except Exception as e:
            logger.warning(f"Redis 批量删除缓存失败 [{symbol}]: {e}")


_storage_manager = None


def get_storage_manager():
    """获取全局 StorageManager 单例。"""
    global _storage_manager
    if _storage_manager is None:
        _storage_manager = StorageManager()
    return _storage_manager