"""
东方财富美股实时股价数据采集模块
提供从东方财富网站采集美股实时股价数据并存储到数据库的功能
功能包括：
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("us_stock_price_collector")


def get_create_us_stock_table_sql() -> str:
    """
    获取创建美股实时股价数据表的SQL语句
    
    Returns:
        创建表的SQL语句
    """
    return """
    CREATE TABLE IF NOT EXISTS us_stock_price_data (
        stock_code VARCHAR(20) PRIMARY KEY COMMENT '美股代码',
        stock_name VARCHAR(100) 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 USStockPriceCollector:
    """东方财富美股实时股价数据采集器类"""
    
    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,en;q=0.8:zh-CN,zh;q=0.9,en;q=0.8",
                            "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 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36:Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36"
                        }
        logger.info("东方财富美股实时股价数据采集器初始化完成")
    
    def _ensure_table_exists(self) -> bool:
        """
        确保数据表存在，如果不存在则创建
        
        Returns:
            是否成功确保表存在
        """
        try:
            create_table_query = text(get_create_us_stock_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_us_stock_code(self, code: str) -> str:
        """
        转换美股代码格式，保持原始格式
        
        Args:
            code: 原始股票代码
            
        Returns:
            转换后的股票代码
        """
        # 美股代码通常直接返回，不需要添加后缀
        return code.upper()
    
    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:
            # 美股数据参数，根据你提供的示例URL调整
            params = {
                "np": 1,
                "fltt": 1,  # 美股使用1
                "invt": 2,
                "fs": "m:105,m:106,m:107",  # 美股市场标识
                "fid": "f12",
                "pn": page,
                "pz": 100,
                "po": 1,  # 美股使用1
                "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_us_stock_code)
            
            # 转换上市日期
            df['list_date'] = df['list_date'].apply(self._parse_list_date)
            
            # 转换价格字段：东方财富API返回的价格需要除以1000转换为美元
            price_columns = ['latest_price', 'change_amount', 'high_price', 'low_price', 'open_price', 'pre_close']
            for col in price_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce') / 1000.0
            
            # 转换市值字段：市值字段也可能需要除以100
            market_value_columns = ['total_market_value', 'float_market_value']
            for col in market_value_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce') / 1000.0

            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)
            
            # 保存数据到数据库
            data.to_sql('us_stock_price_data', self.engine, if_exists='replace', index=False)
            
            logger.info(f"成功保存 {len(data)} 条美股数据到数据库")
            return True
            
        except Exception as e:
            logger.error(f"保存美股数据到数据库失败: {e}")
            return False
    
    def run_collection(self) -> bool:
        """
        运行完整的美股数据采集流程
        
        Returns:
            是否成功完成采集
        """
        try:
            logger.info("开始美股数据采集流程")
            
            # 获取所有数据
            all_data = self.fetch_all_data()
            if all_data.empty:
                logger.warning("未获取到任何美股数据")
                return False
            
            # 保存到数据库
            if self.save_to_database(all_data):
                logger.info("美股数据采集流程完成")
                return True
            else:
                logger.error("美股数据保存失败")
                return False
                
        except Exception as e:
            logger.error(f"美股数据采集流程失败: {e}")
            return False


def main():
    """主函数，用于测试美股数据采集"""
    collector = USStockPriceCollector()
    success = collector.run_collection()
    
    if success:
        print("美股数据采集成功完成")
    else:
        print("美股数据采集失败")


if __name__ == "__main__":
    main()


