"""
第5步：第二轮大模型精标润色
对第一轮标注结果进行精细化和润色
"""

import pandas as pd
import asyncio
from typing import List, Dict, Any, Optional
from ..core import PipelineStep, StepResult, PipelineConfig
from maque.performance import MeasureTime
from flexllm.mllm_client import MllmClient
from flexllm.async_api import ConcurrentExecutor


class MllmRefinementStep(PipelineStep):
    """MLLM精标润色步骤"""
    
    CONFIG_SCHEMA = {
        "type": "object",
        "properties": {
            "mllm_config": {
                "type": "object",
                "properties": {
                    "model_name": {"type": "string", "default": "gpt-4o"},
                    "base_url": {"type": "string"},
                    "api_key": {"type": "string"},
                    "temperature": {"type": "number", "default": 0.3},
                    "max_tokens": {"type": "integer", "default": 3072}
                },
                "required": ["model_name"]
            },
            "refinement_prompts": {
                "type": "object",
                "properties": {
                    "system_prompt": {
                        "type": "string",
                        "default": "你是一个专业的数据质量专家，负责对已有的标注结果进行精细化和质量提升。"
                    },
                    "user_prompt_template": {
                        "type": "string",
                        "default": "请对以下标注结果进行精细化处理：\n\n原始内容：\n文本：{text}\n\n第一轮标注结果：\n摘要：{summary}\n标签：{tags}\n情感：{sentiment}\n质量：{quality}\n\n请提供：\n1. 优化后的摘要\n2. 更准确的标签\n3. 精确的情感分析\n4. 详细的质量评估\n5. 改进建议"
                    }
                }
            },
            "refinement_criteria": {
                "type": "object",
                "properties": {
                    "focus_areas": {
                        "type": "array",
                        "items": {"type": "string"},
                        "default": ["accuracy", "completeness", "consistency", "clarity"],
                        "description": "重点改进的方面"
                    },
                    "quality_threshold": {
                        "type": "number",
                        "default": 0.7,
                        "description": "质量阈值，低于此值的数据会被优先精标"
                    },
                    "skip_high_quality": {
                        "type": "boolean",
                        "default": True,
                        "description": "是否跳过高质量数据的精标"
                    }
                }
            },
            "concurrent_config": {
                "type": "object",
                "properties": {
                    "max_workers": {"type": "integer", "default": 3},
                    "batch_size": {"type": "integer", "default": 5},
                    "rate_limit": {"type": "number", "default": 0.5}
                }
            },
            "output_columns": {
                "type": "object",
                "properties": {
                    "refined_summary": {"type": "string", "default": "refined_summary"},
                    "refined_tags": {"type": "string", "default": "refined_tags"},
                    "refined_sentiment": {"type": "string", "default": "refined_sentiment"},
                    "refined_quality": {"type": "string", "default": "refined_quality"},
                    "improvement_suggestions": {"type": "string", "default": "improvement_suggestions"},
                    "refinement_raw_response": {"type": "string", "default": "refinement_raw_response"}
                }
            },
            "retry_config": {
                "type": "object",
                "properties": {
                    "max_retries": {"type": "integer", "default": 2},
                    "retry_delay": {"type": "number", "default": 2.0}
                }
            }
        }
    }
    
    def __init__(self, name: str = "mllm_refinement", config: Dict[str, Any] = None):
        super().__init__(name, config)
        self.mllm_client: Optional[MllmClient] = None
    
    async def execute(self, data: pd.DataFrame, config: PipelineConfig) -> StepResult:
        """执行MLLM精标润色"""
        with MeasureTime(f"步骤 {self.name}"):
            try:
                step_config = self.get_step_config(config)
                data_copy = data.copy()
                
                # 初始化MLLM客户端
                await self._initialize_mllm_client(step_config)
                
                # 筛选需要精标的数据
                refinement_candidates = await self._select_refinement_candidates(data_copy, step_config)
                
                if not refinement_candidates:
                    self.logger.info("没有需要精标的数据")
                    return StepResult(
                        step_name=self.name,
                        success=True,
                        data=data_copy,
                        metadata={"skipped": True, "reason": "no_candidates"}
                    )
                
                # 准备精标任务
                refinement_tasks = await self._prepare_refinement_tasks(data_copy, refinement_candidates, step_config)
                
                # 执行并发精标
                refinement_results = await self._execute_concurrent_refinement(
                    refinement_tasks, step_config
                )
                
                # 处理精标结果
                await self._process_refinement_results(
                    data_copy, refinement_results, step_config
                )
                
                # 更新处理状态
                data_copy['__processing_status'] = 'refined'
                data_copy['__refined_at'] = pd.Timestamp.now()
                
                # 统计结果
                successful_refinements = len([r for r in refinement_results if r.get('success', False)])
                failed_refinements = len(refinement_results) - successful_refinements
                
                metadata = {
                    "total_rows": len(data_copy),
                    "refinement_candidates": len(refinement_candidates),
                    "successful_refinements": successful_refinements,
                    "failed_refinements": failed_refinements,
                    "success_rate": successful_refinements / len(refinement_results) if refinement_results else 0,
                    "mllm_model": step_config.get("mllm_config", {}).get("model_name", "unknown")
                }
                
                self.logger.info(f"MLLM精标完成，成功率: {metadata['success_rate']:.2%}")
                
                return StepResult(
                    step_name=self.name,
                    success=True,
                    data=data_copy,
                    metadata=metadata
                )
                
            except Exception as e:
                self.logger.error(f"MLLM精标失败: {e}")
                return StepResult(
                    step_name=self.name,
                    success=False,
                    data=data,
                    metadata={},
                    error=str(e)
                )
    
    async def _initialize_mllm_client(self, step_config: Dict[str, Any]):
        """初始化MLLM客户端"""
        mllm_config = step_config.get("mllm_config", {})
        
        self.mllm_client = MllmClient(
            model_name=mllm_config.get("model_name", "gpt-4o"),  
            base_url=mllm_config.get("base_url"),
            api_key=mllm_config.get("api_key"),
            temperature=mllm_config.get("temperature", 0.3),
            max_tokens=mllm_config.get("max_tokens", 3072)
        )
        
        self.logger.info(f"初始化精标MLLM客户端: {mllm_config.get('model_name', 'gpt-4o')}")
    
    async def _select_refinement_candidates(self, data: pd.DataFrame, step_config: Dict[str, Any]) -> List[int]:
        """选择需要精标的数据"""
        criteria = step_config.get("refinement_criteria", {})
        quality_threshold = criteria.get("quality_threshold", 0.7)
        skip_high_quality = criteria.get("skip_high_quality", True)
        
        candidates = []
        
        # 检查是否有第一轮标注结果
        if '__mllm_annotation_success' not in data.columns:
            self.logger.warning("未找到第一轮标注结果，将对所有数据进行精标")
            return list(data.index)
        
        for idx, row in data.iterrows():
            # 只对第一轮标注成功的数据进行精标
            if not row.get('__mllm_annotation_success', False):
                continue
            
            should_refine = True
            
            if skip_high_quality:
                # 评估质量，决定是否需要精标
                quality_score = self._assess_annotation_quality(row)
                if quality_score >= quality_threshold:
                    should_refine = False
            
            if should_refine:
                candidates.append(idx)
        
        self.logger.info(f"选择了 {len(candidates)} 条数据进行精标")
        return candidates
    
    def _assess_annotation_quality(self, row: pd.Series) -> float:
        """评估标注质量"""
        quality_score = 0.0
        factors = 0
        
        # 检查摘要质量
        summary = str(row.get('mllm_summary', ''))
        if summary and len(summary.strip()) >= 20:
            quality_score += 0.25
        factors += 1
        
        # 检查标签质量
        tags = str(row.get('mllm_tags', ''))
        if tags and len(tags.strip()) >= 5:
            quality_score += 0.25
        factors += 1
        
        # 检查情感分析
        sentiment = str(row.get('mllm_sentiment', ''))
        if sentiment and sentiment.strip():
            quality_score += 0.25
        factors += 1
        
        # 检查质量评估
        quality = str(row.get('mllm_quality', ''))
        if quality and quality.strip():
            quality_score += 0.25
        factors += 1
        
        return quality_score
    
    async def _prepare_refinement_tasks(self, data: pd.DataFrame, candidates: List[int], step_config: Dict[str, Any]) -> List[Dict[str, Any]]:
        """准备精标任务"""
        tasks = []
        
        prompts = step_config.get("refinement_prompts", {})
        system_prompt = prompts.get("system_prompt", "你是一个专业的数据质量专家。")
        user_prompt_template = prompts.get("user_prompt_template", "请对以下标注结果进行精细化处理：\n\n原始内容：\n文本：{text}\n\n第一轮标注结果：\n摘要：{summary}\n标签：{tags}\n情感：{sentiment}\n质量：{quality}")
        
        for idx in candidates:
            row = data.loc[idx]
            
            text_content = str(row.get('text', ''))
            current_summary = str(row.get('mllm_summary', ''))
            current_tags = str(row.get('mllm_tags', ''))
            current_sentiment = str(row.get('mllm_sentiment', ''))
            current_quality = str(row.get('mllm_quality', ''))
            
            # 构建用户提示
            user_prompt = user_prompt_template.format(
                text=text_content,
                summary=current_summary,
                tags=current_tags,
                sentiment=current_sentiment,
                quality=current_quality
            )
            
            # 准备图像路径（如果有）
            image_list = []
            images = str(row.get('images', ''))
            if images:
                separator = "|"  # 使用固定分隔符
                image_list = [p.strip() for p in images.split(separator) if p.strip()]
            
            task = {
                "row_index": idx,
                "text": text_content,
                "images": image_list,
                "system_prompt": system_prompt,
                "user_prompt": user_prompt,
                "current_annotations": {
                    "summary": current_summary,
                    "tags": current_tags,
                    "sentiment": current_sentiment,
                    "quality": current_quality
                }
            }
            
            tasks.append(task)
        
        self.logger.info(f"准备了 {len(tasks)} 个精标任务")
        return tasks
    
    async def _execute_concurrent_refinement(self, tasks: List[Dict[str, Any]], step_config: Dict[str, Any]) -> List[Dict[str, Any]]:
        """执行并发精标"""
        concurrent_config = step_config.get("concurrent_config", {})
        max_workers = concurrent_config.get("max_workers", 3)
        rate_limit = concurrent_config.get("rate_limit", 0.5)
        
        # 创建并发执行器
        executor = ConcurrentExecutor(
            max_concurrent=max_workers,
            rate_limit=rate_limit
        )
        
        # 准备异步任务
        async_tasks = []
        for task in tasks:
            async_task = self._refine_single_item(task, step_config)
            async_tasks.append(async_task)
        
        # 执行并发精标
        results = await executor.execute_all(async_tasks)
        
        return results
    
    async def _refine_single_item(self, task: Dict[str, Any], step_config: Dict[str, Any]) -> Dict[str, Any]:
        """精标单个数据项"""
        retry_config = step_config.get("retry_config", {})
        max_retries = retry_config.get("max_retries", 2)
        retry_delay = retry_config.get("retry_delay", 2.0)
        
        for attempt in range(max_retries + 1):
            try:
                # 调用MLLM
                response = await self.mllm_client.chat_async(
                    messages=[
                        {"role": "system", "content": task["system_prompt"]},
                        {"role": "user", "content": task["user_prompt"]}
                    ],
                    images=task["images"]
                )
                
                # 解析响应
                parsed_result = self._parse_refinement_response(response)
                
                return {
                    "row_index": task["row_index"],
                    "success": True,
                    "raw_response": response,
                    **parsed_result
                }
                
            except Exception as e:
                self.logger.warning(f"第 {attempt + 1} 次精标失败 (行 {task['row_index']}): {e}")
                
                if attempt < max_retries:
                    await asyncio.sleep(retry_delay * (attempt + 1))
                else:
                    return {
                        "row_index": task["row_index"],
                        "success": False,
                        "error": str(e),
                        "raw_response": "",
                        "refined_summary": "",
                        "refined_tags": "",
                        "refined_sentiment": "",
                        "refined_quality": "",
                        "improvement_suggestions": ""
                    }
    
    def _parse_refinement_response(self, response: str) -> Dict[str, str]:
        """解析精标响应"""
        result = {
            "refined_summary": "",
            "refined_tags": "",
            "refined_sentiment": "",
            "refined_quality": "",
            "improvement_suggestions": ""
        }
        
        try:
            lines = response.strip().split('\n')
            current_section = None
            
            for line in lines:
                line = line.strip()
                if not line:
                    continue
                
                # 检查是否是新的部分标题
                if "优化后的摘要" in line or "refined summary" in line.lower():
                    current_section = "refined_summary"
                elif "更准确的标签" in line or "refined tag" in line.lower():
                    current_section = "refined_tags"
                elif "精确的情感" in line or "refined sentiment" in line.lower():
                    current_section = "refined_sentiment"
                elif "详细的质量" in line or "refined quality" in line.lower():
                    current_section = "refined_quality"
                elif "改进建议" in line or "improvement" in line.lower():
                    current_section = "improvement_suggestions"
                elif current_section and line:
                    # 移除常见的前缀
                    content = line.replace("：", "").replace(":", "").strip()
                    if content and not content.startswith(("1.", "2.", "3.", "4.", "5.")):
                        if result[current_section]:
                            result[current_section] += " " + content
                        else:
                            result[current_section] = content
            
            # 如果解析失败，将整个响应作为改进建议
            if not any(result.values()):
                result["improvement_suggestions"] = response[:800]
                
        except Exception as e:
            self.logger.warning(f"解析精标响应失败: {e}")
            result["improvement_suggestions"] = response[:800] if response else ""
        
        return result
    
    async def _process_refinement_results(self, data: pd.DataFrame, results: List[Dict[str, Any]], step_config: Dict[str, Any]):
        """处理精标结果"""
        output_columns = step_config.get("output_columns", {})
        
        refined_summary_col = output_columns.get("refined_summary", "refined_summary")
        refined_tags_col = output_columns.get("refined_tags", "refined_tags")
        refined_sentiment_col = output_columns.get("refined_sentiment", "refined_sentiment")
        refined_quality_col = output_columns.get("refined_quality", "refined_quality")
        improvement_suggestions_col = output_columns.get("improvement_suggestions", "improvement_suggestions")
        refinement_raw_response_col = output_columns.get("refinement_raw_response", "refinement_raw_response")
        
        # 初始化新列
        data[refined_summary_col] = ""
        data[refined_tags_col] = ""
        data[refined_sentiment_col] = ""
        data[refined_quality_col] = ""
        data[improvement_suggestions_col] = ""
        data[refinement_raw_response_col] = ""
        data['__mllm_refinement_success'] = False
        
        # 填充结果
        for result in results:
            row_idx = result["row_index"]
            
            data.at[row_idx, refined_summary_col] = result.get("refined_summary", "")
            data.at[row_idx, refined_tags_col] = result.get("refined_tags", "")
            data.at[row_idx, refined_sentiment_col] = result.get("refined_sentiment", "")
            data.at[row_idx, refined_quality_col] = result.get("refined_quality", "")
            data.at[row_idx, improvement_suggestions_col] = result.get("improvement_suggestions", "")
            data.at[row_idx, refinement_raw_response_col] = result.get("raw_response", "")
            data.at[row_idx, '__mllm_refinement_success'] = result.get("success", False)