"""
AI审查引擎
负责执行AI审查，支持审查条目并发处理
"""

import asyncio
from enum import Enum
from dataclasses import dataclass
from typing import Dict, List, Any
from foundation.logger.loggering import server_logger as logger
from core.construction_review.component.reviewers.base_reviewer import BaseReviewer,BaseRAGReviewer
@dataclass
class ReviewResult:
    """审查结果"""
    unit_index: int
    unit_content: Dict[str, Any]
    basic_compliance: Dict[str, Any]
    technical_compliance: Dict[str, Any]
    rag_enhanced: Dict[str, Any]
    overall_risk: str

class Stage(Enum):
    """工作流状态"""
    BASIC = {
        'reviewer_type':'basic',
        'sensitive': 'sensitive_word_check',
        'semantic': 'semantic_logic_check',
        'completeness': 'completeness_check',
        'timeliness': 'timeliness_check',
        'reference': 'reference_check'
    }
    TECHNICAL = {
        'reviewer_type':'technical',
        'mandatory': 'mandatory_standards_check',
        'technical': 'technical_parameters_check',
        'design': 'design_values_check'
    }
    RAG = {
        'reviewer_type':'rag',
        'rag': 'rag_enhanced_review',
        'vector': 'vector_search_review',
        'hybrid': 'hybrid_search_review'
    }
    AI = {
        'reviewer_type':'ai',
        'professional': 'professional_suggestion',
        'standardization': 'standardization_suggestion',
        'completeness': 'completeness_suggestion',
        'readability': 'readability_suggestion'
    }


