# -*- coding: utf-8 -*-
"""
题目生成服务

负责题目生成的核心业务逻辑
"""
import os
import datetime
import traceback
import threading
import logging
from typing import Dict, Any, Optional, List
from pathlib import Path
from flask import current_app

from app.core.exceptions import GenerationError, ValidationError, ConfigurationError


class ChallengeGeneratorService:
    """题目生成服务类"""
    
    def __init__(self, user_id: int, category_id: str, task_id: Optional[str] = None):
        """
        Args:
            user_id: 用户ID
            category_id: 方向ID
            task_id: 任务ID（可选）
        """
        self.user_id = user_id
        self.category_id = category_id
        self.task_id = task_id
        self.logger = logging.getLogger(__name__)
    
    def generate(
        self,
        language: Optional[str],
        vulnerabilities: Optional[List],
        scene: Optional[Dict[str, Any]],
        difficulty: str,
        extra_requirements: str = '',
        ai_mode: str = 'augment',
        app=None,
        form_data: Optional[Dict[str, Any]] = None,
        ai_config_id: Optional[int] = None
    ) -> Dict[str, Any]:
        """生成题目
        
        Args:
            language: 编程语言
            vulnerabilities: 漏洞列表
            scene: 场景
            difficulty: 难度
            extra_requirements: 额外要求
            ai_mode: AI 模式
            app: Flask 应用对象
            form_data: 完整的表单数据（可选，用于传递所有表单字段）
            
        Returns:
            {
                'status': 'success' | 'error',
                'message': str,
                'challenge_id': Optional[int],
                'output_dir': Optional[str],
                'log_file': Optional[str]
            }
        """
        if app is None:
            app = current_app._get_current_object()
        
        with app.app_context():
            try:
                self.logger.info("生成题目主进程开始")
                self.logger.info(f"用户ID: {self.user_id}")
                
                # 重置日志ID
                from app.routes.generator.utils import reset_log_id, set_current_user_id
                reset_log_id()
                set_current_user_id(self.user_id)
                
                # 根据 AI 模式选择生成方式
                result = self._execute_generation(
                    language=language,
                    vulnerabilities=vulnerabilities,
                    scene=scene,
                    difficulty=difficulty,
                    extra_requirements=extra_requirements,
                    ai_mode=ai_mode,
                    form_data=form_data,
                    ai_config_id=ai_config_id
                )
                
                if result["status"] == "success":
                    return self._handle_success(result)
                else:
                    return self._handle_failure(result)
                    
            except Exception as e:
                self.logger.error(f"生成进程中发生错误: {str(e)}", exc_info=True)
                return self._handle_exception(e)
    
    def _execute_generation(
        self,
        language: Optional[str],
        vulnerabilities: Optional[List],
        scene: Optional[Dict[str, Any]],
        difficulty: str,
        extra_requirements: str,
        ai_mode: str,
        form_data: Optional[Dict[str, Any]] = None,
        ai_config_id: Optional[int] = None
    ) -> Dict[str, Any]:
        """执行生成逻辑"""
        if ai_mode == 'api':
            return self._generate_with_api(
                language, vulnerabilities, scene, difficulty, extra_requirements, form_data, ai_config_id
            )
        elif ai_mode in ('anyrouter', 'agentrouter'):
            return self._generate_with_claude_router(
                ai_mode, language, vulnerabilities, scene, difficulty, extra_requirements, form_data
            )
        else:
            return self._generate_with_augment(
                language, vulnerabilities, scene, difficulty, extra_requirements, form_data
            )
    
    def _generate_with_api(
        self,
        language: Optional[str],
        vulnerabilities: Optional[List],
        scene: Optional[Dict[str, Any]],
        difficulty: str,
        extra_requirements: str,
        form_data: Optional[Dict[str, Any]] = None,
        ai_config_id: Optional[int] = None
    ) -> Dict[str, Any]:
        """使用 API 模式生成"""
        self.logger.info("使用 API 模式生成题目")
        if ai_config_id:
            self.logger.info(f"使用指定的 AI 配置 ID: {ai_config_id}")
        from app.services.ai.service import AIService
        
        # 获取输出目录
        ctf_dir = Path(__file__).parent.parent.parent.parent
        category_output_dir = ctf_dir / 'ge10' / self.category_id / 'output'
        category_output_dir.mkdir(parents=True, exist_ok=True)
        
        ai_service = AIService(
            user_id=self.user_id,
            output_dir=str(category_output_dir),
            log_callback=lambda level, msg: self.logger.log(
                getattr(logging, level.upper(), logging.INFO), msg
            ),
            category_id=self.category_id,
            ai_config_id=ai_config_id
        )
        
        if not ai_service.is_available():
            return {
                "status": "error",
                "message": "AI 服务不可用，请检查 AI 配置"
            }
        
        provider_info = ai_service.get_provider_info()
        self.logger.info(f"已创建 AI 服务实例: {provider_info['name']} ({provider_info['model']})")
        
        return ai_service.generate_ctf_challenge(
            language=language,
            vulnerabilities=vulnerabilities,
            scene=scene,
            difficulty=difficulty,
            extra_requirements=extra_requirements,
            category_id=self.category_id,
            form_data=form_data or {},
            task_id=self.task_id
        )
    
    def _generate_with_claude_router(
        self,
        router_type: str,
        language: Optional[str],
        vulnerabilities: Optional[List],
        scene: Optional[Dict[str, Any]],
        difficulty: str,
        extra_requirements: str,
        form_data: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """使用 Claude Router CLI 模式生成"""
        self.logger.info(f"使用 {router_type} CLI 模式生成题目")
        from app.services.ai.core.claude_router_service import ClaudeRouterService
        
        verbose_mode = os.environ.get('CLAUDE_VERBOSE', 'false').lower() == 'true'
        
        claude_service = ClaudeRouterService(
            router_type=router_type,
            user_id=self.user_id,
            verbose=verbose_mode
        )
        self.logger.info(f"已创建 {router_type} 服务实例，用户ID: {self.user_id}")
        
        return claude_service.generate_ctf_challenge(
            language=language,
            vulnerabilities=vulnerabilities,
            scene=scene,
            difficulty=difficulty,
            extra_requirements=extra_requirements,
            category_id=self.category_id,
            form_data=form_data,
            task_id=self.task_id
        )
    
    def _generate_with_augment(
        self,
        language: Optional[str],
        vulnerabilities: Optional[List],
        scene: Optional[Dict[str, Any]],
        difficulty: str,
        extra_requirements: str,
        form_data: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """使用 Augment CLI 模式生成"""
        self.logger.info("使用 Augment CLI 模式生成题目")
        from app.services.ai.core.augment_service import AugmentService
        
        verbose_mode = os.environ.get('AUGMENT_VERBOSE', 'false').lower() == 'true'
        
        augment_service = AugmentService(user_id=self.user_id, verbose=verbose_mode)
        self.logger.info(f"已创建 Augment 服务实例，用户ID: {self.user_id}，详细模式: {verbose_mode}")
        
        result = augment_service.generate_ctf_challenge(
            language=language,
            vulnerabilities=vulnerabilities,
            scene=scene,
            difficulty=difficulty,
            extra_requirements=extra_requirements,
            category_id=self.category_id,
            form_data=form_data,
            task_id=self.task_id
        )
        
        # 保存 augment_service 用于后续检查超时恢复
        result['_augment_service'] = augment_service
        
        return result
    
    def _handle_success(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """处理生成成功的情况"""
        self.logger.info("题目生成成功!")
        
        # 查找输出目录
        output_dir, challenge_info, estimated_time = self._find_output_directory()
        
        # 获取日志文件路径
        log_file = self._get_log_file_path()
        
        # 移动日志文件到题目目录
        if log_file and output_dir:
            try:
                from app.services.storage import LogManager
                log_manager = LogManager()
                moved_log_file = log_manager.move_log_to_challenge(log_file, output_dir)
                if moved_log_file:
                    log_file = moved_log_file
                    self.logger.info(f"日志文件已移动到题目目录: {moved_log_file}")
                    # 更新 generation_statuses 中的日志文件路径
                    if self.task_id:
                        from app.routes.generator.utils import generation_statuses, generation_lock
                        with generation_lock:
                            if self.task_id in generation_statuses:
                                generation_statuses[self.task_id]['log_file'] = moved_log_file
            except Exception as e:
                self.logger.warning(f"移动日志文件到题目目录失败: {e}")
        
        # 检查超时恢复（仅 Augment 模式）
        augment_service = result.get('_augment_service')
        generation_status_value = 'completed'
        generation_warning = None
        if augment_service and hasattr(augment_service, 'timeout_recovered') and augment_service.timeout_recovered:
            generation_status_value = 'timeout_recovered'
            generation_warning = augment_service.timeout_warning
            self.logger.warning(f"检测到超时恢复: {generation_warning}")
        
        # 创建题目记录
        challenge_id = self._create_challenge_record(
            output_dir=output_dir,
            log_file=log_file,
            challenge_info=challenge_info,
            estimated_time=estimated_time,
            generation_status=generation_status_value,
            generation_warning=generation_warning
        )
        
        # 移动会话历史（仅 Augment 模式）
        if output_dir and augment_service:
            try:
                augment_service._move_session_to_challenge_dir(output_dir)
            except Exception as e:
                self.logger.warning(f"移动会话历史失败: {e}")
        
        return {
            'status': 'success',
            'challenge_id': challenge_id,
            'output_dir': output_dir,
            'log_file': log_file,
            'challenge_info': challenge_info
        }
    
    def _handle_failure(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """处理生成失败的情况"""
        error_msg = result.get('message') or result.get('error') or '未知错误'
        self.logger.error(f"生成失败: {error_msg}")
        
        # 移动日志文件到失败目录
        log_file = self._get_log_file_path()
        if log_file:
            try:
                from app.services.storage import LogManager
                log_manager = LogManager()
                moved_log_file = log_manager.move_log_to_failed(log_file, self.task_id)
                if moved_log_file:
                    self.logger.info(f"日志文件已移动到失败目录: {moved_log_file}")
            except Exception as e:
                self.logger.warning(f"移动日志文件到失败目录失败: {e}")
        
        # 更新任务状态
        if self.task_id:
            from app.routes.generator.tasks import update_task
            update_task(self.task_id, status='failed', error=error_msg)
        
        return {
            'status': 'error',
            'message': error_msg
        }
    
    def _handle_exception(self, exception: Exception) -> Dict[str, Any]:
        """处理异常情况"""
        error_msg = str(exception)
        self.logger.error(f"生成出错: {error_msg}", exc_info=True)
        
        # 移动日志文件到失败目录
        log_file = self._get_log_file_path()
        if log_file:
            try:
                from app.services.storage import LogManager
                log_manager = LogManager()
                moved_log_file = log_manager.move_log_to_failed(log_file, self.task_id)
                if moved_log_file:
                    self.logger.info(f"日志文件已移动到失败目录: {moved_log_file}")
            except Exception as e:
                self.logger.warning(f"移动日志文件到失败目录失败: {e}")
        
        # 更新任务状态
        if self.task_id:
            from app.routes.generator.tasks import update_task
            update_task(self.task_id, status='failed', error=error_msg)
        
        return {
            'status': 'error',
            'message': f"生成出错: {error_msg}"
        }
    
    def _find_output_directory(self) -> tuple:
        """查找输出目录"""
        from pathlib import Path
        from app.routes.generator.core import _extract_challenge_info_from_writeup
        
        output_dir = None
        challenge_info = {}
        estimated_time = None
        
        self.logger.info("查找输出目录...")
        ctf_dir = Path(__file__).parent.parent.parent.parent
        category_output_dir = ctf_dir / "ge10" / self.category_id / "output"
        
        if category_output_dir.exists():
            output_dirs = sorted(
                [x for x in category_output_dir.glob("*") if x.is_dir()],
                key=lambda x: x.stat().st_mtime,
                reverse=True
            )
            if output_dirs:
                output_dir = str(output_dirs[0])
                self.logger.info(f"找到输出目录: {output_dir}")
                
                # 从目录名提取题目名称
                import re
                dir_name = output_dirs[0].name
                name_match = re.search(r'\d{8}_\d{6}_(.+)', dir_name)
                if name_match:
                    challenge_name = name_match.group(1)
                    challenge_info['name'] = challenge_name
                    self.logger.info(f"从目录名提取题目名称: {challenge_name}")
                
                # 从 writeup.md 提取信息（优先使用writeup中的信息）
                writeup_file = output_dirs[0] / "writeup.md"
                if writeup_file.exists():
                    with open(writeup_file, 'r', encoding='utf-8') as f:
                        writeup_content = f.read()
                    writeup_info = _extract_challenge_info_from_writeup(writeup_content, dir_name)
                    # writeup中的题目名优先于目录名
                    if writeup_info.get('name') and writeup_info.get('name') != '未命名题目':
                        challenge_info['name'] = writeup_info.get('name')
                    if 'description' not in challenge_info or not challenge_info.get('description'):
                        challenge_info['description'] = writeup_info.get('description', '')
                    if 'estimated_time' not in challenge_info or not challenge_info.get('estimated_time'):
                        challenge_info['estimated_time'] = writeup_info.get('estimated_time', '1-2小时')
                    estimated_time = challenge_info.get('estimated_time', '1-2小时')
        
        return output_dir, challenge_info, estimated_time
    
    def _get_log_file_path(self) -> Optional[str]:
        """获取日志文件路径"""
        from app.routes.generator.utils import get_generation_status
        
        if not self.task_id:
            return None
        
        task_status = get_generation_status(self.task_id)
        if task_status and task_status.get('log_file'):
            return task_status['log_file']
        return None
    
    def _check_deployable(self, output_dir: Optional[str]) -> Optional[bool]:
        """检查题目是否可部署"""
        if not output_dir:
            return None
        
        try:
            from app.services.deployment.deployment_validator import check_challenge_deployable
            return check_challenge_deployable(output_dir)
        except Exception as e:
            self.logger.warning(f"检测部署能力失败: {str(e)}")
            return None
    
    def _create_challenge_record(
        self,
        output_dir: Optional[str],
        log_file: Optional[str],
        challenge_info: Dict[str, Any],
        estimated_time: Optional[str],
        generation_status: str,
        generation_warning: Optional[str]
    ) -> Optional[int]:
        """创建题目记录"""
        try:
            from app.models.database.operations import save_challenge_record
            
            # 准备数据
            generation_result = {
                'output_dir': output_dir,
                'log_file': log_file,
                'challenge_info': challenge_info,
                'estimated_time': estimated_time or '1-2小时',
                'generation_status': generation_status,
                'generation_warning': generation_warning,
                'timestamp': datetime.datetime.now()
            }
            
            # 从 challenge_info 中提取信息
            if challenge_info:
                # 优先使用从目录名或writeup提取的题目名
                challenge_name = challenge_info.get('name')
                if not challenge_name or challenge_name == '未命名题目':
                    # 如果还没有题目名，尝试从output_dir提取
                    if output_dir:
                        import re
                        from pathlib import Path
                        dir_name = Path(output_dir).name
                        name_match = re.search(r'\d{8}_\d{6}_(.+)', dir_name)
                        if name_match:
                            challenge_name = name_match.group(1)
                
                generation_result['name'] = challenge_name or f"CTF题目-{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}"
                generation_result['description'] = challenge_info.get('description', '')
            
            self.logger.info(f"准备创建题目记录: name={generation_result.get('name')}")
            
            # 检测是否可部署（在创建记录之前）
            deployable = self._check_deployable(output_dir)
            if deployable is not None:
                generation_result['deployable'] = deployable
                self.logger.info(f"检测部署能力: deployable={deployable}")
            
            # 创建题目记录
            challenge_id = save_challenge_record(generation_result, self.user_id)
            
            if challenge_id:
                self.logger.info(f"已成功创建题目记录，ID: {challenge_id}")
                
                # 更新任务状态和名称（使用题目文件夹名）
                if self.task_id:
                    from app.routes.generator.tasks import update_task
                    
                    # 从输出目录提取题目名称（时间戳后面的部分）
                    challenge_name = None
                    if output_dir:
                        import re
                        from pathlib import Path
                        dir_name = Path(output_dir).name
                        name_match = re.search(r'\d{8}_\d{6}_(.+)', dir_name)
                        if name_match:
                            challenge_name = name_match.group(1)
                    
                    # 如果提取到了题目名，更新任务名称
                    if challenge_name:
                        update_task(
                            self.task_id,
                            challenge_id=challenge_id,
                            status='completed',
                            progress=100,
                            name=challenge_name
                        )
                        self.logger.info(f"已更新任务名称为: {challenge_name}")
                    else:
                        # 如果没有提取到，使用默认名称
                        update_task(
                            self.task_id,
                            challenge_id=challenge_id,
                            status='completed',
                            progress=100
                        )
                
                # 更新对话日志关联
                from app.routes.generator.utils import update_conversation_logs_challenge
                update_conversation_logs_challenge(challenge_id)
                
                return challenge_id
            else:
                self.logger.error("创建题目记录失败")
                return None
                
        except Exception as e:
            self.logger.error(f"创建题目记录失败: {str(e)}", exc_info=True)
            return None