"""
东方财富实时股价数据采集模块
提供从东方财富网站采集实时股价数据并存储到数据库的功能
功能包括：
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, Optional, Union
from decimal import Decimal

# 添加项目根目录到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

# 条件导入代理管理器
proxy_manager = None
try:
    sys.path.append(os.path.dirname(project_root))
    from src.scripts.ProxyIP import EnhancedProxyManager
    proxy_manager = EnhancedProxyManager()
except ImportError:
    logger = logging.getLogger("stock_price_collector")
    logger.warning("代理管理器导入失败，将使用直接请求模式")
    proxy_manager = None

# 获取项目根目录
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 = {
                            "accept": "*/*:*/*",
                            "accept-encoding": "gzip, deflate, br, zstd:gzip, deflate, br, zstd",
                            "accept-language": "zh-CN,zh;q=0.9",
                            "connection": "keep-alive:keep-alive",
                            "cookie": "qgqp_b_id=3d1376c8b3609b85ddd613ff70797b21; st_nvi=OkigGRrjz_ZnTtL_7eE4-4595; nid=0820fd3f15fe4f7c0164f8a7eff85c5f; nid_create_time=1754027308372; gvi=p9WxDMFjnfkCUiL8CQ3ex015a; gvi_create_time=1754027308372; fullscreengg=1; fullscreengg2=1; st_si=11531712238902; st_asi=delete; st_pvi=49494569999152; st_sp=2025-05-12%2009%3A11%3A41; st_inirUrl=https%3A%2F%2Fwww.baidu.com%2Flink; st_sn=11; st_psi=20250904172534993-113200301321-0920486271:qgqp_b_id=3d1376c8b3609b85ddd613ff70797b21; st_nvi=OkigGRrjz_ZnTtL_7eE4-4595; nid=0820fd3f15fe4f7c0164f8a7eff85c5f; nid_create_time=1754027308372; gvi=p9WxDMFjnfkCUiL8CQ3ex015a; gvi_create_time=1754027308372; fullscreengg=1; fullscreengg2=1; st_si=11531712238902; st_asi=delete; st_pvi=49494569999152; st_sp=2025-05-12%2009%3A11%3A41; st_inirUrl=https%3A%2F%2Fwww.baidu.com%2Flink; st_sn=11; st_psi=20250904172534993-113200301321-0920486271",
                            "host": "push2.eastmoney.com:push2.eastmoney.com",
                            "referer": "https://quote.eastmoney.com/center/gridlist.html",
                            "sec-ch-ua": "\"Chromium\";v=\"134\", \"Not:A-Brand\";v=\"24\", \"Google Chrome\";v=\"134\"",
                            "sec-ch-ua-mobile": "?0:?0",
                            "sec-ch-ua-platform": "\"Windows\":\"Windows\"",
                            "sec-fetch-dest": "script:script",
                            "sec-fetch-mode": "no-cors:no-cors",
                            "sec-fetch-site": "same-site:same-site",
                            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
                        }
        logger.info("东方财富实时股价数据采集器初始化完成")

    @staticmethod
    def _safe_float(value: Union[str, int, float, Decimal]) -> Optional[float]:
        """安全地将值转换为浮点数"""
        if value is None:
            return None
        try:
            if isinstance(value, Decimal):
                return float(value)
            if isinstance(value, (int, float)):
                return float(value)
            if isinstance(value, str):
                stripped = value.strip()
                if stripped in {"", "-"}:
                    return None
                return float(stripped)
        except (ValueError, TypeError):
            return None
        return None

    @staticmethod
    def _safe_int(value: Union[str, int, float, Decimal]) -> Optional[int]:
        """安全地将值转换为整数"""
        if value is None:
            return None
        try:
            if isinstance(value, Decimal):
                return int(value)
            if isinstance(value, int):
                return value
            if isinstance(value, float):
                return int(value)
            if isinstance(value, str):
                stripped = value.strip()
                if stripped in {"", "-"}:
                    return None
                return int(float(stripped))
        except (ValueError, TypeError):
            return None
        return None

    @staticmethod
    def _format_decimal(value: Optional[float], precision: int = 2) -> Optional[str]:
        """按照指定精度格式化小数"""
        if value is None:
            return None
        return format(value, f".{precision}f")

    @staticmethod
    def _to_db_symbol(stock_code: str) -> Optional[str]:
        """
        转换股票代码为数据库格式（如 603986.SH -> SH603986）
        """
        if not stock_code:
            return None
        code = stock_code.upper().strip()
        if "." in code:
            number, market = code.split(".")
            return f"{market}{number}"
        if code.startswith(("SH", "SZ", "BJ")):
            return code
        if code.isdigit() and len(code) == 6:
            if code.startswith(("6", "9")):
                return f"SH{code}"
            if code.startswith(("0", "3")):
                return f"SZ{code}"
            if code.startswith("4"):
                return f"BJ{code}"
        return code
    
    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, use_proxy: bool = False) -> pd.DataFrame:
        """
        获取指定页码的实时股价数据
        
        Args:
            page: 页码
            use_proxy: 是否使用代理
            
        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} 页数据")
            
            # 根据配置选择是否使用代理
            if use_proxy and proxy_manager:
                response = proxy_manager.request_with_proxy('get', self.base_url, headers=self.headers, params=params)
            else:
                response = requests.get(self.base_url, params=params, headers=self.headers, timeout=10)
            
            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
            db_symbol = self._to_db_symbol(formatted_code)

            if not db_symbol:
                logger.warning(f"无法转换股票代码 {stock_code}")
                return None

            query = text("""
                SELECT 
                    d.symbol AS symbol,
                    d.timestamp AS day_timestamp,
                    d.open AS day_open,
                    d.high AS day_high,
                    d.low AS day_low,
                    d.close AS day_close,
                    d.chg AS day_chg,
                    d.percent AS day_percent,
                    d.volume AS day_volume,
                    d.amount AS day_amount,
                    d.pre_close AS day_pre_close,
                    d.turnoverrate AS day_turnoverrate,
                    b.timestamp AS basic_timestamp,
                    b.turnover_rate AS basic_turnover_rate,
                    b.pe_ttm AS basic_pe_ttm,
                    b.pb AS basic_pb,
                    b.total_mv AS basic_total_mv,
                    b.circ_mv AS basic_circ_mv
                FROM (
                    SELECT 
                        symbol,
                        timestamp,
                        open,
                        high,
                        low,
                        close,
                        chg,
                        percent,
                        volume,
                        amount,
                        pre_close,
                        turnoverrate
                    FROM gp_day_data
                    WHERE symbol = :symbol
                    ORDER BY timestamp DESC
                    LIMIT 1
                ) d
                JOIN (
                    SELECT 
                        symbol,
                        timestamp,
                        turnover_rate,
                        pe_ttm,
                        pb,
                        total_mv,
                        circ_mv
                    FROM gp_day_basic
                    WHERE symbol = :symbol
                    ORDER BY timestamp DESC
                    LIMIT 1
                ) b ON d.symbol = b.symbol
            """)

            with self.engine.connect() as conn:
                result = conn.execute(query, {"symbol": db_symbol}).fetchone()

            if not result:
                logger.warning(f"未找到股票 {stock_code} 的日线或基础数据")
                return None

            row = dict(result._mapping)

            close_price = self._safe_float(row.get("day_close"))
            change_amount = self._safe_float(row.get("day_chg"))
            change_percent = self._safe_float(row.get("day_percent"))
            high_price = self._safe_float(row.get("day_high"))
            low_price = self._safe_float(row.get("day_low"))
            pre_close = self._safe_float(row.get("day_pre_close"))
            turnover_rate_daily = self._safe_float(row.get("day_turnoverrate"))
            amount_value = self._safe_float(row.get("day_amount"))
            volume_value = self._safe_int(row.get("day_volume"))

            turnover_rate = self._safe_float(row.get("basic_turnover_rate")) or turnover_rate_daily
            pe_ratio = self._safe_float(row.get("basic_pe_ttm"))
            pb_ratio = self._safe_float(row.get("basic_pb"))

            total_mv = self._safe_float(row.get("basic_total_mv"))
            circ_mv = self._safe_float(row.get("basic_circ_mv"))

            amplitude = None
            if all(value is not None for value in (high_price, low_price, pre_close)) and pre_close != 0:
                amplitude = (high_price - low_price) / abs(pre_close) * 100

            update_timestamp = row.get("day_timestamp") or row.get("basic_timestamp")
            update_time_str = (
                update_timestamp.strftime('%Y-%m-%d %H:%M:%S')
                if update_timestamp else None
            )

            result = {
                'stock_code': formatted_code,
                'stock_name': None,
                'latest_price': self._format_decimal(close_price),
                'change_percent': self._format_decimal(change_percent),
                'change_amount': self._format_decimal(change_amount),
                'volume': volume_value,
                'amount': self._format_decimal(amount_value) if amount_value is not None else None,
                'amplitude': self._format_decimal(amplitude),
                'turnover_rate': self._format_decimal(turnover_rate),
                'pe_ratio': self._format_decimal(pe_ratio),
                'high_price': self._format_decimal(high_price),
                'low_price': self._format_decimal(low_price),
                'total_market_value': self._format_decimal(total_mv * 10000) if total_mv is not None else None,
                'float_market_value': self._format_decimal(circ_mv * 10000) if circ_mv is not None else None,
                'pb_ratio': self._format_decimal(pb_ratio),
                'list_date': None,
                'update_time': update_time_str
            }

            return result
        except Exception as e:
            logger.error(f"获取股票价格数据失败: {e}")
            return None


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