"""流水线协调器模块"""

from typing import Dict, List, Tuple
import os
import tempfile
import json
from hos_m2f.converters.docx_to_xml_lxml import DOCXToXMLLXMLConverter
from .xml_templater import XMLTemplater
from .markdown_converter import MarkdownConverter
from .template_renderer import TemplateRenderer
from .xml_modifier import XMLModifier
from .schema_generator import SchemaGenerator
from .ai_prompt_manager import AIPromptManager
from .log_manager import LogManager


class PipelineCoordinator:
    """流水线协调器
    
    功能：
    1. 协调各个模块的工作
    2. 实现端到端的文档处理流程
    3. 提供完整的AI文档自动化流水线
    """
    
    def __init__(self, log_dir: str = None):
        """初始化流水线协调器
        
        Args:
            log_dir: 日志目录
        """
        # 初始化日志管理器
        self.logger = LogManager(log_dir)
        
        # 初始化各个模块
        self.docx_to_xml = DOCXToXMLLXMLConverter()
        self.xml_templater = XMLTemplater()
        self.markdown_converter = MarkdownConverter()
        self.template_renderer = TemplateRenderer()
        self.xml_modifier = XMLModifier()
        self.schema_generator = SchemaGenerator()
        self.ai_prompt_manager = AIPromptManager()
    
    def process_docx(self, docx_path: str, output_dir: str = None) -> Dict[str, str]:
        """处理DOCX文档，生成模板和映射表
        
        Args:
            docx_path: DOCX文件路径
            output_dir: 输出目录
            
        Returns:
            Dict[str, str]: 处理结果，包含模板XML路径、映射表路径等
        """
        if output_dir is None:
            output_dir = tempfile.mkdtemp()
        else:
            os.makedirs(output_dir, exist_ok=True)
        
        # 1. DOCX → XML
        with open(docx_path, 'rb') as f:
            docx_content = f.read()
        
        xml_content = self.docx_to_xml.convert(docx_content)
        
        # 确保xml_content是字符串
        if isinstance(xml_content, bytes):
            xml_content = xml_content.decode('utf-8')
        
        # 保存原始XML
        raw_xml_path = os.path.join(output_dir, 'raw.xml')
        with open(raw_xml_path, 'w', encoding='utf-8') as f:
            f.write(xml_content)
        
        # 生成Schema
        schema = self.schema_generator.generate_from_xml(xml_content, format='json_schema')
        
        # 保存Schema
        schema_path = os.path.join(output_dir, 'schema.json')
        with open(schema_path, 'w', encoding='utf-8') as f:
            json.dump(schema, f, ensure_ascii=False, indent=2)
        
        # 2. XML → 模板
        template_xml, variable_map = self.xml_templater.template_auto(xml_content)
        
        # 保存模板XML
        template_xml_path = os.path.join(output_dir, 'template.xml')
        with open(template_xml_path, 'w', encoding='utf-8') as f:
            f.write(template_xml)
        
        # 保存变量映射表
        variable_map_path = os.path.join(output_dir, 'variable_map.json')
        with open(variable_map_path, 'w', encoding='utf-8') as f:
            json.dump(variable_map, f, ensure_ascii=False, indent=2)
        
        return {
            'raw_xml': raw_xml_path,
            'schema': schema_path,
            'template_xml': template_xml_path,
            'variable_map': variable_map_path,
            'output_dir': output_dir
        }
    
    def generate_content(self, template_xml_path: str, variable_map_path: str, ai_content: str, output_dir: str = None) -> Dict[str, str]:
        """使用AI生成的内容填充模板
        
        Args:
            template_xml_path: 模板XML路径
            variable_map_path: 变量映射表路径
            ai_content: AI生成的Markdown内容
            output_dir: 输出目录
            
        Returns:
            Dict[str, str]: 处理结果，包含填充后的XML路径等
        """
        if output_dir is None:
            output_dir = os.path.dirname(template_xml_path)
        
        # 读取模板和映射表
        with open(template_xml_path, 'r', encoding='utf-8') as f:
            template_xml = f.read()
        
        with open(variable_map_path, 'r', encoding='utf-8') as f:
            variable_map = json.load(f)
        
        # 智能处理：根据变量映射表创建变量
        variables = {}
        markdown_fields = []
        
        # 检查AI内容是否为JSON格式
        try:
            # 尝试解析为JSON
            ai_data = json.loads(ai_content)
            # 如果是JSON，使用键值对
            for key, value in ai_data.items():
                variables[key] = value
                markdown_fields.append(key)
        except json.JSONDecodeError:
            # 如果不是JSON，将内容分配给所有字段
            for field in variable_map.keys():
                variables[field] = ai_content
                markdown_fields.append(field)
        
        # 渲染模板
        rendered_xml = self.template_renderer.render_with_markdown(
            template_xml, 
            variables, 
            markdown_fields=markdown_fields
        )
        
        # 保存渲染后的XML
        rendered_xml_path = os.path.join(output_dir, 'rendered.xml')
        with open(rendered_xml_path, 'w', encoding='utf-8') as f:
            f.write(rendered_xml)
        
        # 保存字段值
        field_values_path = os.path.join(output_dir, 'field_values.json')
        with open(field_values_path, 'w', encoding='utf-8') as f:
            json.dump(variables, f, ensure_ascii=False, indent=2)
        
        return {
            'rendered_xml': rendered_xml_path,
            'field_values': field_values_path,
            'output_dir': output_dir
        }
    
    def refine_xml(self, xml_path: str, ai_instructions: str, output_dir: str = None) -> Dict[str, str]:
        """使用AI指令微调XML
        
        Args:
            xml_path: XML文件路径
            ai_instructions: AI生成的修改指令（JSON Patch格式）
            output_dir: 输出目录
            
        Returns:
            Dict[str, str]: 处理结果，包含修改后的XML路径等
        """
        if output_dir is None:
            output_dir = os.path.dirname(xml_path)
        
        # 读取XML
        encodings = ['utf-8', 'gbk', 'ansi']
        xml_content = None
        
        for encoding in encodings:
            try:
                with open(xml_path, 'r', encoding=encoding) as f:
                    xml_content = f.read()
                break
            except UnicodeDecodeError:
                continue
        
        if xml_content is None:
            raise Exception("无法读取XML文件：所有编码尝试都失败")
        
        # 应用AI指令
        try:
            modified_xml = self.xml_modifier.modify_by_json_patch(xml_content, ai_instructions)
        except Exception as e:
            # 如果修改失败，返回原始XML
            modified_xml = xml_content
            self.logger.warning(f"XML修改失败: {str(e)}")
        
        # 保存修改后的XML
        modified_xml_path = os.path.join(output_dir, 'final.xml')
        with open(modified_xml_path, 'w', encoding='utf-8') as f:
            f.write(modified_xml)
        
        # 保存AI指令
        instructions_path = os.path.join(output_dir, 'ai_instructions.json')
        with open(instructions_path, 'w', encoding='utf-8') as f:
            # 尝试解析为JSON
            try:
                instructions_data = json.loads(ai_instructions)
                json.dump(instructions_data, f, ensure_ascii=False, indent=2)
            except json.JSONDecodeError:
                # 如果不是JSON，保存为字符串
                f.write(ai_instructions)
        
        return {
            'final_xml': modified_xml_path,
            'instructions': instructions_path,
            'output_dir': output_dir
        }
    
    def full_pipeline(self, docx_path: str, ai_content: str, ai_instructions: str, output_dir: str = None) -> Dict[str, str]:
        """完整流水线处理
        
        Args:
            docx_path: DOCX文件路径
            ai_content: AI生成的Markdown内容
            ai_instructions: AI生成的修改指令（JSON Patch格式）
            output_dir: 输出目录
            
        Returns:
            Dict[str, str]: 完整处理结果
        """
        if output_dir is None:
            output_dir = tempfile.mkdtemp()
        
        self.logger.info(f"开始完整流水线处理: {docx_path}")
        self.logger.info(f"输出目录: {output_dir}")
        
        try:
            # 1. 处理DOCX
            self.logger.info("步骤1: 处理DOCX文档...")
            step1_result = self.process_docx(docx_path, output_dir)
            self.logger.info(f"  - 生成原始XML: {step1_result['raw_xml']}")
            self.logger.info(f"  - 生成Schema: {step1_result['schema']}")
            self.logger.info(f"  - 生成模板XML: {step1_result['template_xml']}")
            self.logger.info(f"  - 生成变量映射表: {step1_result['variable_map']}")
            
            # 2. 生成内容
            self.logger.info("步骤2: 生成内容并填充模板...")
            step2_result = self.generate_content(
                step1_result['template_xml'],
                step1_result['variable_map'],
                ai_content,
                output_dir
            )
            self.logger.info(f"  - 生成渲染XML: {step2_result['rendered_xml']}")
            self.logger.info(f"  - 生成字段值: {step2_result['field_values']}")
            
            # 3. 微调XML
            self.logger.info("步骤3: 微调XML...")
            step3_result = self.refine_xml(
                step2_result['rendered_xml'],
                ai_instructions,
                output_dir
            )
            self.logger.info(f"  - 生成最终XML: {step3_result['final_xml']}")
            self.logger.info(f"  - 保存AI指令: {step3_result['instructions']}")
            
            # 合并结果
            result = {
                **step1_result,
                **step2_result,
                **step3_result,
                'status': 'success'
            }
            
            self.logger.info("完整流水线处理完成！")
            return result
            
        except Exception as e:
            self.logger.error(f"流水线处理失败: {str(e)}", exc_info=True)
            return {
                'output_dir': output_dir,
                'status': 'failed',
                'error': str(e)
            }
    
    def generate_ai_prompt(self, task_type: str, **kwargs) -> str:
        """生成AI提示词
        
        Args:
            task_type: 任务类型
            **kwargs: 额外参数
            
        Returns:
            str: AI提示词
        """
        return self.ai_prompt_manager.get_full_pipeline_prompt(task_type, **kwargs)
    
    def get_prompt_for_schema(self, schema_path: str, task_type: str) -> str:
        """根据Schema生成提示词
        
        Args:
            schema_path: Schema文件路径
            task_type: 任务类型
            
        Returns:
            str: 基于Schema的提示词
        """
        with open(schema_path, 'r', encoding='utf-8') as f:
            schema = json.load(f)
        
        return self.ai_prompt_manager.get_schema_based_prompt(schema, task_type)
