# -*- coding: utf-8 -*-
"""
东方财富财务数据采集器 V2.0--每次季度财报更新之后执行这个脚本
适配2025年新版接口

从东方财富网自动采集A股上市公司的财务报表数据，包括：
1. 资产负债表
2. 利润表  
3. 现金流量表

数据存储到MongoDB数据库中，保留所有原始字段
"""

import requests
import pymongo
import logging
import datetime
import time
import json
import pandas as pd
from typing import List, Dict, Optional
import sys
import os
from sqlalchemy import create_engine, text

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

# 导入配置
try:
    from valuation_analysis.config import MONGO_CONFIG2, DB_URL
except ImportError:
    # 如果上面的导入失败，尝试直接导入
    import importlib.util
    config_path = os.path.join(project_root, 'valuation_analysis', 'config.py')
    spec = importlib.util.spec_from_file_location("config", config_path)
    config_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(config_module)
    MONGO_CONFIG2 = config_module.MONGO_CONFIG2
    DB_URL = config_module.DB_URL

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('financial_data_collector_v2.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 为StreamHandler也设置编码（如果可能）
try:
    if hasattr(sys.stdout, 'reconfigure'):
        sys.stdout.reconfigure(encoding='utf-8')
    if hasattr(sys.stderr, 'reconfigure'):
        sys.stderr.reconfigure(encoding='utf-8')
except:
    pass  # 如果设置失败就忽略


class FinancialDataCollectorV2:
    """财务数据采集器 V2.0 - 适配新版东方财富接口"""
    
    def __init__(self):
        """初始化"""
        self.mongo_client = None
        self.db = None
        # 使用新的集合名称存储新版本数据
        self.collection_name = 'eastmoney_financial_data_v2'
        self.collection = None
        
        # 初始化MySQL连接
        self.mysql_engine = None
        
        self.connect_mongodb()
        self.connect_mysql()
        
    def connect_mongodb(self):
        """连接MongoDB数据库"""
        try:
            # 使用参数形式连接MongoDB，而不是连接字符串
            self.mongo_client = pymongo.MongoClient(
                host=MONGO_CONFIG2['host'],
                port=MONGO_CONFIG2['port'],
                username=MONGO_CONFIG2['username'],
                password=MONGO_CONFIG2['password']
            )
            self.db = self.mongo_client[MONGO_CONFIG2['db']]
            self.collection = self.db[self.collection_name]
            
            # 测试连接
            self.mongo_client.admin.command('ping')
            logger.info(f"MongoDB连接成功，使用集合: {self.collection_name}")
            
        except Exception as e:
            logger.error(f"MongoDB连接失败: {str(e)}")
            raise
    
    def connect_mysql(self):
        """连接MySQL数据库"""
        try:
            self.mysql_engine = create_engine(
                DB_URL,
                pool_size=5,
                max_overflow=10,
                pool_recycle=3600
            )
            
            # 测试连接
            with self.mysql_engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            
            logger.info("MySQL数据库连接成功")
            
        except Exception as e:
            logger.error(f"MySQL数据库连接失败: {str(e)}")
            raise
    
    def get_all_stock_codes(self) -> List[str]:
        """
        从数据库中获取所有股票代码
        
        Returns:
            List[str]: 股票代码列表
        """
        try:
            query = "SELECT DISTINCT gp_code_two FROM gp_code_all WHERE gp_code_two IS NOT NULL AND gp_code_two != ''"
            
            with self.mysql_engine.connect() as conn:
                df = pd.read_sql(text(query), conn)
            
            stock_codes = df['gp_code_two'].tolist()
            logger.info(f"从数据库获取到 {len(stock_codes)} 只股票")
            
            return stock_codes
            
        except Exception as e:
            logger.error(f"获取股票代码列表失败: {str(e)}")
            return []
    
    def build_date_filter(self, stock_code: str, periods: int = 21) -> str:
        """
        构建日期过滤条件 - 适配新版接口
        
        Args:
            stock_code: 股票代码，如'300750.SZ'
            periods: 获取多少个报告期，默认21个季度
            
        Returns:
            str: 日期过滤字符串
        """
        # 计算需要的日期范围，从当前日期往前推
        current_date = datetime.datetime.now()
        dates = []
        
        # 生成最近21个季度的日期
        for i in range(periods):
            # 从当前季度往前推
            year = current_date.year
            month = current_date.month
            
            # 计算季度
            if month <= 3:
                quarter_month = 3
                year_offset = i // 4
                quarter_offset = i % 4
                if quarter_offset > 0:
                    year_offset += 1
                    quarter_month = [12, 9, 6, 3][quarter_offset - 1]
                target_year = year - year_offset
            elif month <= 6:
                quarter_month = 6
                year_offset = i // 4
                quarter_offset = i % 4
                if quarter_offset >= 1:
                    if quarter_offset == 1:
                        quarter_month = 3
                    elif quarter_offset == 2:
                        quarter_month = 12
                        target_year = year - 1
                    elif quarter_offset == 3:
                        quarter_month = 9
                        target_year = year - 1
                else:
                    target_year = year
                if quarter_offset >= 2:
                    year_offset += 1
                target_year = year - year_offset
            elif month <= 9:
                quarter_month = 9
                year_offset = i // 4
                quarter_offset = i % 4
                if quarter_offset >= 1:
                    if quarter_offset == 1:
                        quarter_month = 6
                    elif quarter_offset == 2:
                        quarter_month = 3
                    elif quarter_offset == 3:
                        quarter_month = 12
                        target_year = year - 1
                else:
                    target_year = year
                if quarter_offset >= 3:
                    year_offset += 1
                target_year = year - year_offset
            else:
                quarter_month = 12
                year_offset = i // 4
                quarter_offset = i % 4
                if quarter_offset >= 1:
                    if quarter_offset == 1:
                        quarter_month = 9
                    elif quarter_offset == 2:
                        quarter_month = 6
                    elif quarter_offset == 3:
                        quarter_month = 3
                else:
                    target_year = year
                target_year = year - year_offset
            
            # 简化逻辑：直接从2025年开始往前推21个季度
            base_year = 2025
            base_quarters = [
                (2025, 9), (2025, 6), (2025, 3), (2024, 12), (2024, 9), (2024, 6), (2024, 3),
                (2023, 12), (2023, 9), (2023, 6), (2023, 3),
                (2022, 12), (2022, 9), (2022, 6), (2022, 3),
                (2021, 12), (2021, 9), (2021, 6), (2021, 3),
                (2020, 12), (2020, 9)
            ]
            
            if i < len(base_quarters):
                year, month = base_quarters[i]
                if month == 3:
                    date_str = f'{year}-03-31'
                elif month == 6:
                    date_str = f'{year}-06-30'
                elif month == 9:
                    date_str = f'{year}-09-30'
                else:
                    date_str = f'{year}-12-31'
                dates.append(date_str)
        
        # 构建过滤字符串
        date_filter = f'(SECUCODE%3D%22{stock_code}%22)(REPORT_DATE%20in%20('
        date_filter += '%2C'.join([f'%27{date}%27' for date in dates])
        date_filter += '))'
        
        return date_filter
    
    def fetch_profit_statement(self, stock_code: str, periods: int = 21) -> List[Dict]:
        """获取利润表数据"""
        date_filter = self.build_date_filter(stock_code, periods)
        url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_GINCOME&sty=APP_F10_GINCOME&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC'
        
        headers = {"Content-Type": "application/json"}
        
        try:
            response = requests.get(url, headers=headers, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            if 'result' in data and 'data' in data['result']:
                logger.info(f"成功获取利润表数据，共 {len(data['result']['data'])} 个报告期")
                return data['result']['data']
            else:
                logger.warning("利润表数据格式异常")
                return []
                
        except Exception as e:
            logger.error(f"获取利润表失败: {str(e)}")
            return []
    
    def fetch_balance_sheet(self, stock_code: str, periods: int = 21) -> List[Dict]:
        """获取资产负债表数据 - 支持API自动切换 (G→B→S→I)"""
        date_filter = self.build_date_filter(stock_code, periods)
        
        # 先尝试G系列API
        g_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_GBALANCE&sty=F10_FINANCE_GBALANCE&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=012481899342117453'
        
        headers = {"Content-Type": "application/json"}
        
        try:
            response = requests.get(g_url, headers=headers, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            # 检查G系列API是否返回有效数据
            if data.get('code') == 9201 or data.get('success') == False or not data.get('result') or not data.get('result', {}).get('data'):
                logger.info(f"G系列API返回空数据，尝试B系列API - {stock_code}")
                
                # 切换到B系列API
                b_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_BBALANCE&sty=F10_FINANCE_BBALANCE&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=08691713756184818'
                
                response = requests.get(b_url, headers=headers, timeout=30)
                response.raise_for_status()
                data = response.json()
                
                # 检查B系列API是否返回有效数据
                if data.get('code') == 9201 or data.get('success') == False or not data.get('result') or not data.get('result', {}).get('data'):
                    logger.info(f"B系列API也返回空数据，尝试S系列API - {stock_code}")
                    
                    # 切换到S系列API (证券公司)
                    s_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_SBALANCE&sty=F10_FINANCE_SBALANCE&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=08691713756184818'
                    
                    response = requests.get(s_url, headers=headers, timeout=30)
                    response.raise_for_status()
                    data = response.json()
                    
                    # 检查S系列API是否返回有效数据
                    if data.get('code') == 9201 or data.get('success') == False or not data.get('result') or not data.get('result', {}).get('data'):
                        logger.info(f"S系列API也返回空数据，尝试I系列API - {stock_code}")
                        
                        # 切换到I系列API
                        i_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_IBALANCE&sty=F10_FINANCE_IBALANCE&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=08691713756184818'
                        
                        response = requests.get(i_url, headers=headers, timeout=30)
                        response.raise_for_status()
                        data = response.json()
                        
                        if 'result' in data and data['result'] and 'data' in data['result']:
                            logger.info(f"I系列API成功获取资产负债表数据，共 {len(data['result']['data'])} 个报告期")
                            return data['result']['data']
                        else:
                            logger.warning(f"I系列API也无法获取资产负债表数据 - {stock_code}")
                            return []
                    else:
                        # S系列API成功
                        if 'result' in data and data['result'] and 'data' in data['result']:
                            logger.info(f"S系列API成功获取资产负债表数据，共 {len(data['result']['data'])} 个报告期")
                            return data['result']['data']
                        else:
                            logger.warning(f"S系列API资产负债表数据格式异常 - {stock_code}")
                            return []
                else:
                    # B系列API成功
                    if 'result' in data and data['result'] and 'data' in data['result']:
                        logger.info(f"B系列API成功获取资产负债表数据，共 {len(data['result']['data'])} 个报告期")
                        return data['result']['data']
                    else:
                        logger.warning(f"B系列API资产负债表数据格式异常 - {stock_code}")
                        return []
            else:
                # G系列API成功
                if 'result' in data and 'data' in data['result']:
                    logger.info(f"G系列API成功获取资产负债表数据，共 {len(data['result']['data'])} 个报告期")
                    return data['result']['data']
                else:
                    logger.warning("G系列API资产负债表数据格式异常")
                    return []
                
        except Exception as e:
            logger.error(f"获取资产负债表失败: {str(e)}")
            return []
    
    def fetch_cash_flow_statement(self, stock_code: str, periods: int = 21) -> List[Dict]:
        """获取现金流量表数据 - 支持API自动切换 (G→B→S→I)"""
        date_filter = self.build_date_filter(stock_code, periods)
        
        # 先尝试G系列API
        g_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_GCASHFLOW&sty=APP_F10_GCASHFLOW&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=04664977872701077'
        
        headers = {"Content-Type": "application/json"}
        
        try:
            response = requests.get(g_url, headers=headers, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            # 检查G系列API是否返回有效数据
            if data.get('code') == 9201 or data.get('success') == False or not data.get('result') or not data.get('result', {}).get('data'):
                logger.info(f"G系列API返回空数据，尝试B系列API - {stock_code}")
                
                # 切换到B系列API
                b_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_BCASHFLOW&sty=APP_F10_BCASHFLOW&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=03391901368547232'
                
                response = requests.get(b_url, headers=headers, timeout=30)
                response.raise_for_status()
                data = response.json()
                
                # 检查B系列API是否返回有效数据
                if data.get('code') == 9201 or data.get('success') == False or not data.get('result') or not data.get('result', {}).get('data'):
                    logger.info(f"B系列API也返回空数据，尝试S系列API - {stock_code}")
                    
                    # 切换到S系列API (证券公司)
                    s_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_SCASHFLOW&sty=APP_F10_SCASHFLOW&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=03391901368547232'
                    
                    response = requests.get(s_url, headers=headers, timeout=30)
                    response.raise_for_status()
                    data = response.json()
                    
                    # 检查S系列API是否返回有效数据
                    if data.get('code') == 9201 or data.get('success') == False or not data.get('result') or not data.get('result', {}).get('data'):
                        logger.info(f"S系列API也返回空数据，尝试I系列API - {stock_code}")
                        
                        # 切换到I系列API
                        i_url = f'https://datacenter.eastmoney.com/securities/api/data/get?type=RPT_F10_FINANCE_ICASHFLOW&sty=APP_F10_ICASHFLOW&filter={date_filter}&p=1&ps={periods}&sr=-1&st=REPORT_DATE&source=HSF10&client=PC&v=03391901368547232'
                        
                        response = requests.get(i_url, headers=headers, timeout=30)
                        response.raise_for_status()
                        data = response.json()
                        
                        if 'result' in data and data['result'] and 'data' in data['result']:
                            logger.info(f"I系列API成功获取现金流量表数据，共 {len(data['result']['data'])} 个报告期")
                            return data['result']['data']
                        else:
                            logger.warning(f"I系列API也无法获取现金流量表数据 - {stock_code}")
                            return []
                    else:
                        # S系列API成功
                        if 'result' in data and data['result'] and 'data' in data['result']:
                            logger.info(f"S系列API成功获取现金流量表数据，共 {len(data['result']['data'])} 个报告期")
                            return data['result']['data']
                        else:
                            logger.warning(f"S系列API现金流量表数据格式异常 - {stock_code}")
                            return []
                else:
                    # B系列API成功
                    if 'result' in data and data['result'] and 'data' in data['result']:
                        logger.info(f"B系列API成功获取现金流量表数据，共 {len(data['result']['data'])} 个报告期")
                        return data['result']['data']
                    else:
                        logger.warning(f"B系列API现金流量表数据格式异常 - {stock_code}")
                        return []
            else:
                # G系列API成功
                if 'result' in data and 'data' in data['result']:
                    logger.info(f"G系列API成功获取现金流量表数据，共 {len(data['result']['data'])} 个报告期")
                    return data['result']['data']
                else:
                    logger.warning("G系列API现金流量表数据格式异常")
                    return []
                
        except Exception as e:
            logger.error(f"获取现金流量表失败: {str(e)}")
            return []
    
    def process_financial_data(self, stock_code: str, profit_data: List[Dict], 
                             balance_data: List[Dict], cash_data: List[Dict]) -> List[Dict]:
        """
        处理财务数据，按报告期合并三张表的数据
        
        Args:
            stock_code: 股票代码
            profit_data: 利润表数据
            balance_data: 资产负债表数据
            cash_data: 现金流量表数据
            
        Returns:
            List[Dict]: 处理后的财务数据列表
        """
        financial_data_list = []
        
        # 创建按报告日期索引的字典
        profit_dict = {item['REPORT_DATE']: item for item in profit_data}
        balance_dict = {item['REPORT_DATE']: item for item in balance_data}
        cash_dict = {item['REPORT_DATE']: item for item in cash_data}
        
        # 获取所有报告日期
        all_dates = set()
        all_dates.update(profit_dict.keys())
        all_dates.update(balance_dict.keys())
        all_dates.update(cash_dict.keys())
        
        # 按日期排序
        sorted_dates = sorted(all_dates, reverse=True)
        
        for report_date in sorted_dates:
            try:
                # 构建完整的财务数据记录
                financial_record = {
                    'stock_code': stock_code,
                    'report_date': report_date[:10] if report_date else None,  # 只取日期部分
                    'collect_time': datetime.datetime.now(),
                    'data_source': 'eastmoney_v2',
                    'profit_statement': profit_dict.get(report_date, {}),
                    'balance_sheet': balance_dict.get(report_date, {}),
                    'cash_flow_statement': cash_dict.get(report_date, {})
                }
                
                # 只有当至少有一张表有数据时才添加记录
                if (financial_record['profit_statement'] or 
                    financial_record['balance_sheet'] or 
                    financial_record['cash_flow_statement']):
                    financial_data_list.append(financial_record)
                
            except Exception as e:
                logger.error(f"处理报告期 {report_date} 数据时出错: {str(e)}")
                continue
        
        logger.info(f"成功处理 {len(financial_data_list)} 个报告期的数据")
        return financial_data_list
    
    def save_to_mongodb(self, financial_data_list: List[Dict]) -> bool:
        """
        保存数据到MongoDB，如果数据已存在则跳过，只新增不存在的数据
        
        Args:
            financial_data_list: 财务数据列表
            
        Returns:
            bool: 是否保存成功
        """
        try:
            if not financial_data_list:
                logger.warning("没有数据需要保存")
                return False
            
            # 统计插入和跳过的数量
            inserted_count = 0
            skipped_count = 0
            
            for data in financial_data_list:
                # 使用股票代码和报告日期作为唯一标识
                filter_condition = {
                    'stock_code': data['stock_code'],
                    'report_date': data['report_date']
                }
                
                # 检查数据是否已存在
                existing_record = self.collection.find_one(filter_condition)
                
                if existing_record:
                    # 数据已存在，跳过
                    skipped_count += 1
                    logger.debug(f"跳过已存在的数据: {data['stock_code']} - {data['report_date']}")
                else:
                    # 数据不存在，插入新数据
                    self.collection.insert_one(data)
                    inserted_count += 1
                    logger.debug(f"插入新数据: {data['stock_code']} - {data['report_date']}")
            
            logger.info(f"数据保存完成 - 新增: {inserted_count} 条，跳过: {skipped_count} 条")
            return True
            
        except Exception as e:
            logger.error(f"保存数据到MongoDB失败: {str(e)}")
            return False

    def collect_financial_data(self, stock_code: str, periods: int = 21) -> bool:
        """
        采集单只股票的财务数据
        
        Args:
            stock_code: 股票代码，如'300750.SZ'
            periods: 获取多少个报告期，默认21个季度
            
        Returns:
            bool: 是否采集成功
        """
        try:
            logger.info(f"开始采集股票 {stock_code} 的财务数据（{periods}个报告期）")

            # 获取三张财务报表数据
            profit_data = self.fetch_profit_statement(stock_code, periods)
            time.sleep(1)  # 避免请求过于频繁

            balance_data = self.fetch_balance_sheet(stock_code, periods)
            time.sleep(1)

            cash_data = self.fetch_cash_flow_statement(stock_code, periods)
            time.sleep(1)
            
            # 检查至少有一张表有数据
            if not any([profit_data, balance_data, cash_data]):
                logger.error(f"股票 {stock_code} 没有获取到任何财务数据")
                return False
            
            # 处理财务数据
            financial_data_list = self.process_financial_data(
                stock_code, profit_data, balance_data, cash_data
            )
            
            if not financial_data_list:
                logger.error(f"股票 {stock_code} 的财务数据处理失败")
                return False

            # 保存到MongoDB
            success = self.save_to_mongodb(financial_data_list)
            
            if success:
                logger.info(f"股票 {stock_code} 的财务数据采集完成")
            
            return success
            
        except Exception as e:
            logger.error(f"采集股票 {stock_code} 的财务数据失败: {str(e)}")
            return False
    
    def batch_collect_financial_data(self, stock_codes: List[str], periods: int = 21) -> Dict:
        """
        批量采集多只股票的财务数据
        
        Args:
            stock_codes: 股票代码列表
            periods: 获取多少个报告期，默认21个季度
            
        Returns:
            Dict: 采集结果统计
        """
        results = {'success': 0, 'failed': 0, 'failed_stocks': []}
        total_stocks = len(stock_codes)
        
        logger.info(f"开始批量采集 {total_stocks} 只股票的财务数据")
        
        for index, stock_code in enumerate(stock_codes, 1):
            try:
                # 显示进度
                progress = (index / total_stocks) * 100
                logger.info(f"进度 [{index}/{total_stocks}] ({progress:.1f}%) - 正在处理: {stock_code}")
                
                success = self.collect_financial_data(stock_code, periods)
                if success:
                    results['success'] += 1
                    logger.info(f"SUCCESS [{index}/{total_stocks}] {stock_code} 采集成功")
                else:
                    results['failed'] += 1
                    results['failed_stocks'].append(stock_code)
                    logger.warning(f"FAILED [{index}/{total_stocks}] {stock_code} 采集失败")
                
                # 每只股票之间暂停一下，避免请求过于频繁
                time.sleep(2)
                
                # 每100只股票输出一次统计信息
                if index % 100 == 0:
                    current_success_rate = (results['success'] / index) * 100
                    logger.info(f"阶段性统计: 已处理 {index}/{total_stocks} 只股票，成功率: {current_success_rate:.2f}%")
                
            except Exception as e:
                logger.error(f"处理股票 {stock_code} 时出错: {str(e)}")
                results['failed'] += 1
                results['failed_stocks'].append(stock_code)
                
                # 继续处理下一只股票，不中断整个流程
                continue
        
        success_rate = (results['success'] / total_stocks) * 100
        logger.info(f"批量采集完成: 成功{results['success']}只，失败{results['failed']}只，成功率: {success_rate:.2f}%")
        
        if results['failed_stocks']:
            logger.info(f"失败的股票数量: {len(results['failed_stocks'])}")
            # 记录前10个失败的股票到日志
            failed_sample = results['failed_stocks'][:10]
            logger.info(f"失败股票示例: {', '.join(failed_sample)}")
        
        return results
    
    def close_connection(self):
        """关闭数据库连接"""
        if self.mongo_client:
            self.mongo_client.close()
            logger.info("MongoDB连接已关闭")
        
        if self.mysql_engine:
            self.mysql_engine.dispose()
            logger.info("MySQL连接已关闭")


def main():
    """主函数 - 批量采集所有股票的财务数据"""
    collector = FinancialDataCollectorV2()
    
    try:
        # 从数据库获取所有股票代码
        logger.info("正在从数据库获取股票列表...")
        stock_codes = collector.get_all_stock_codes()
        
        if not stock_codes:
            logger.error("未获取到任何股票代码，程序退出")
            return
        
        logger.info(f"从数据库获取到 {len(stock_codes)} 只股票")
        
        # 可以选择采集所有股票或者部分股票进行测试
        # 如果要测试，可以取前几只股票
        # 测试模式：只采集前10只股票
        TEST_MODE = False  # 设置为False将采集所有股票
        
        if TEST_MODE:
            test_count = min(10, len(stock_codes))  # 最多取10只股票测试
            stock_codes = stock_codes[:test_count]
            logger.info(f"TEST MODE: 仅采集前 {test_count} 只股票")
        else:
            logger.info(f"PRODUCTION MODE: 将采集全部 {len(stock_codes)} 只股票")
        
        logger.info(f"开始批量采集 {len(stock_codes)} 只股票的财务数据")
        
        # 批量采集
        results = collector.batch_collect_financial_data(stock_codes, periods=21)
        
        # 输出最终结果
        print(f"\n{'='*50}")
        print(f"批量采集完成统计")
        print(f"{'='*50}")
        print(f"SUCCESS 成功采集: {results['success']} 只股票")
        print(f"FAILED 采集失败: {results['failed']} 只股票")
        print(f"SUCCESS RATE 成功率: {(results['success'] / len(stock_codes) * 100):.2f}%")
        
        if results['failed_stocks']:
            print(f"\n失败的股票列表:")
            for i, stock in enumerate(results['failed_stocks'][:10], 1):  # 只显示前10个
                print(f"  {i}. {stock}")
            if len(results['failed_stocks']) > 10:
                print(f"  ... 还有 {len(results['failed_stocks']) - 10} 只股票")
        
        logger.info("所有任务已完成")
        
    except KeyboardInterrupt:
        logger.info("用户中断程序执行")
        print("\n警告: 程序被用户中断")
    except Exception as e:
        logger.error(f"采集过程中出现错误: {str(e)}")
        print(f"\n错误: 程序执行出错: {str(e)}")
    finally:
        collector.close_connection()


if __name__ == "__main__":
    main()
