"""
恐贪指数(Fear & Greed Index)数据管理模块

提供恐贪指数的数据库操作功能，包括：
1. 创建恐贪指数数据表
2. 新增恐贪指数数据
3. 查询恐贪指数数据
"""

import pandas as pd
from sqlalchemy import create_engine, text
import datetime
import logging
from typing import Dict, List, Optional, Union

from .config import DB_URL, LOG_FILE

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(LOG_FILE),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("fear_greed_index")


class FearGreedIndexManager:
    """恐贪指数数据管理器类"""
    
    def __init__(self, db_url: str = DB_URL):
        """
        初始化恐贪指数数据管理器
        
        Args:
            db_url: 数据库连接URL
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        # 确保数据表存在
        self._ensure_table_exists()
        logger.info("恐贪指数数据管理器初始化完成")
    
    def _ensure_table_exists(self) -> bool:
        """
        确保恐贪指数数据表存在，如果不存在则创建
        
        Returns:
            是否成功确保表存在
        """
        try:
            # 创建恐贪指数表的SQL语句
            create_table_sql = """
            CREATE TABLE IF NOT EXISTS `fear_greed_index` (
              `id` int(11) NOT NULL AUTO_INCREMENT,
              `index_value` decimal(5,2) NOT NULL COMMENT '恐贪指数值(0-100)',
              `trading_date` date NOT NULL COMMENT '交易日期',
              `update_time` datetime NOT NULL COMMENT '更新时间',
              PRIMARY KEY (`id`),
              UNIQUE KEY `uk_trading_date` (`trading_date`)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='市场恐贪指数数据';
            """
            
            # 执行建表语句
            with self.engine.connect() as conn:
                conn.execute(text(create_table_sql))
                conn.commit()
                
            logger.info("恐贪指数数据表创建成功或已存在")
            return True
            
        except Exception as e:
            logger.error(f"确保恐贪指数数据表存在失败: {e}")
            return False
    
    def add_index_data(self, index_value: float, trading_date: str) -> bool:
        """
        添加恐贪指数数据
        
        Args:
            index_value: 恐贪指数值，范围0-100，保留两位小数
            trading_date: 交易日期，格式为 YYYY-MM-DD
            
        Returns:
            是否成功添加数据
        """
        try:
            # 验证指数值范围
            if not 0 <= index_value <= 100:
                logger.error(f"恐贪指数值 {index_value} 超出范围(0-100)")
                return False
            
            # 验证日期格式
            try:
                formatted_date = datetime.datetime.strptime(trading_date, "%Y-%m-%d").date()
            except ValueError:
                logger.error(f"无效的日期格式: {trading_date}，应为YYYY-MM-DD格式")
                return False
            
            # 当前时间作为更新时间
            update_time = datetime.datetime.now()
            
            # 构建插入SQL
            insert_sql = """
            INSERT INTO fear_greed_index (index_value, trading_date, update_time)
            VALUES (:index_value, :trading_date, :update_time)
            ON DUPLICATE KEY UPDATE
            index_value = :index_value,
            update_time = :update_time
            """
            
            # 执行插入操作
            with self.engine.connect() as conn:
                conn.execute(
                    text(insert_sql),
                    {
                        "index_value": round(index_value, 2),
                        "trading_date": trading_date,
                        "update_time": update_time
                    }
                )
                conn.commit()
                
            logger.info(f"恐贪指数数据添加成功: 日期={trading_date}, 值={index_value}")
            return True
            
        except Exception as e:
            logger.error(f"添加恐贪指数数据失败: {e}")
            return False
    
    def get_index_data(self, start_date: str = None, end_date: str = None, limit: int = 730) -> Dict:
        """
        获取恐贪指数数据
        
        Args:
            start_date: 开始日期，格式为YYYY-MM-DD
            end_date: 结束日期，格式为YYYY-MM-DD
            limit: 限制返回的记录数量，默认为730条（约两年的交易日数据）
            
        Returns:
            包含恐贪指数数据的字典
        """
        try:
            # 构建查询条件
            conditions = []
            params = {}
            
            if start_date:
                conditions.append("trading_date >= :start_date")
                params["start_date"] = start_date
                
            if end_date:
                conditions.append("trading_date <= :end_date")
                params["end_date"] = end_date
                
            # 构建查询SQL
            where_clause = " AND ".join(conditions) if conditions else "1=1"
            query = f"""
            SELECT id, index_value, trading_date, update_time
            FROM fear_greed_index
            WHERE {where_clause}
            ORDER BY trading_date DESC
            LIMIT :limit
            """
            
            params["limit"] = limit
            
            # 执行查询
            with self.engine.connect() as conn:
                result = conn.execute(text(query), params)
                rows = result.fetchall()
                
            if not rows:
                logger.warning("未找到恐贪指数数据")
                return {
                    "success": False,
                    "message": "未找到数据"
                }
            
            # 处理查询结果
            data = []
            for row in rows:
                data.append({
                    "id": row[0],
                    "index_value": float(row[1]),
                    "trading_date": row[2].strftime("%Y-%m-%d"),
                    "update_time": row[3].strftime("%Y-%m-%d %H:%M:%S")
                })
            
            # 按日期升序排序，方便生成图表
            data.reverse()
            
            # 提取日期和指数值列表
            dates = [item["trading_date"] for item in data]
            values = [item["index_value"] for item in data]
            
            # 最新数据是最大日期的数据
            data_by_date = sorted(data, key=lambda x: x["trading_date"], reverse=True)
            latest = data_by_date[0] if data_by_date else None
            
            # 计算恐贪指数状态
            def get_index_status(value):
                if 0 <= value < 25:
                    return "极度恐慌"
                elif 25 <= value < 40:
                    return "恐慌"
                elif 40 <= value < 50:
                    return "偏向恐慌"
                elif 50 <= value < 60:
                    return "中性"
                elif 60 <= value < 75:
                    return "偏向贪婪"
                elif 75 <= value < 90:
                    return "贪婪"
                else:
                    return "极度贪婪"
            
            # 计算最新状态
            latest_status = get_index_status(latest["index_value"]) if latest else None
            
            return {
                "success": True,
                "dates": dates,
                "values": values,
                "data": data,
                "latest": latest,
                "latest_status": latest_status,
                "update_time": latest["update_time"] if latest else datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
            
        except Exception as e:
            logger.error(f"获取恐贪指数数据失败: {e}")
            return {
                "success": False,
                "message": f"获取数据失败: {str(e)}"
            } 