"""
标签关联分析数据库操作模块
用于连接 MongoDB 并存储分析结果
"""

import logging
import pymongo
from pymongo import ReturnDocument
from typing import Dict, Any, Optional, List
from datetime import datetime

# 导入配置
import sys
import os

# 添加项目路径
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(ROOT_DIR)

try:
    from valuation_analysis.config import MONGO_CONFIG222
except ImportError:
    try:
        from src.valuation_analysis.config import MONGO_CONFIG222
    except ImportError:
        # 尝试从当前目录导入
        sys.path.append(os.path.dirname(ROOT_DIR))
        from src.valuation_analysis.config import MONGO_CONFIG222

# 设置日志记录
logger = logging.getLogger(__name__)


class TagRelationDatabase:
    """标签关联分析数据库操作类"""
    
    def __init__(self, collection_name: Optional[str] = None):
        """初始化数据库连接
        
        Args:
            collection_name: 集合名称，如果不指定则使用配置中的默认集合名
        """
        self.mongo_client = None
        self.db = None
        self.collection_name = collection_name or MONGO_CONFIG222.get('collection', 'tag_relation_analysis')
        self.collection = None
        
        # 标签集合和关联集合
        self.tags_collection = None
        self.tag_stock_relations_collection = None
        self.counters_collection = None
        
        self.connect_mongodb()
    
    def connect_mongodb(self):
        """连接MongoDB数据库"""
        try:
            self.mongo_client = pymongo.MongoClient(
                host=MONGO_CONFIG222['host'],
                port=MONGO_CONFIG222['port'],
                username=MONGO_CONFIG222['username'],
                password=MONGO_CONFIG222['password'],
                authSource='admin'  # 认证数据库
            )
            self.db = self.mongo_client[MONGO_CONFIG222['db']]
            self.collection = self.db[self.collection_name]
            
            # 初始化标签集合和关联集合
            self.tags_collection = self.db['stock_pool_tags']  # 股票池标签集合
            self.tag_stock_relations_collection = self.db['tag_stock_relations']
            self.counters_collection = self.db['counters']

            # 确保索引（唯一约束）
            self._ensure_indexes()
            
            # 测试连接
            self.mongo_client.admin.command('ping')
            logger.info(f"MongoDB连接成功，数据库: {MONGO_CONFIG222['db']}, 集合: {self.collection_name}")
            
        except Exception as e:
            logger.error(f"MongoDB连接失败: {str(e)}")
            raise
    
    def _ensure_indexes(self):
        """为标签集合创建唯一索引，防止 tag_name / tag_code 冲突"""
        try:
            # tag_name 唯一
            self.tags_collection.create_index(
                [('tag_name', pymongo.ASCENDING)],
                unique=True,
                name='idx_tag_name_unique'
            )
            # tag_code 唯一
            self.tags_collection.create_index(
                [('tag_code', pymongo.ASCENDING)],
                unique=True,
                name='idx_tag_code_unique'
            )
        except Exception as e:
            logger.warning(f"创建索引时发生异常: {str(e)}")
    
    def _get_next_tag_code(self) -> str:
        """使用 counters 集合原子递增生成 tag_code"""
        try:
            result = self.counters_collection.find_one_and_update(
                {'_id': 'tag_code'},
                {'$inc': {'seq': 1}},
                upsert=True,
                return_document=ReturnDocument.AFTER
            )
            seq = result.get('seq', 1)
            return str(seq)
        except Exception as e:
            logger.error(f"生成tag_code失败: {str(e)}", exc_info=True)
            # 兜底：使用时间戳
            return str(int(datetime.now().timestamp() * 1000))
    
    def save_analysis_result(self, analysis_result: Dict[str, Any]) -> bool:
        """保存分析结果到MongoDB
        
        Args:
            analysis_result: 分析结果字典，包含所有分析字段
            
        Returns:
            bool: 是否保存成功
        """
        try:
            if not analysis_result:
                logger.warning("分析结果为空，无法保存")
                return False
            
            # 添加保存时间
            analysis_result['save_time'] = datetime.now().isoformat()
            
            # 使用 stock_code 和 tag_name 作为唯一标识
            filter_condition = {
                'stock_code': analysis_result.get('stock_code'),
                'tag_name': analysis_result.get('tag_name')
            }
            
            # 检查是否已存在
            existing_record = self.collection.find_one(filter_condition)
            
            if existing_record:
                # 更新现有记录
                update_result = self.collection.update_one(
                    filter_condition,
                    {'$set': analysis_result}
                )
                logger.info(f"更新分析结果: {analysis_result.get('stock_code')} - {analysis_result.get('tag_name')}")
                return update_result.modified_count > 0
            else:
                # 插入新记录
                insert_result = self.collection.insert_one(analysis_result)
                logger.info(f"插入新分析结果: {analysis_result.get('stock_code')} - {analysis_result.get('tag_name')}, ID: {insert_result.inserted_id}")
                return insert_result.inserted_id is not None
                
        except Exception as e:
            logger.error(f"保存分析结果失败: {str(e)}", exc_info=True)
            return False
    
    def get_analysis_result(self, stock_code: str, tag_name: str) -> Optional[Dict[str, Any]]:
        """获取指定股票和标签的分析结果
        
        Args:
            stock_code: 股票代码
            tag_name: 标签名称
            
        Returns:
            Optional[Dict[str, Any]]: 分析结果，如果不存在返回None
        """
        try:
            result = self.collection.find_one({
                'stock_code': stock_code,
                'tag_name': tag_name
            })
            
            if result:
                # 移除 MongoDB 的 _id 字段（如果需要）
                result.pop('_id', None)
            
            return result
            
        except Exception as e:
            logger.error(f"获取分析结果失败: {str(e)}")
            return None
    
    def get_stock_analyses(self, stock_code: str) -> list:
        """获取指定股票的所有标签分析结果
        
        Args:
            stock_code: 股票代码
            
        Returns:
            list: 分析结果列表
        """
        try:
            results = list(self.collection.find({'stock_code': stock_code}))
            
            # 移除 MongoDB 的 _id 字段
            for result in results:
                result.pop('_id', None)
            
            return results
            
        except Exception as e:
            logger.error(f"获取股票分析结果失败: {str(e)}")
            return []
    
    def get_tag_analyses(self, tag_name: str) -> list:
        """获取指定标签的所有股票分析结果
        
        Args:
            tag_name: 标签名称
            
        Returns:
            list: 分析结果列表
        """
        try:
            results = list(self.collection.find({'tag_name': tag_name}))
            
            # 移除 MongoDB 的 _id 字段
            for result in results:
                result.pop('_id', None)
            
            return results
            
        except Exception as e:
            logger.error(f"获取标签分析结果失败: {str(e)}")
            return []
    
    def save_tag(self, tag_data: Dict[str, Any]) -> Optional[str]:
        """保存标签到tags集合
        
        Args:
            tag_data: 标签数据字典，包含tag_id, tag_name, tag_type等
            
        Returns:
            Optional[str]: 标签ID（tag_code），如果保存失败返回None
        """
        try:
            # 检查标签是否已存在（根据tag_name）
            tag_name = tag_data.get('tag_name')
            if not tag_name:
                logger.warning("标签名称为空，无法保存")
                return None
            
            existing_tag = self.tags_collection.find_one({'tag_name': tag_name})
            
            if existing_tag:
                tag_code = existing_tag.get('tag_code')
                # 保持原 tag_code，不覆盖 create_time
                update_data = dict(tag_data)
                update_data['tag_code'] = tag_code
                # 避免与 create_time 冲突
                update_data.pop('create_time', None)
                self.tags_collection.update_one(
                    {'tag_name': tag_name},
                    {'$set': update_data}
                )
                logger.info(f"更新标签: {tag_name}, tag_code: {tag_code}")
                return tag_code
            else:
                # 生成新的 tag_code（原子递增）
                tag_code = tag_data.get('tag_code') or self._get_next_tag_code()
                tag_data = dict(tag_data)
                # 新增时通过 $setOnInsert 写入 tag_code，避免冲突
                tag_data.pop('tag_code', None)
                create_time = tag_data.pop('create_time', datetime.now().isoformat())
                
                self.tags_collection.update_one(
                    {'tag_name': tag_name},
                    {
                        '$set': tag_data,
                        '$setOnInsert': {
                            'create_time': create_time,
                            'tag_code': tag_code
                        }
                    },
                    upsert=True
                )
                logger.info(f"插入新标签: {tag_name}, tag_code: {tag_code}")
                return tag_code
                
        except Exception as e:
            logger.error(f"保存标签失败: {str(e)}", exc_info=True)
            return None
    
    def get_tag_by_name(self, tag_name: str) -> Optional[Dict[str, Any]]:
        """根据标签名称获取标签信息
        
        Args:
            tag_name: 标签名称
            
        Returns:
            Optional[Dict[str, Any]]: 标签信息，如果不存在返回None
        """
        try:
            tag = self.tags_collection.find_one({'tag_name': tag_name})
            if tag:
                tag.pop('_id', None)
            return tag
        except Exception as e:
            logger.error(f"获取标签失败: {str(e)}")
            return None
    
    def save_tag_stock_relations(self, tag_code: str, stock_codes: List[str], replace_all: bool = False) -> bool:
        """保存标签与股票的关联关系
        
        Args:
            tag_code: 标签代码
            stock_codes: 股票代码列表
            replace_all: 是否替换所有关联关系（True=全量替换，False=增量添加），默认False
            
        Returns:
            bool: 是否保存成功
        """
        try:
            if not stock_codes:
                logger.warning("股票代码列表为空")
                return False
            
            # 如果replace_all为True，先删除该标签的所有旧关联关系
            if replace_all:
                self.tag_stock_relations_collection.delete_many({'tag_code': tag_code})
            
            # 批量插入新关联关系（使用upsert避免重复）
            relations = []
            for stock_code in stock_codes:
                # 检查是否已存在
                existing = self.tag_stock_relations_collection.find_one({
                    'tag_code': tag_code,
                    'stock_code': stock_code
                })
                
                if not existing:
                    relations.append({
                        'tag_code': tag_code,
                        'stock_code': stock_code,
                        'create_time': datetime.now().isoformat()
                    })
            
            if relations:
                self.tag_stock_relations_collection.insert_many(relations)
                logger.info(f"保存标签关联关系: tag_code={tag_code}, 新增股票数量={len(relations)}")
            
            return True
            
        except Exception as e:
            logger.error(f"保存标签关联关系失败: {str(e)}", exc_info=True)
            return False
    
    def get_tag_stocks(self, tag_code: str) -> List[str]:
        """获取标签关联的所有股票代码
        
        Args:
            tag_code: 标签代码
            
        Returns:
            List[str]: 股票代码列表
        """
        try:
            relations = list(self.tag_stock_relations_collection.find({'tag_code': tag_code}))
            stock_codes = [r.get('stock_code') for r in relations if r.get('stock_code')]
            return stock_codes
        except Exception as e:
            logger.error(f"获取标签关联股票失败: {str(e)}")
            return []
    
    def update_tag_status(self, tag_code: str, status: str, progress: Optional[float] = None):
        """更新标签处理状态和进度
        
        Args:
            tag_code: 标签代码
            status: 处理状态（pending/processing/completed/failed）
            progress: 处理进度（0-100），可选
        """
        try:
            update_data = {
                'status': status,
                'update_time': datetime.now().isoformat()
            }
            if progress is not None:
                update_data['progress'] = progress
            
            self.tags_collection.update_one(
                {'tag_code': tag_code},
                {'$set': update_data}
            )
            logger.info(f"更新标签状态: tag_code={tag_code}, status={status}, progress={progress}")
        except Exception as e:
            logger.error(f"更新标签状态失败: {str(e)}")
    
    def delete_tag(self, tag_code: str, tag_name: str, delete_analysis: bool = False) -> Dict[str, Any]:
        """删除标签及其相关数据
        
        Args:
            tag_code: 标签代码
            tag_name: 标签名称（用于删除分析结果）
            delete_analysis: 是否删除分析结果，默认False
            
        Returns:
            Dict[str, Any]: 删除结果
                - success: 是否成功
                - deleted_tag: 是否删除了标签
                - deleted_relations: 删除的关联关系数量
                - deleted_analyses: 删除的分析结果数量（如果delete_analysis=True）
                - error: 错误信息
        """
        result = {
            'success': False,
            'deleted_tag': False,
            'deleted_relations': 0,
            'deleted_analyses': 0,
            'error': None
        }
        
        try:
            # 1. 删除标签与股票的关联关系（中间表）
            relations_result = self.tag_stock_relations_collection.delete_many({'tag_code': tag_code})
            result['deleted_relations'] = relations_result.deleted_count
            logger.info(f"删除标签关联关系: tag_code={tag_code}, 删除数量={relations_result.deleted_count}")
            
            # 2. 删除标签本身
            tag_result = self.tags_collection.delete_one({'tag_code': tag_code})
            result['deleted_tag'] = tag_result.deleted_count > 0
            logger.info(f"删除标签: tag_code={tag_code}, tag_name={tag_name}, 是否删除={result['deleted_tag']}")
            
            # 3. 可选：删除分析结果
            if delete_analysis:
                analyses_result = self.collection.delete_many({'tag_name': tag_name})
                result['deleted_analyses'] = analyses_result.deleted_count
                logger.info(f"删除分析结果: tag_name={tag_name}, 删除数量={analyses_result.deleted_count}")
            
            result['success'] = True
            return result
            
        except Exception as e:
            logger.error(f"删除标签失败: {str(e)}", exc_info=True)
            result['error'] = str(e)
            return result
    
    def delete_tag_by_name(self, tag_name: str, delete_analysis: bool = False) -> Dict[str, Any]:
        """根据标签名称删除标签及其相关数据
        
        Args:
            tag_name: 标签名称
            delete_analysis: 是否删除分析结果，默认False
            
        Returns:
            Dict[str, Any]: 删除结果
                - success: 是否成功
                - tag_code: 标签代码（如果找到）
                - deleted_tag: 是否删除了标签
                - deleted_relations: 删除的关联关系数量
                - deleted_analyses: 删除的分析结果数量（如果delete_analysis=True）
                - error: 错误信息
        """
        result = {
            'success': False,
            'tag_code': None,
            'deleted_tag': False,
            'deleted_relations': 0,
            'deleted_analyses': 0,
            'error': None
        }
        
        try:
            # 先查找标签
            tag = self.get_tag_by_name(tag_name)
            if not tag:
                result['error'] = f"标签不存在: {tag_name}"
                logger.warning(result['error'])
                return result
            
            tag_code = tag.get('tag_code')
            result['tag_code'] = tag_code
            
            # 调用根据tag_code删除的方法
            delete_result = self.delete_tag(tag_code, tag_name, delete_analysis)
            
            result.update(delete_result)
            return result
            
        except Exception as e:
            logger.error(f"根据标签名称删除标签失败: {str(e)}", exc_info=True)
            result['error'] = str(e)
            return result
    
    def clear_all_tags_and_relations(self) -> Dict[str, Any]:
        """清空行业/概念标签及其与个股的中间关系（不动其他标签、不动分析结果/Redis）
        
        Returns:
            Dict[str, Any]: 删除统计
        """
        result = {
            'success': False,
            'deleted_tags': 0,
            'deleted_relations': 0,
            'error': None
        }
        try:
            # 仅删除行业/概念标签
            tag_filter = {
                '$or': [
                    {'source': {'$in': ['hybk', 'gnbk']}},
                    {'tag_type': {'$in': ['行业标签', '概念标签']}}
                ]
            }
            tags = list(self.tags_collection.find(tag_filter, {'tag_code': 1}))
            tag_codes = [t.get('tag_code') for t in tags if t.get('tag_code')]
            
            tag_res = self.tags_collection.delete_many(tag_filter)
            rel_res = self.tag_stock_relations_collection.delete_many({'tag_code': {'$in': tag_codes}}) if tag_codes else type("Dummy", (), {"deleted_count": 0})()
            
            result['deleted_tags'] = tag_res.deleted_count
            result['deleted_relations'] = rel_res.deleted_count if hasattr(rel_res, 'deleted_count') else 0
            result['success'] = True
            return result
        except Exception as e:
            logger.error(f"清空标签及关系失败: {str(e)}", exc_info=True)
            result['error'] = str(e)
            return result
    
    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)}")

