#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
供应商客户占比数据更新器
用于将前五供应商占比和前五客户占比数据更新到MongoDB的eastmoney_financial_data_v2集合中
这个脚本要在财务脚本之后执行
"""

import os
import sys
import pandas as pd
import pymongo
import datetime
import logging
from typing import Dict, List, Optional
from pathlib import Path

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

# 导入配置
from src.valuation_analysis.config import MONGO_CONFIG2

# 设置日志
logging.basicConfig(
    level=logging.ERROR,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class SupplierCustomerRatioUpdater:
    """供应商客户占比数据更新器"""
    
    def __init__(self):
        """初始化"""
        self.mongo_client = None
        self.db = None
        self.collection_name = 'eastmoney_financial_data_v2'
        self.collection = None
        
        # CSV文件路径 - 使用绝对路径
        current_dir = Path(__file__).parent
        self.supplier_file = current_dir / 'data' / '前五供应商占比.csv'
        self.customer_file = current_dir / 'data' / '前五客户占比.csv'
        
        # 季度日期映射 - 根据CSV文件中的列名
        self.quarter_mapping = {
            '20200331': '2020-03-31',
            '20200630': '2020-06-30', 
            '20200930': '2020-09-30',
            '20201231': '2020-12-31',
            '20210331': '2021-03-31',
            '20210630': '2021-06-30',
            '20210930': '2021-09-30', 
            '20211231': '2021-12-31',
            '20220331': '2022-03-31',
            '20220630': '2022-06-30',
            '20220930': '2022-09-30',
            '20221231': '2022-12-31',
            '20230331': '2023-03-31',
            '20230630': '2023-06-30',
            '20230930': '2023-09-30',
            '20231231': '2023-12-31',
            '20240331': '2024-03-31',
            '20240630': '2024-06-30',
            '20240930': '2024-09-30',
            '20241231': '2024-12-31',
            '20250331': '2025-03-31'
        }
        
        self.connect_mongodb()
    
    def connect_mongodb(self):
        """连接MongoDB数据库"""
        try:
            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 load_csv_data(self, file_path: str) -> pd.DataFrame:
        """加载CSV文件数据"""
        try:
            if not os.path.exists(file_path):
                logger.error(f"文件不存在: {file_path}")
                return pd.DataFrame()
            
            df = pd.read_csv(file_path, encoding='utf-8')
            logger.info(f"成功加载文件: {file_path}, 数据行数: {len(df)}")
            return df
            
        except Exception as e:
            logger.error(f"加载CSV文件失败 {file_path}: {str(e)}")
            return pd.DataFrame()
    
    def parse_ratio_data(self, df: pd.DataFrame, data_type: str) -> Dict[str, Dict[str, float]]:
        """
        解析占比数据
        
        Args:
            df: CSV数据DataFrame
            data_type: 数据类型，'supplier' 或 'customer'
            
        Returns:
            Dict: {stock_code: {report_date: ratio_value}}
        """
        ratio_data = {}
        
        try:
            for _, row in df.iterrows():
                stock_code = row['stock_code']
                if pd.isna(stock_code):
                    continue
                
                ratio_data[stock_code] = {}
                
                # 遍历所有季度列
                for quarter_col, report_date in self.quarter_mapping.items():
                    if quarter_col in row:
                        ratio_value = row[quarter_col]
                        
                        # 处理数据：如果值为0或nan，跳过；否则除以100转换为小数
                        if pd.notna(ratio_value) and ratio_value > 0:
                            # 将百分比数值转换为小数（如56.22 -> 0.5622）
                            ratio_data[stock_code][report_date] = float(ratio_value) / 100.0
                        else:
                            # 0值或空值不添加到字典中
                            continue
            
            logger.info(f"解析{data_type}数据完成，共处理 {len(ratio_data)} 只股票")
            return ratio_data
            
        except Exception as e:
            logger.error(f"解析{data_type}数据失败: {str(e)}")
            return {}
    
    def update_mongodb_record(self, stock_code: str, report_date: str, 
                            supplier_ratio: Optional[float], customer_ratio: Optional[float]) -> bool:
        """
        更新MongoDB中的单条记录
        
        Args:
            stock_code: 股票代码
            report_date: 报告日期
            supplier_ratio: 前五供应商占比
            customer_ratio: 前五客户占比
            
        Returns:
            bool: 是否更新成功
        """
        try:
            # 查找记录
            filter_condition = {
                'stock_code': stock_code,
                'report_date': report_date
            }
            
            existing_record = self.collection.find_one(filter_condition)
            
            if not existing_record:
                # 记录不存在，跳过
                logger.debug(f"记录不存在，跳过: {stock_code} - {report_date}")
                return False
            
            # 准备更新字段
            update_fields = {}
            
            if supplier_ratio is not None:
                update_fields['top_five_suppliers_ratio'] = supplier_ratio
            
            if customer_ratio is not None:
                update_fields['top_five_customers_ratio'] = customer_ratio
            
            if not update_fields:
                return False
            
            # 添加更新时间
            update_fields['ratio_update_time'] = datetime.datetime.now()
            
            # 执行更新
            result = self.collection.update_one(
                filter_condition,
                {'$set': update_fields}
            )
            
            if result.modified_count > 0:
                logger.debug(f"更新成功: {stock_code} - {report_date}")
                return True
            else:
                logger.debug(f"无需更新: {stock_code} - {report_date}")
                return False
                
        except Exception as e:
            logger.error(f"更新记录失败 {stock_code} - {report_date}: {str(e)}")
            return False
    
    def batch_update_ratios(self):
        """批量更新所有占比数据"""
        try:
            # 加载CSV数据
            logger.info("开始加载CSV文件...")
            supplier_df = self.load_csv_data(self.supplier_file)
            customer_df = self.load_csv_data(self.customer_file)
            
            if supplier_df.empty and customer_df.empty:
                logger.error("没有可用的CSV数据")
                return False
            
            # 解析数据
            logger.info("解析数据中...")
            supplier_data = self.parse_ratio_data(supplier_df, 'supplier') if not supplier_df.empty else {}
            customer_data = self.parse_ratio_data(customer_df, 'customer') if not customer_df.empty else {}
            
            # 获取所有涉及的股票代码
            all_stock_codes = set()
            all_stock_codes.update(supplier_data.keys())
            all_stock_codes.update(customer_data.keys())
            
            logger.info(f"开始更新数据，共涉及 {len(all_stock_codes)} 只股票")
            
            # 统计
            total_count = 0
            updated_count = 0
            
            # 逐个股票更新
            for stock_code in all_stock_codes:
                supplier_ratios = supplier_data.get(stock_code, {})
                customer_ratios = customer_data.get(stock_code, {})
                
                # 获取所有日期
                all_dates = set()
                all_dates.update(supplier_ratios.keys())
                all_dates.update(customer_ratios.keys())
                
                # 逐个日期更新
                for report_date in all_dates:
                    supplier_ratio = supplier_ratios.get(report_date)
                    customer_ratio = customer_ratios.get(report_date)
                    
                    total_count += 1
                    
                    if self.update_mongodb_record(stock_code, report_date, supplier_ratio, customer_ratio):
                        updated_count += 1
                    
                    # 每处理100条记录输出一次进度
                    if total_count % 100 == 0:
                        logger.info(f"已处理 {total_count} 条记录，更新 {updated_count} 条")
            
            logger.info(f"数据更新完成！总计处理 {total_count} 条记录，成功更新 {updated_count} 条")
            return True
            
        except Exception as e:
            logger.error(f"批量更新失败: {str(e)}")
            return False
    
    def close_connection(self):
        """关闭数据库连接"""
        try:
            if self.mongo_client:
                self.mongo_client.close()
                logger.info("MongoDB连接已关闭")
        except Exception as e:
            logger.error(f"关闭MongoDB连接失败: {str(e)}")


def main():
    """主函数"""
    logger.info("=== 供应商客户占比数据更新器启动 ===")
    
    updater = None
    try:
        # 创建更新器实例
        updater = SupplierCustomerRatioUpdater()
        
        # 执行批量更新
        success = updater.batch_update_ratios()
        
        if success:
            logger.info("=== 数据更新成功完成 ===")
        else:
            logger.error("=== 数据更新失败 ===")
            
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
    finally:
        if updater:
            updater.close_connection()


if __name__ == "__main__":
    main() 