class AIReviewEngine(BaseReviewer):
    """AI审查引擎 - 支持审查条目并发"""

    def __init__(self, max_concurrent_reviews: int = 8):
        super().__init__()  
        self.max_concurrent_reviews = max_concurrent_reviews
        self.semaphore = asyncio.Semaphore(max_concurrent_reviews)



    
    async def basic_compliance_check(self,trace_id_idx: str, unit_content: Dict[str, Any]) -> Dict[str, Any]:
        """基础合规性检查"""
        review_content = unit_content['content']
        review_references = unit_content.get('review_references')

        logger.info(f"开始基础合规性检查, 内容长度: {len(review_content)}")

        async def check_with_semaphore(check_func, *args):
            async with self.semaphore:
                return await check_func(*args)

        basic_tasks = [
            check_with_semaphore(self.check_grammar, trace_id_idx, review_content, review_references),
            check_with_semaphore(self.check_semantic_logic, trace_id_idx, review_content, review_references),
            check_with_semaphore(self.check_completeness, trace_id_idx, review_content, review_references)
        ]


        grammar_result, semantic_result, completeness_result = await asyncio.gather(*basic_tasks, return_exceptions=True)



        if isinstance(grammar_result, Exception):
            grammar_result = {"error": str(grammar_result), "success": False}
        if isinstance(semantic_result, Exception):
            semantic_result = {"error": str(semantic_result), "success": False}
        if isinstance(completeness_result, Exception):
            completeness_result = {"error": str(completeness_result), "success": False}

        return {
            'grammar_check': grammar_result,
            'semantic_check': semantic_result,
            'completeness_check': completeness_result,
            'overall_score': self._calculate_basic_score(grammar_result, semantic_result, completeness_result)
        }

    async def technical_compliance_check(self, trace_id_idx: str, unit_content: Dict[str, Any]) -> Dict[str, Any]:
        """技术性合规检查"""
        review_content = unit_content['content']
        review_references = unit_content.get('review_references')
        logger.info(f"开始技术性合规检查，内容长度: {len(review_content)}")
        async def check_with_semaphore(check_func, *args):
            async with self.semaphore:
                return await check_func(*args)

        technical_tasks = [
            check_with_semaphore(self.check_mandatory_standards, trace_id_idx, review_content,review_references),
            check_with_semaphore(self.check_design_values, trace_id_idx, review_content,review_references),
            check_with_semaphore(self.check_technical_parameters, trace_id_idx, review_content,review_references)
        ]

        mandatory_result, design_value_result, technical_param_result = await asyncio.gather(*technical_tasks, return_exceptions=True)

        # 处理异常结果
        if isinstance(mandatory_result, Exception):
            mandatory_result = {"error": str(mandatory_result), "success": False}
        if isinstance(design_value_result, Exception):
            design_value_result = {"error": str(design_value_result), "success": False}
        if isinstance(technical_param_result, Exception):
            technical_param_result = {"error": str(technical_param_result), "success": False}

        return {
            'mandatory_standards': mandatory_result,
            'design_values': design_value_result,
            'technical_parameters': technical_param_result,
            'overall_score': self._calculate_technical_score(mandatory_result, design_value_result, technical_param_result)
        }

    async def rag_enhanced_check(self, unit_content: Dict[str, Any]) -> Dict[str, Any]:
        """RAG增强审查"""
        # 向量检索
        vector_results = await self.vector_search(unit_content['content'])

        # 混合检索
        hybrid_results = await self.hybrid_search(unit_content['content'])

        # 重排序
        reranked_results = await self.rerank_results(unit_content['content'], hybrid_results)

        return {
            'vector_search': vector_results,
            'hybrid_search': hybrid_results,
            'reranked_results': reranked_results,
            'enhanced_suggestions': self.generate_enhanced_suggestions(reranked_results)
        }


    async def check_grammar(self, trace_id_idx: str, review_content: str = None, review_references: str = None) -> Dict[str, Any]:
        """语法检查"""
        reviewer_type = Stage.BASIC.value['reviewer_type']
        prompt_name = Stage.BASIC.value['sensitive']
        trace_id = prompt_name+trace_id_idx
        return await self.review("语法检查", trace_id, reviewer_type, prompt_name, review_content,review_references)

    async def check_semantic_logic(self, trace_id_idx: str, review_content: str = None, review_references: str = None) -> Dict[str, Any]:
        """语义逻辑检查"""
        reviewer_type = Stage.BASIC.value['reviewer_type']
        prompt_name = Stage.BASIC.value['semantic']
        trace_id = prompt_name+trace_id_idx
        return await self.review("语义逻辑检查", trace_id, reviewer_type, prompt_name, review_content,review_references)

    async def check_completeness(self, trace_id_idx: str, review_content: str = None, review_references: str = None) -> Dict[str, Any]:
        """完整性检查"""
        reviewer_type = Stage.BASIC.value['reviewer_type']
        prompt_name = Stage.BASIC.value['completeness']
        trace_id = prompt_name+trace_id_idx
        return await self.review("完整性检查", trace_id, reviewer_type, prompt_name, review_content,review_references)

    async def check_mandatory_standards(self, trace_id_idx: str, review_content: str = None, review_references: str = None) -> Dict[str, Any]:
        """强制性标准检查"""
        reviewer_type = Stage.TECHNICAL.value['reviewer_type']
        prompt_name = Stage.TECHNICAL.value['mandatory']
        trace_id = prompt_name+trace_id_idx
        return await self.review("强制性标准检查", trace_id, reviewer_type, prompt_name, review_content,review_references)

    async def check_design_values(self, trace_id_idx: str, review_content: str = None, review_references: str = None) -> Dict[str, Any]:
        """设计值检查"""
        reviewer_type = Stage.TECHNICAL.value['reviewer_type']
        prompt_name = Stage.TECHNICAL.value['design']
        trace_id = prompt_name+trace_id_idx
        return await self.review("设计值检查", trace_id, reviewer_type, prompt_name, review_content,review_references)

    async def check_technical_parameters(self, trace_id_idx: str, review_content: str = None, review_references: str = None) -> Dict[str, Any]:
        """技术参数检查"""
        reviewer_type = Stage.TECHNICAL.value['reviewer_type']
        prompt_name = Stage.TECHNICAL.value['technical']
        trace_id = prompt_name+trace_id_idx
        return await self.review("技术参数检查", trace_id, reviewer_type, prompt_name, review_content,review_references)

    # RAG检索增强
    async def vector_search(self, content: str) -> List[Dict[str, Any]]:
        """向量检索"""
        await asyncio.sleep(0.1)
        return [{"similarity": 0.85, "content": "相关标准1"}, {"similarity": 0.78, "content": "相关标准2"}]

    async def hybrid_search(self, content: str) -> List[Dict[str, Any]]:
        """混合检索"""
        await asyncio.sleep(0.2)
        return [{"score": 0.88, "content": "混合检索结果1"}, {"score": 0.82, "content": "混合检索结果2"}]

    async def rerank_results(self, content: str, results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """重排序"""
        await asyncio.sleep(0.1)
        return sorted(results, key=lambda x: x.get('score', 0), reverse=True)[:5]

    def generate_enhanced_suggestions(self, results: List[Dict[str, Any]]) -> List[str]:
        """生成增强建议"""
        suggestions = []
        for result in results:
            suggestions.append(f"基于{result.get('content', '相关内容')}的建议")
        return suggestions

    def _calculate_basic_score(self, grammar: Dict, semantic: Dict, completeness: Dict) -> float:
        """计算基础合规性得分"""
        return (grammar.get('score', 0) + semantic.get('score', 0) + completeness.get('score', 0)) / 3

    def _calculate_technical_score(self, mandatory: Dict, design: Dict, technical: Dict) -> float:
        """计算技术合规性得分"""
        return (mandatory.get('compliance_rate', 0) + design.get('accuracy', 0) + technical.get('precision', 0)) / 3

    def _calculate_overall_risk(self, basic: Dict, technical: Dict, rag: Dict) -> str:
        """计算总体风险等级"""
        basic_score = basic.get('overall_score', 0)
        technical_score = technical.get('overall_score', 0)

        # 如果RAG被禁用或有错误，忽略它
        if 'error' in rag:
            avg_score = (basic_score + technical_score) / 2
        else:
            rag_score = rag.get('overall_score', 0)
            avg_score = (basic_score + technical_score + rag_score) / 3

        if avg_score >= 90:
            return "low"
        elif avg_score >= 70:
            return "medium"
        else:
            return "high"

    def _aggregate_results(self, results: List[ReviewResult]) -> Dict[str, Any]:
        """汇总审查结果"""
        risk_counts = {"high": 0, "medium": 0, "low": 0}

        for result in results:
            risk_counts[result.overall_risk] += 1

        return {
            "risk_distribution": risk_counts,
            "total_issues": len([r for r in results if r.overall_risk != "low"]),
            "high_risk_count": risk_counts["high"],
            "medium_risk_count": risk_counts["medium"],
            "low_risk_count": risk_counts["low"]
        }