"""
东方财富实时股价数据采集模块
提供从东方财富网站采集实时股价数据并存储到数据库的功能
功能包括：
1. 采集实时股价数据
2. 存储数据到数据库
3. 定时自动更新数据
"""

import requests
import pandas as pd
import datetime
import logging
import time
import os
import sys
from pathlib import Path
from sqlalchemy import create_engine, text
from typing import Dict

# 添加项目根目录到Python路径
current_file = Path(__file__)
project_root = current_file.parent.parent.parent
sys.path.append(str(project_root))

from src.valuation_analysis.config import DB_URL, LOG_FILE

# 获取项目根目录
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 确保日志目录存在
os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)

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


def get_create_table_sql() -> str:
    """
    获取创建实时股价数据表的SQL语句
    
    Returns:
        创建表的SQL语句
    """
    return """
    CREATE TABLE IF NOT EXISTS stock_price_data (
        stock_code VARCHAR(10) PRIMARY KEY COMMENT '股票代码',
        stock_name VARCHAR(50) COMMENT '股票名称',
        latest_price DECIMAL(10,2) COMMENT '最新价',
        change_percent DECIMAL(10,2) COMMENT '涨跌幅',
        change_amount DECIMAL(10,2) COMMENT '涨跌额',
        volume BIGINT COMMENT '成交量(手)',
        amount DECIMAL(20,2) COMMENT '成交额',
        amplitude DECIMAL(10,2) COMMENT '振幅',
        turnover_rate DECIMAL(10,2) COMMENT '换手率',
        pe_ratio DECIMAL(10,2) COMMENT '市盈率',
        high_price DECIMAL(10,2) COMMENT '最高价',
        low_price DECIMAL(10,2) COMMENT '最低价',
        open_price DECIMAL(10,2) COMMENT '开盘价',
        pre_close DECIMAL(10,2) COMMENT '昨收价',
        total_market_value DECIMAL(20,2) COMMENT '总市值',
        float_market_value DECIMAL(20,2) COMMENT '流通市值',
        pb_ratio DECIMAL(10,2) COMMENT '市净率',
        list_date DATE COMMENT '上市日期',
        update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间'
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='实时股价数据表';
    """


