#!/usr/bin/env python3
"""
MongoDB + Redis 数据库缓存管理器
提供高性能的股票数据缓存和持久化存储
"""

import os
import json
import pickle
import hashlib
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List, Union
import pandas as pd

# MongoDB
try:
    from pymongo import MongoClient
    from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
    MONGODB_AVAILABLE = True
except ImportError:
    MONGODB_AVAILABLE = False
    print("⚠️ pymongo 未安装，MongoDB功能不可用")

# Redis
try:
    import redis
    from redis.exceptions import ConnectionError as RedisConnectionError
    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False
    print("⚠️ redis 未安装，Redis功能不可用")


class DatabaseCacheManager:
    """MongoDB + Redis 数据库缓存管理器"""
    
    def __init__(self,
                 mongodb_url: Optional[str] = None,
                 redis_url: Optional[str] = None,
                 mongodb_db: str = "tradingagents",
                 redis_db: int = 0):
        """
        初始化数据库缓存管理器

        Args:
            mongodb_url: MongoDB连接URL，默认使用配置文件端口
            redis_url: Redis连接URL，默认使用配置文件端口
            mongodb_db: MongoDB数据库名
            redis_db: Redis数据库编号
        """
        # 从配置文件获取正确的端口
        mongodb_port = os.getenv("MONGODB_PORT", "27018")
        redis_port = os.getenv("REDIS_PORT", "6380")
        mongodb_password = os.getenv("MONGODB_PASSWORD", "tradingagents123")
        redis_password = os.getenv("REDIS_PASSWORD", "tradingagents123")

        self.mongodb_url = mongodb_url or os.getenv("MONGODB_URL", f"mongodb://admin:{mongodb_password}@localhost:{mongodb_port}")
        self.redis_url = redis_url or os.getenv("REDIS_URL", f"redis://:{redis_password}@localhost:{redis_port}")
        self.mongodb_db_name = mongodb_db
        self.redis_db = redis_db
        
        # 初始化连接
        self.mongodb_client = None
        self.mongodb_db = None
        self.redis_client = None
        
        self._init_mongodb()
        self._init_redis()
        
        print(f"🗄️ 数据库缓存管理器初始化完成")
        print(f"   MongoDB: {'✅ 已连接' if self.mongodb_client else '❌ 未连接'}")
        print(f"   Redis: {'✅ 已连接' if self.redis_client else '❌ 未连接'}")
    
    def _init_mongodb(self):
        """初始化MongoDB连接"""
        if not MONGODB_AVAILABLE:
            return
        
        try:
            self.mongodb_client = MongoClient(
                self.mongodb_url,
                serverSelectionTimeoutMS=5000,  # 5秒超时
                connectTimeoutMS=5000
            )
            # 测试连接
            self.mongodb_client.admin.command('ping')
            self.mongodb_db = self.mongodb_client[self.mongodb_db_name]
            
            # 创建索引
            self._create_mongodb_indexes()
            
            print(f"✅ MongoDB连接成功: {self.mongodb_url}")
            
        except Exception as e:
            print(f"❌ MongoDB连接失败: {e}")
            self.mongodb_client = None
            self.mongodb_db = None
    
    def _init_redis(self):
        """初始化Redis连接"""
        if not REDIS_AVAILABLE:
            return
        
        try:
            self.redis_client = redis.from_url(
                self.redis_url,
                db=self.redis_db,
                socket_timeout=5,
                socket_connect_timeout=5,
                decode_responses=True
            )
            # 测试连接
            self.redis_client.ping()
            
            print(f"✅ Redis连接成功: {self.redis_url}")
            
        except Exception as e:
            print(f"❌ Redis连接失败: {e}")
            self.redis_client = None
    
    def _create_mongodb_indexes(self):
        """创建MongoDB索引"""
        if self.mongodb_db is None:
            return
        
        try:
            # 股票数据集合索引
            stock_collection = self.mongodb_db.stock_data
            stock_collection.create_index([
                ("symbol", 1),
                ("data_source", 1),
                ("start_date", 1),
                ("end_date", 1)
            ])
            stock_collection.create_index([("created_at", 1)])
            
            # 新闻数据集合索引
            news_collection = self.mongodb_db.news_data
            news_collection.create_index([
                ("symbol", 1),
                ("data_source", 1),
                ("date_range", 1)
            ])
            news_collection.create_index([("created_at", 1)])
            
            # 基本面数据集合索引
            fundamentals_collection = self.mongodb_db.fundamentals_data
            fundamentals_collection.create_index([
                ("symbol", 1),
                ("data_source", 1),
                ("analysis_date", 1)
            ])
            fundamentals_collection.create_index([("created_at", 1)])
            
            print("✅ MongoDB索引创建完成")
            
        except Exception as e:
            print(f"⚠️ MongoDB索引创建失败: {e}")
    
    def _generate_cache_key(self, data_type: str, symbol: str, **kwargs) -> str:
        """生成缓存键"""
        params_str = f"{data_type}_{symbol}"
        for key, value in sorted(kwargs.items()):
            params_str += f"_{key}_{value}"
        
        cache_key = hashlib.md5(params_str.encode()).hexdigest()[:16]
        return f"{data_type}:{symbol}:{cache_key}"
    
    def save_stock_data(self, symbol: str, data: Union[pd.DataFrame, str],
                       start_date: str = None, end_date: str = None,
                       data_source: str = "unknown", market_type: str = None) -> str:
        """
        保存股票数据到MongoDB和Redis
        
        Args:
            symbol: 股票代码
            data: 股票数据
            start_date: 开始日期
            end_date: 结束日期
            data_source: 数据源
            market_type: 市场类型 (us/china)
        
        Returns:
            cache_key: 缓存键
        """
        cache_key = self._generate_cache_key("stock", symbol,
                                           start_date=start_date,
                                           end_date=end_date,
                                           source=data_source)
        
        # 自动推断市场类型
        if market_type is None:
            # 根据股票代码格式推断市场类型
            import re
            if re.match(r'^\d{6}$', symbol):  # 6位数字为A股
                market_type = "china"
            else:  # 其他格式为美股
                market_type = "us"
        
        # 准备文档数据
        doc = {
            "_id": cache_key,
            "symbol": symbol,
            "market_type": market_type,
            "data_type": "stock_data",
            "start_date": start_date,
            "end_date": end_date,
            "data_source": data_source,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        }
        
        # 处理数据格式
        if isinstance(data, pd.DataFrame):
            doc["data"] = data.to_json(orient='records', date_format='iso')
            doc["data_format"] = "dataframe_json"
        else:
            doc["data"] = str(data)
            doc["data_format"] = "text"
        
        # 保存到MongoDB（持久化）
        if self.mongodb_db is not None:
            try:
                collection = self.mongodb_db.stock_data
                collection.replace_one({"_id": cache_key}, doc, upsert=True)
                print(f"💾 股票数据已保存到MongoDB: {symbol} -> {cache_key}")
            except Exception as e:
                print(f"⚠️ MongoDB保存失败: {e}")
        
        # 保存到Redis（快速缓存，6小时过期）
        if self.redis_client:
            try:
                redis_data = {
                    "data": doc["data"],
                    "data_format": doc["data_format"],
                    "symbol": symbol,
                    "data_source": data_source,
                    "created_at": doc["created_at"].isoformat()
                }
                self.redis_client.setex(
                    cache_key,
                    6 * 3600,  # 6小时过期
                    json.dumps(redis_data, ensure_ascii=False)
                )
                print(f"⚡ 股票数据已缓存到Redis: {symbol} -> {cache_key}")
            except Exception as e:
                print(f"⚠️ Redis缓存失败: {e}")
        
        return cache_key
    
    def load_stock_data(self, cache_key: str) -> Optional[Union[pd.DataFrame, str]]:
        """从Redis或MongoDB加载股票数据"""
        
        # 首先尝试从Redis加载（更快）
        if self.redis_client:
            try:
                redis_data = self.redis_client.get(cache_key)
                if redis_data:
                    data_dict = json.loads(redis_data)
                    print(f"⚡ 从Redis加载数据: {cache_key}")
                    
                    if data_dict["data_format"] == "dataframe_json":
                        return pd.read_json(data_dict["data"], orient='records')
                    else:
                        return data_dict["data"]
            except Exception as e:
                print(f"⚠️ Redis加载失败: {e}")
        
        # 如果Redis没有，从MongoDB加载
        if self.mongodb_db is not None:
            try:
                collection = self.mongodb_db.stock_data
                doc = collection.find_one({"_id": cache_key})
                
                if doc:
                    print(f"💾 从MongoDB加载数据: {cache_key}")
                    
                    # 同时更新到Redis缓存
                    if self.redis_client:
                        try:
                            redis_data = {
                                "data": doc["data"],
                                "data_format": doc["data_format"],
                                "symbol": doc["symbol"],
                                "data_source": doc["data_source"],
                                "created_at": doc["created_at"].isoformat()
                            }
                            self.redis_client.setex(
                                cache_key,
                                6 * 3600,
                                json.dumps(redis_data, ensure_ascii=False)
                            )
                            print(f"⚡ 数据已同步到Redis缓存")
                        except Exception as e:
                            print(f"⚠️ Redis同步失败: {e}")
                    
                    if doc["data_format"] == "dataframe_json":
                        return pd.read_json(doc["data"], orient='records')
                    else:
                        return doc["data"]
                        
            except Exception as e:
                print(f"⚠️ MongoDB加载失败: {e}")
        
        return None
    
    def find_cached_stock_data(self, symbol: str, start_date: str = None,
                              end_date: str = None, data_source: str = None,
                              max_age_hours: int = 6) -> Optional[str]:
        """查找匹配的缓存数据"""
        
        # 生成精确匹配的缓存键
        exact_key = self._generate_cache_key("stock", symbol,
                                           start_date=start_date,
                                           end_date=end_date,
                                           source=data_source)
        
        # 检查Redis中是否有精确匹配
        if self.redis_client and self.redis_client.exists(exact_key):
            print(f"⚡ Redis中找到精确匹配: {symbol} -> {exact_key}")
            return exact_key
        
        # 检查MongoDB中的匹配项
        if self.mongodb_db is not None:
            try:
                collection = self.mongodb_db.stock_data
                cutoff_time = datetime.utcnow() - timedelta(hours=max_age_hours)
                
                query = {
                    "symbol": symbol,
                    "created_at": {"$gte": cutoff_time}
                }
                
                if data_source:
                    query["data_source"] = data_source
                if start_date:
                    query["start_date"] = start_date
                if end_date:
                    query["end_date"] = end_date
                
                doc = collection.find_one(query, sort=[("created_at", -1)])
                
                if doc:
                    cache_key = doc["_id"]
                    print(f"💾 MongoDB中找到匹配: {symbol} -> {cache_key}")
                    return cache_key
                    
            except Exception as e:
                print(f"⚠️ MongoDB查询失败: {e}")
        
        print(f"❌ 未找到有效缓存: {symbol}")
        return None

    def save_news_data(self, symbol: str, news_data: str,
                      start_date: str = None, end_date: str = None,
                      data_source: str = "unknown") -> str:
        """保存新闻数据到MongoDB和Redis"""
        cache_key = self._generate_cache_key("news", symbol,
                                           start_date=start_date,
                                           end_date=end_date,
                                           source=data_source)

        doc = {
            "_id": cache_key,
            "symbol": symbol,
            "data_type": "news_data",
            "date_range": f"{start_date}_{end_date}",
            "start_date": start_date,
            "end_date": end_date,
            "data_source": data_source,
            "data": news_data,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        }

        # 保存到MongoDB
        if self.mongodb_db is not None:
            try:
                collection = self.mongodb_db.news_data
                collection.replace_one({"_id": cache_key}, doc, upsert=True)
                print(f"📰 新闻数据已保存到MongoDB: {symbol} -> {cache_key}")
            except Exception as e:
                print(f"⚠️ MongoDB保存失败: {e}")

        # 保存到Redis（24小时过期）
        if self.redis_client:
            try:
                redis_data = {
                    "data": news_data,
                    "symbol": symbol,
                    "data_source": data_source,
                    "created_at": doc["created_at"].isoformat()
                }
                self.redis_client.setex(
                    cache_key,
                    24 * 3600,  # 24小时过期
                    json.dumps(redis_data, ensure_ascii=False)
                )
                print(f"⚡ 新闻数据已缓存到Redis: {symbol} -> {cache_key}")
            except Exception as e:
                print(f"⚠️ Redis缓存失败: {e}")

        return cache_key

    def save_fundamentals_data(self, symbol: str, fundamentals_data: str,
                              analysis_date: str = None,
                              data_source: str = "unknown") -> str:
        """保存基本面数据到MongoDB和Redis"""
        if not analysis_date:
            analysis_date = datetime.now().strftime("%Y-%m-%d")

        cache_key = self._generate_cache_key("fundamentals", symbol,
                                           date=analysis_date,
                                           source=data_source)

        doc = {
            "_id": cache_key,
            "symbol": symbol,
            "data_type": "fundamentals_data",
            "analysis_date": analysis_date,
            "data_source": data_source,
            "data": fundamentals_data,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        }

        # 保存到MongoDB
        if self.mongodb_db is not None:
            try:
                collection = self.mongodb_db.fundamentals_data
                collection.replace_one({"_id": cache_key}, doc, upsert=True)
                print(f"💼 基本面数据已保存到MongoDB: {symbol} -> {cache_key}")
            except Exception as e:
                print(f"⚠️ MongoDB保存失败: {e}")

        # 保存到Redis（24小时过期）
        if self.redis_client:
            try:
                redis_data = {
                    "data": fundamentals_data,
                    "symbol": symbol,
                    "data_source": data_source,
                    "analysis_date": analysis_date,
                    "created_at": doc["created_at"].isoformat()
                }
                self.redis_client.setex(
                    cache_key,
                    24 * 3600,  # 24小时过期
                    json.dumps(redis_data, ensure_ascii=False)
                )
                print(f"⚡ 基本面数据已缓存到Redis: {symbol} -> {cache_key}")
            except Exception as e:
                print(f"⚠️ Redis缓存失败: {e}")

        return cache_key

    def get_cache_stats(self) -> Dict[str, Any]:
        """获取缓存统计信息"""
        stats = {
            "mongodb": {"available": self.mongodb_db is not None, "collections": {}},
            "redis": {"available": self.redis_client is not None, "keys": 0, "memory_usage": "N/A"}
        }

        # MongoDB统计
        if self.mongodb_db is not None:
            try:
                for collection_name in ["stock_data", "news_data", "fundamentals_data"]:
                    collection = self.mongodb_db[collection_name]
                    count = collection.count_documents({})
                    size = self.mongodb_db.command("collStats", collection_name).get("size", 0)
                    stats["mongodb"]["collections"][collection_name] = {
                        "count": count,
                        "size_mb": round(size / (1024 * 1024), 2)
                    }
            except Exception as e:
                print(f"⚠️ MongoDB统计获取失败: {e}")

        # Redis统计
        if self.redis_client:
            try:
                info = self.redis_client.info()
                stats["redis"]["keys"] = info.get("db0", {}).get("keys", 0)
                stats["redis"]["memory_usage"] = f"{info.get('used_memory_human', 'N/A')}"
            except Exception as e:
                print(f"⚠️ Redis统计获取失败: {e}")

        return stats

    def clear_old_cache(self, max_age_days: int = 7):
        """清理过期缓存"""
        cutoff_time = datetime.utcnow() - timedelta(days=max_age_days)
        cleared_count = 0

        # 清理MongoDB
        if self.mongodb_db is not None:
            try:
                for collection_name in ["stock_data", "news_data", "fundamentals_data"]:
                    collection = self.mongodb_db[collection_name]
                    result = collection.delete_many({"created_at": {"$lt": cutoff_time}})
                    cleared_count += result.deleted_count
                    print(f"🧹 MongoDB {collection_name} 清理了 {result.deleted_count} 条记录")
            except Exception as e:
                print(f"⚠️ MongoDB清理失败: {e}")

        # Redis会自动过期，不需要手动清理
        print(f"🧹 总共清理了 {cleared_count} 条过期记录")
        return cleared_count

    def close(self):
        """关闭数据库连接"""
        if self.mongodb_client:
            self.mongodb_client.close()
            print("🔒 MongoDB连接已关闭")

        if self.redis_client:
            self.redis_client.close()
            print("🔒 Redis连接已关闭")


# 全局数据库缓存实例
_db_cache_instance = None

def get_db_cache() -> DatabaseCacheManager:
    """获取全局数据库缓存实例"""
    global _db_cache_instance
    if _db_cache_instance is None:
        _db_cache_instance = DatabaseCacheManager()
    return _db_cache_instance
