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

"""
Topic Generator
主题生成器 - 重构版本，使用动态提示词加载和JSON修复
"""

import logging
from typing import Dict, Any, List, Optional, Tuple
import json

from ..config import AlgorithmConfig
from ..core import AIService, OutputManager, PromptManager, JSONProcessor
from ..exceptions import ContentGenerationError

logger = logging.getLogger(__name__)


class TopicGenerator:
    """
    主题生成器 - 重构版本
    负责生成旅游相关的选题，支持配置化参数和动态提示词
    """

    def __init__(self, config: AlgorithmConfig):
        """
        初始化主题生成器

        Args:
            config: 算法配置
        """
        self.config = config
        self.ai_service = AIService(config.ai_model)
        self.output_manager = OutputManager(config.output)
        self.prompt_manager = PromptManager(config.prompts, config.resources)
        
        # 初始化JSON处理器
        self.json_processor = JSONProcessor(
            enable_repair=config.content_generation.enable_json_repair,
            max_repair_attempts=config.content_generation.json_repair_attempts
        )
        
        # 获取任务特定的模型配置和字段配置
        self.task_model_config = config.ai_model.get_task_config("topic_generation")
        self.field_config = config.content_generation.result_field_mapping.get("topic_generation", {})
        
        logger.info(f"主题生成器初始化完成，使用模型参数: {self.task_model_config}")

    async def generate_topics(
        self, 
        creative_materials: str,
        num_topics: int,
        month: str,
        **kwargs
    ) -> Tuple[str, List[Dict[str, Any]]]:
        """
        生成主题列表

        Args:
            creative_materials: 创作材料
            num_topics: 生成数量
            month: 选题月份
            **kwargs: 其他参数

        Returns:
            Tuple[请求ID, 主题列表]

        Raises:
            ContentGenerationError: 生成失败时抛出
        """
        try:
            logger.info(f"开始执行主题生成流程，数量: {num_topics}, 月份: {month}")

            # 1. 构建提示词
            system_prompt = self.prompt_manager.get_prompt("topic_generation", "system")
            user_prompt_template = self.prompt_manager.get_prompt("topic_generation", "user")
            
            # 格式化用户提示词
            user_prompt = self.prompt_manager.format_prompt(
                user_prompt_template,
                creative_materials=creative_materials,
                numTopics=num_topics,
                month=month,
                **kwargs
            )

            # 保存提示词（如果配置允许）
            if self.config.output.save_prompts:
                self.output_manager.save_text(system_prompt, "system_prompt", "topic_generation")
                self.output_manager.save_text(user_prompt, "user_prompt", "topic_generation")

            # 2. 调用AI生成
            content, input_tokens, output_tokens, elapsed_time = await self.ai_service.generate_text(
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                stage="主题生成",
                **self.task_model_config
            )

            # 保存原始响应（如果配置允许）
            if self.config.output.save_raw_responses:
                self.output_manager.save_text(content, "raw_response", "topic_generation")

            # 3. 解析结果（使用JSON处理器）
            topics = self._parse_topics_result(content)
            if not topics:
                raise ContentGenerationError("未能从AI响应中解析出任何有效主题")

            # 4. 保存结果
            self.output_manager.save_json(topics, "topics")
            
            # 5. 保存元数据
            metadata = {
                "creative_materials": creative_materials[:200] + "..." if len(creative_materials) > 200 else creative_materials,
                "num_topics": num_topics,
                "month": month,
                "generated_count": len(topics),
                "model_config": self.task_model_config,
                "field_config": self.field_config,
                "tokens": {
                    "input": input_tokens,
                    "output": output_tokens
                },
                "elapsed_time": elapsed_time
            }
            self.output_manager.save_metadata(metadata, "topic_generation")

            logger.info(f"成功生成并保存 {len(topics)} 个主题")
            return self.output_manager.run_id, topics

        except Exception as e:
            error_msg = f"主题生成失败: {str(e)}"
            logger.error(error_msg, exc_info=True)
            raise ContentGenerationError(error_msg)

    def _parse_topics_result(self, raw_output: str) -> List[Dict[str, Any]]:
        """
        解析主题生成结果

        Args:
            raw_output: AI原始输出

        Returns:
            解析后的主题列表
        """
        try:
            # 使用JSON处理器解析
            parsed_data = self.json_processor.parse_llm_output(
                raw_output=raw_output,
                expected_fields=self.field_config.get("expected_fields"),
                required_fields=self.field_config.get("required_fields")
            )

            # 确保返回列表格式
            if isinstance(parsed_data, dict):
                # 如果是单个对象，转为列表
                topics = [parsed_data]
            elif isinstance(parsed_data, list):
                topics = parsed_data
            else:
                logger.error(f"解析结果格式错误: {type(parsed_data)}")
                return []

            # 验证和标准化主题数据
            validated_topics = []
            for i, topic in enumerate(topics):
                validated_topic = self._validate_and_normalize_topic(topic, i)
                if validated_topic:
                    validated_topics.append(validated_topic)

            return validated_topics

        except Exception as e:
            logger.error(f"主题结果解析失败: {e}")
            return []

    def _validate_and_normalize_topic(self, topic: Dict[str, Any], index: int) -> Optional[Dict[str, Any]]:
        """
        验证和标准化单个主题数据

        Args:
            topic: 原始主题数据
            index: 主题索引

        Returns:
            标准化后的主题数据
        """
        try:
            if not isinstance(topic, dict):
                logger.warning(f"主题 {index + 1} 不是字典格式，跳过")
                return None

            # 根据实际提示词的字段进行标准化
            normalized_topic = {
                "index": topic.get("index", str(index + 1)),
                "date": topic.get("date", ""),
                "logic": topic.get("logic", ""),
                "object": topic.get("object", "None"),
                "product": topic.get("product", "None"),
                "productLogic": topic.get("productLogic", ""),
                "style": topic.get("style", ""),
                "styleLogic": topic.get("styleLogic", ""),
                "targetAudience": topic.get("targetAudience", ""),
                "targetAudienceLogic": topic.get("targetAudienceLogic", ""),
                # 额外的元数据
                "metadata": {
                    "original_data": topic,
                    "generated_at": self.output_manager.run_id,
                    "validation_passed": True
                }
            }

            # 验证必需字段
            required_fields = self.field_config.get("required_fields", ["index", "date", "logic", "object"])
            missing_fields = [field for field in required_fields if not normalized_topic.get(field)]
            
            if missing_fields:
                logger.warning(f"主题 {index + 1} 缺少必需字段: {missing_fields}")
                # 尝试填充默认值
                for field in missing_fields:
                    if field == "index":
                        normalized_topic["index"] = str(index + 1)
                    elif field == "date":
                        normalized_topic["date"] = "未指定"
                    elif field == "logic":
                        normalized_topic["logic"] = f"主题{index + 1}的策划逻辑"
                    elif field == "object":
                        normalized_topic["object"] = "None"

            return normalized_topic

        except Exception as e:
            logger.error(f"主题 {index + 1} 验证失败: {e}")
            return None

    async def generate_topics_batch(
        self, 
        materials_list: List[str],
        num_topics_per_batch: int,
        month: str
    ) -> Dict[str, List[Dict[str, Any]]]:
        """
        批量生成主题

        Args:
            materials_list: 创作材料列表
            num_topics_per_batch: 每批生成数量
            month: 月份

        Returns:
            材料索引->主题列表的字典
        """
        results = {}
        
        for i, materials in enumerate(materials_list):
            try:
                logger.info(f"批量生成主题 {i+1}/{len(materials_list)}")
                
                request_id, topics = await self.generate_topics(
                    creative_materials=materials,
                    num_topics=num_topics_per_batch,
                    month=month
                )
                
                results[f"batch_{i+1}"] = topics
                
            except Exception as e:
                logger.error(f"批量生成第 {i+1} 项失败: {e}")
                results[f"batch_{i+1}"] = []
        
        return results

    def get_generation_stats(self) -> Dict[str, Any]:
        """
        获取生成统计信息

        Returns:
            统计信息字典
        """
        return {
            "task_model_config": self.task_model_config,
            "field_config": self.field_config,
            "output_directory": str(self.output_manager.run_output_dir),
            "ai_model_info": self.ai_service.get_model_info(),
            "prompt_templates": self.prompt_manager.get_available_templates().get("topic_generation", {}),
            "json_processor_enabled": self.json_processor.enable_repair
        }

    async def test_generation(self) -> bool:
        """
        测试主题生成功能

        Returns:
            测试是否成功
        """
        try:
            test_materials = """
            选题数量：1
            选题日期：2024-12-01
            
            测试创作材料
            """
            
            _, topics = await self.generate_topics(
                creative_materials=test_materials,
                num_topics=1,
                month="2024-12-01"
            )
            return len(topics) > 0
        except Exception as e:
            logger.error(f"主题生成测试失败: {e}")
            return False 