class StockPriceCollector:
    """东方财富实时股价数据采集器类"""
    
    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.base_url = "https://push2.eastmoney.com/api/qt/clist/get"
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
            "Accept": "application/json, text/plain, */*",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
            "Origin": "https://quote.eastmoney.com",
            "Referer": "https://quote.eastmoney.com/",
        }
        logger.info("东方财富实时股价数据采集器初始化完成")
    
    def _ensure_table_exists(self) -> bool:
        """
        确保数据表存在，如果不存在则创建
        
        Returns:
            是否成功确保表存在
        """
        try:
            create_table_query = text(get_create_table_sql())
            
            with self.engine.connect() as conn:
                conn.execute(create_table_query)
                conn.commit()
                
            logger.info("实时股价数据表创建成功")
            return True
            
        except Exception as e:
            logger.error(f"确保数据表存在失败: {e}")
            return False
    
    def _convert_stock_code(self, code: str) -> str:
        """
        转换股票代码格式
        
        Args:
            code: 原始股票代码
            
        Returns:
            转换后的股票代码
        """
        if code.startswith(('0', '3')):
            return f"{code}.SZ"
        else:
            return f"{code}.SH"
    
    def _parse_list_date(self, date_str: str) -> datetime.date:
        """
        解析上市日期
        
        Args:
            date_str: 日期字符串
            
        Returns:
            日期对象
        """
        if not date_str or date_str == '-':
            return None
        try:
            # 如果输入是整数，先转换为字符串
            if isinstance(date_str, int):
                date_str = str(date_str)
            return datetime.datetime.strptime(date_str, "%Y%m%d").date()
        except ValueError:
            logger.warning(f"无法解析日期: {date_str}")
            return None
    
    def fetch_data(self, page: int = 1) -> pd.DataFrame:
        """
        获取指定页码的实时股价数据
        
        Args:
            page: 页码
            
        Returns:
            包含实时股价数据的DataFrame
        """
        try:
            params = {
                "np": 1,
                "fltt": 2,
                "invt": 2,
                "fs": "m:0+t:6,m:0+t:80,m:1+t:2,m:1+t:23,m:0+t:81+s:2048",
                "fid": "f12",
                "pn": page,
                "pz": 100,
                "po": 0,
                "dect": 1
            }
            
            logger.info(f"开始获取第 {page} 页数据")
            
            response = requests.get(self.base_url, params=params, headers=self.headers)
            if response.status_code != 200:
                logger.error(f"获取第 {page} 页数据失败: HTTP {response.status_code}")
                return pd.DataFrame()
            
            data = response.json()
            if not data.get("rc") == 0:
                logger.error(f"获取数据失败: {data.get('message', '未知错误')}")
                return pd.DataFrame()
            
            # 提取数据列表
            items = data.get("data", {}).get("diff", [])
            if not items:
                logger.warning(f"第 {page} 页未找到有效数据")
                return pd.DataFrame()
            
            # 转换为DataFrame
            df = pd.DataFrame(items)
            
            # 重命名列
            column_mapping = {
                "f12": "stock_code",
                "f14": "stock_name",
                "f2": "latest_price",
                "f3": "change_percent",
                "f4": "change_amount",
                "f5": "volume",
                "f6": "amount",
                "f7": "amplitude",
                "f8": "turnover_rate",
                "f9": "pe_ratio",
                "f15": "high_price",
                "f16": "low_price",
                "f17": "open_price",
                "f18": "pre_close",
                "f20": "total_market_value",
                "f21": "float_market_value",
                "f23": "pb_ratio",
                "f26": "list_date"
            }
            
            df = df.rename(columns=column_mapping)
            
            # 转换股票代码格式
            df['stock_code'] = df['stock_code'].apply(self._convert_stock_code)
            
            # 转换上市日期
            df['list_date'] = df['list_date'].apply(self._parse_list_date)

            logger.info(f"第 {page} 页数据获取成功，包含 {len(df)} 条记录")
            return df
            
        except Exception as e:
            logger.error(f"获取第 {page} 页数据失败: {e}")
            return pd.DataFrame()
    
    def fetch_all_data(self) -> pd.DataFrame:
        """
        获取所有页的实时股价数据
        
        Returns:
            包含所有实时股价数据的DataFrame
        """
        all_data = []
        page = 1
        
        while True:
            page_data = self.fetch_data(page)
            if page_data.empty:
                logger.info(f"第 {page} 页数据为空，停止采集")
                break
                
            all_data.append(page_data)
            
            # 如果返回的数据少于100条，说明是最后一页
            if len(page_data) < 100:
                break
                
            page += 1
            # 添加延迟，避免请求过于频繁
            time.sleep(1)
        
        if all_data:
            combined_df = pd.concat(all_data, ignore_index=True)
            logger.info(f"数据采集完成，共采集 {len(combined_df)} 条记录")
            return combined_df
        else:
            logger.warning("未获取到任何有效数据")
            return pd.DataFrame()
    
    def save_to_database(self, data: pd.DataFrame) -> bool:
        """
        将数据保存到数据库
        
        Args:
            data: 要保存的数据DataFrame
            
        Returns:
            是否成功保存数据
        """
        if data.empty:
            logger.warning("没有数据需要保存")
            return False
        
        try:
            # 确保数据表存在
            if not self._ensure_table_exists():
                return False
            data = data.replace('-', None)
            # 将nan值转换为None（在SQL中会变成NULL）
            data = data.replace({pd.NA: None, pd.NaT: None})
            data = data.where(pd.notnull(data), None)
            
            # 添加数据或更新已有数据
            inserted_count = 0
            updated_count = 0
            
            with self.engine.connect() as conn:
                for _, row in data.iterrows():
                    # 将Series转换为dict，并处理nan值
                    row_dict = {k: (None if pd.isna(v) else v) for k, v in row.items()}
                    
                    # 检查该股票的数据是否已存在
                    check_query = text("""
                    SELECT COUNT(*) FROM stock_price_data WHERE stock_code = :stock_code
                    """)
                    result = conn.execute(check_query, {"stock_code": row_dict['stock_code']}).scalar()
                    
                    if result > 0:  # 数据已存在，执行更新
                        update_query = text("""
                        UPDATE stock_price_data SET
                            stock_name = :stock_name,
                            latest_price = :latest_price,
                            change_percent = :change_percent,
                            change_amount = :change_amount,
                            volume = :volume,
                            amount = :amount,
                            amplitude = :amplitude,
                            turnover_rate = :turnover_rate,
                            pe_ratio = :pe_ratio,
                            high_price = :high_price,
                            low_price = :low_price,
                            open_price = :open_price,
                            pre_close = :pre_close,
                            total_market_value = :total_market_value,
                            float_market_value = :float_market_value,
                            pb_ratio = :pb_ratio,
                            list_date = :list_date
                        WHERE stock_code = :stock_code
                        """)
                        conn.execute(update_query, row_dict)
                        updated_count += 1
                    else:  # 数据不存在，执行插入
                        insert_query = text("""
                        INSERT INTO stock_price_data (
                            stock_code, stock_name, latest_price, change_percent,
                            change_amount, volume, amount, amplitude, turnover_rate,
                            pe_ratio, high_price, low_price, open_price, pre_close,
                            total_market_value, float_market_value, pb_ratio, list_date
                        ) VALUES (
                            :stock_code, :stock_name, :latest_price, :change_percent,
                            :change_amount, :volume, :amount, :amplitude, :turnover_rate,
                            :pe_ratio, :high_price, :low_price, :open_price, :pre_close,
                            :total_market_value, :float_market_value, :pb_ratio, :list_date
                        )
                        """)
                        conn.execute(insert_query, row_dict)
                        inserted_count += 1
                
                conn.commit()
            
            logger.info(f"数据保存成功：新增 {inserted_count} 条记录，更新 {updated_count} 条记录")
            return True
            
        except Exception as e:
            logger.error(f"保存数据到数据库失败: {e}")
            return False
    
    def update_latest_data(self) -> bool:
        """
        更新最新实时股价数据
        
        Returns:
            是否成功更新最新数据
        """
        try:
            logger.info("开始更新最新实时股价数据")
            
            # 获取所有数据
            df = self.fetch_all_data()
            if df.empty:
                logger.warning("未获取到最新数据")
                return False
            
            # 保存数据到数据库
            result = self.save_to_database(df)
            
            if result:
                logger.info(f"最新数据更新成功，共更新 {len(df)} 条记录")
            else:
                logger.warning("最新数据更新失败")
                
            return result
            
        except Exception as e:
            logger.error(f"更新最新数据失败: {e}")
            return False

    def get_stock_price_data(self, stock_code: str, convert_code: bool = False) -> Dict:
        """
        获取指定股票的最新价格数据
        
        Args:
            stock_code: 股票代码
            convert_code: 是否需要转换股票代码格式，默认为False
            
        Returns:
            包含股票价格数据的字典
        """
        try:
            # 转换股票代码格式
            formatted_code = self._convert_stock_code(stock_code) if convert_code else stock_code
            
            query = text("""
                SELECT 
                    stock_code,
                    stock_name,
                    latest_price,
                    change_percent,
                    change_amount,
                    volume,
                    amount,
                    amplitude,
                    turnover_rate,
                    pe_ratio,
                    high_price,
                    low_price,
                    total_market_value,
                    float_market_value,
                    pb_ratio,
                    list_date,
                    update_time
                FROM 
                    stock_price_data 
                WHERE 
                    stock_code = :stock_code
            """)
            
            with self.engine.connect() as conn:
                result = conn.execute(query, {"stock_code": formatted_code}).fetchone()
                
            if result:
                # 将结果转换为字典
                data = dict(result._mapping)
                # 处理日期类型
                if data['list_date']:
                    data['list_date'] = data['list_date'].strftime('%Y-%m-%d')
                if data['update_time']:
                    data['update_time'] = data['update_time'].strftime('%Y-%m-%d %H:%M:%S')
                return data
            else:
                logger.warning(f"未找到股票 {stock_code} 的价格数据")
                return None
                
        except Exception as e:
            logger.error(f"获取股票价格数据失败: {e}")
            return None


# 示例使用方式
if __name__ == "__main__":
    # 创建实时股价数据采集器
    collector = StockPriceCollector()
    
    # 更新最新数据
    logger.info("开始更新最新实时股价数据...")
    collector.update_latest_data() 