"""
AI 助手路由 - 题目详情页面的 AI 对话功能

支持工具调用，可以读写文件、执行命令等
"""

from flask import Blueprint, Response, jsonify, request, g, stream_with_context
import json
import logging
import os
import traceback
from pathlib import Path
from app.services.auth.decorators import login_required

logger = logging.getLogger(__name__)

bp_assistant = Blueprint('generate_assistant', __name__)

# 内存缓存对话历史（加速访问）
# 格式: {user_id: {challenge_id: [messages]}}
_conversation_cache = {}


def _get_history_file_path(challenge_id: int) -> Path:
    """获取对话历史文件路径（存储在题目目录下）"""
    from app.models.database.operations import get_challenge_record
    
    challenge = get_challenge_record(challenge_id)
    if not challenge:
        return None
    
    output_dir = challenge.get('output_dir', '')
    if not output_dir:
        return None
    
    # 确保是绝对路径
    if not output_dir.startswith('/'):
        project_root = Path(__file__).parent.parent.parent.parent
        output_dir = str(project_root / output_dir)
    
    # 存储在题目目录下的 chat_history.json
    return Path(output_dir) / 'chat_history.json'


def _load_history_from_file(challenge_id: int) -> list:
    """从文件加载对话历史"""
    try:
        file_path = _get_history_file_path(challenge_id)
        logger.debug(f"加载对话历史文件: {file_path}")
        if file_path and file_path.exists():
            with open(file_path, 'r', encoding='utf-8') as f:
                history = json.load(f)
                logger.info(f"成功加载 {len(history)} 条对话历史 (challenge_id={challenge_id})")
                return history
        else:
            logger.debug(f"对话历史文件不存在: {file_path}")
    except Exception as e:
        logger.error(f"加载对话历史失败: {e}")
        import traceback
        traceback.print_exc()
    return []


def _save_history_to_file(challenge_id: int, history: list):
    """保存对话历史到文件"""
    try:
        file_path = _get_history_file_path(challenge_id)
        if file_path:
            file_path.parent.mkdir(parents=True, exist_ok=True)
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(history, f, ensure_ascii=False, indent=2)
            logger.info(f"保存 {len(history)} 条对话历史到 {file_path}")
    except Exception as e:
        logger.error(f"保存对话历史失败: {e}")
        import traceback
        traceback.print_exc()


def _get_history(user_id: int, challenge_id: int) -> list:
    """获取对话历史（优先从缓存，否则从文件加载）"""
    # 检查缓存
    if user_id not in _conversation_cache:
        _conversation_cache[user_id] = {}
    
    if challenge_id not in _conversation_cache[user_id]:
        # 从文件加载
        _conversation_cache[user_id][challenge_id] = _load_history_from_file(challenge_id)
    
    return _conversation_cache[user_id][challenge_id]


def _add_message(user_id: int, challenge_id: int, role: str, content: str):
    """添加消息到历史（同时更新缓存和文件）"""
    history = _get_history(user_id, challenge_id)
    history.append({'role': role, 'content': content})
    
    # 限制历史长度
    if len(history) > 50:
        history = history[-50:]
        _conversation_cache[user_id][challenge_id] = history
    
    # 保存到文件
    _save_history_to_file(challenge_id, history)


def _clear_history(user_id: int, challenge_id: int):
    """清空对话历史（同时清除缓存和文件）"""
    # 清除缓存
    if user_id in _conversation_cache and challenge_id in _conversation_cache[user_id]:
        _conversation_cache[user_id][challenge_id] = []
    
    # 清除文件
    _save_history_to_file(challenge_id, [])


def _get_challenge_context(challenge_id: int) -> dict:
    """获取题目上下文信息"""
    try:
        from app.models.database.operations import get_challenge_record
        from pathlib import Path
        
        challenge = get_challenge_record(challenge_id)
        if not challenge:
            return None
        
        # 读取 writeup
        output_dir = challenge.get('output_dir', '')
        if output_dir:
            # 如果是相对路径，转换为绝对路径（相对于项目根目录）
            if not output_dir.startswith('/') and not os.path.isabs(output_dir):
                project_root = Path(__file__).parent.parent.parent.parent
                output_dir = str(project_root / output_dir)
        
        context = {
            'name': challenge.get('name', '未命名题目'),
            'description': challenge.get('description', ''),
            'output_dir': output_dir,  # 使用转换后的绝对路径
            'writeup': '',
            'files': []
        }
        
        # 读取 writeup 和文件列表
        if output_dir:
            
            writeup_path = Path(output_dir) / 'writeup.md'
            if writeup_path.exists():
                try:
                    with open(writeup_path, 'r', encoding='utf-8') as f:
                        context['writeup'] = f.read()[:5000]  # 限制长度
                except:
                    pass
            
            # 获取文件列表（从输出目录根目录读取，不限制在docker目录）
            output_path = Path(output_dir)
            if output_path.exists():
                # 递归遍历输出目录下的所有文件
                for f in output_path.rglob('*'):
                    if f.is_file() and not any(part.startswith('.') for part in f.parts):
                        rel_path = str(f.relative_to(output_path))
                        context['files'].append(rel_path)
        
        return context
    except Exception as e:
        logger.error(f"获取题目上下文失败: {e}")
        return None


def _build_system_prompt(challenge_context: dict, abs_output_dir: str = None) -> str:
    """构建系统提示词
    
    Args:
        challenge_context: 题目上下文
        abs_output_dir: 绝对路径的输出目录
    """
    files_str = '\n'.join(f"- {f}" for f in challenge_context.get('files', [])[:20])
    # 使用传入的绝对路径，否则使用上下文中的路径
    output_dir = abs_output_dir or challenge_context.get('output_dir', '')
    docker_dir = f"{output_dir}/docker" if output_dir else ''
    
    return f"""你是一个 CTF 题目助手，帮助用户优化和改进已生成的 CTF 题目。

## 当前题目信息
- 题目名称: {challenge_context.get('name', '未知')}
- 题目描述: {challenge_context.get('description', '无')}
- 题目目录: {docker_dir}

## 题目文件
{files_str}

## 重要：你必须使用工具来完成任务
你有以下工具可用，当用户请求涉及文件操作或命令执行时，你**必须**调用相应的工具：

1. **read_file** - 读取文件内容。参数: path (文件绝对路径)
2. **write_file** - 写入文件。参数: path (文件绝对路径), content (文件内容)
3. **run_command** - 执行命令。参数: command (命令字符串), cwd (工作目录，可选)
4. **list_directory** - 列出目录。参数: path (目录绝对路径)

## 工作流程
- 当用户要求查看文件时，调用 read_file 工具
- 当用户要求修改文件时，先调用 read_file 读取，然后调用 write_file 写入
- 当用户要求执行命令时，调用 run_command 工具
- 不要假装执行工具，必须真正调用工具函数

## 题目目录
所有文件操作的基础路径是: {docker_dir}
"""


def _create_ai_service(user_id: int):
    """创建 AI 服务实例"""
    try:
        from app.services.ai.providers import AIProviderFactory
        
        provider = AIProviderFactory.create_for_user(user_id)
        if provider:
            return provider
        return None
    except Exception as e:
        logger.error(f"创建 AI 服务失败: {e}")
        return None


def _create_tool_executor(output_dir: str):
    """创建工具执行器"""
    try:
        from app.services.ai.tools import ToolExecutor, SandboxConfig
        
        # 创建沙箱配置，限制在题目目录
        docker_dir = str(Path(output_dir) / 'docker') if output_dir else None
        
        # 注意：参数名是 allowed_directories，不是 allowed_paths
        allowed_dirs = [output_dir, docker_dir] if output_dir else []
        # 过滤掉 None 值
        allowed_dirs = [d for d in allowed_dirs if d]
        
        sandbox = SandboxConfig(
            allowed_directories=allowed_dirs,
            max_output_size=10000
        )
        
        print(f"[AI助手] 创建工具执行器: working_dir={docker_dir}, allowed_dirs={allowed_dirs}", flush=True)
        
        return ToolExecutor(sandbox=sandbox, working_dir=docker_dir)
    except Exception as e:
        logger.error(f"创建工具执行器失败: {e}")
        import traceback
        traceback.print_exc()
        return None


def _get_assistant_tools():
    """获取 AI 助手可用的工具"""
    from app.services.ai.providers.base import ToolDefinition
    return [
        ToolDefinition.read_file(),
        ToolDefinition.write_file(),
        ToolDefinition.run_command(),
        ToolDefinition.list_directory(),
    ]


@bp_assistant.route('/challenge/<int:challenge_id>/continue/stream', methods=['POST'])
@login_required
def chat_stream(challenge_id: int):
    """流式对话接口"""
    try:
        user_id = g.user.id
        data = request.get_json() or {}
        instruction = data.get('instruction', '').strip()
        
        if not instruction:
            return jsonify({'success': False, 'message': '请输入指令'}), 400
        
        # 获取题目上下文
        challenge_context = _get_challenge_context(challenge_id)
        if not challenge_context:
            return jsonify({'success': False, 'message': '题目不存在'}), 404
        
        # 权限检查
        from app.models.database.operations import get_challenge_record
        challenge = get_challenge_record(challenge_id)
        if challenge and challenge.get('user_id') != user_id and g.user.role != 'admin':
            return jsonify({'success': False, 'message': '无权访问此题目'}), 403
        
        # 创建 AI 服务
        provider = _create_ai_service(user_id)
        if not provider:
            return jsonify({'success': False, 'message': '未配置 AI 服务，请先在个人中心或 AI 配置页面设置'}), 400
        
        # 添加用户消息到历史
        _add_message(user_id, challenge_id, 'user', instruction)
        
        # 创建工具执行器
        output_dir = challenge_context.get('output_dir', '')
        
        # 确保 output_dir 是绝对路径
        if output_dir and not output_dir.startswith('/'):
            # 相对路径，转换为绝对路径（相对于项目根目录）
            import os
            project_root = Path(__file__).parent.parent.parent.parent
            output_dir = str(project_root / output_dir)
        
        print(f"[AI助手] output_dir (绝对路径): {output_dir}", flush=True)
        
        tool_executor = _create_tool_executor(output_dir)
        tools = _get_assistant_tools() if tool_executor else None
        
        print(f"[AI助手] tool_executor 创建结果: {tool_executor is not None}", flush=True)
        print(f"[AI助手] tools 数量: {len(tools) if tools else 0}", flush=True)
        if tools:
            print(f"[AI助手] 工具列表: {[t['function']['name'] for t in tools]}", flush=True)
        
        def generate():
            """生成器函数，流式输出，支持工具调用"""
            try:
                # 构建消息列表
                messages = []
                
                # 系统提示词（传入绝对路径）
                system_prompt = _build_system_prompt(challenge_context, output_dir)
                messages.append({'role': 'system', 'content': system_prompt})
                
                print(f"\n{'='*60}", flush=True)
                print(f"[AI助手] 开始对话", flush=True)
                print(f"[AI助手] 题目目录: {output_dir}", flush=True)
                print(f"[AI助手] 工具数量: {len(tools) if tools else 0}", flush=True)
                print(f"[AI助手] 工具执行器: {tool_executor}", flush=True)
                print(f"[AI助手] Provider类型: {type(provider).__name__}", flush=True)
                print(f"{'='*60}", flush=True)
                
                # 添加 writeup 作为上下文（简化版本）
                if challenge_context.get('writeup'):
                    # 只添加简短的上下文，不要太长
                    writeup_summary = challenge_context['writeup'][:1500]
                    messages.append({
                        'role': 'user', 
                        'content': f"题目 Writeup 摘要：\n{writeup_summary}"
                    })
                    messages.append({
                        'role': 'assistant',
                        'content': '好的，我已了解题目信息。'
                    })
                
                # 添加对话历史
                history = _get_history(user_id, challenge_id)
                for msg in history[-10:]:  # 只保留最近10条
                    messages.append(msg)
                
                # 工具调用循环（不限制次数）
                max_iterations = 100
                iteration = 0
                full_content = ""
                
                # 检查是否是 CLI 模式（anyrouter/agentrouter）
                is_cli_mode = hasattr(provider, 'router_type') or hasattr(provider, '_call_claude_chat')
                docker_dir = f"{output_dir}/docker" if output_dir else None
                
                print(f"[AI助手] CLI模式: {is_cli_mode}", flush=True)
                print(f"[AI助手] Docker目录: {docker_dir}", flush=True)
                
                while iteration < max_iterations:
                    iteration += 1
                    
                    print(f"\n[AI助手] === 轮次 {iteration} ===", flush=True)
                    print(f"[AI助手] 消息数量: {len(messages)}", flush=True)
                    print(f"[AI助手] 最后一条消息: {messages[-1]['role']} - {messages[-1].get('content', '')[:100]}...", flush=True)
                    
                    # 调用 AI（带工具）
                    if is_cli_mode and docker_dir:
                        # CLI 模式：传递工作目录，让 Claude 自己执行工具
                        print(f"[AI助手] 使用 CLI 模式调用，工作目录: {docker_dir}", flush=True)
                        response = provider.chat(messages, tools=tools, workspace_dir=docker_dir)
                    else:
                        # API 模式：正常调用
                        print(f"[AI助手] 使用 API 模式调用，工具: {[t['function']['name'] for t in tools] if tools else []}", flush=True)
                        response = provider.chat(messages, tools=tools)
                    
                    print(f"[AI助手] AI 响应:", flush=True)
                    print(f"  - tool_calls: {response.get('tool_calls')}", flush=True)
                    print(f"  - content长度: {len(response.get('content', ''))}", flush=True)
                    print(f"  - content预览: {response.get('content', '')[:200]}...", flush=True)
                    
                    # CLI 模式：Claude 会自己执行工具，直接返回结果
                    if is_cli_mode:
                        content = response.get('content', '')
                        full_content = content
                        
                        # 分块发送
                        chunk_size = 100
                        for i in range(0, len(content), chunk_size):
                            chunk = content[i:i+chunk_size]
                            yield f"data: {json.dumps({'type': 'chunk', 'content': chunk}, ensure_ascii=False)}\n\n"
                        
                        # 保存响应到历史
                        _add_message(user_id, challenge_id, 'assistant', full_content)
                        yield f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"
                        print(f"[AI助手] CLI模式完成", flush=True)
                        break
                    
                    # API 模式：检查是否有工具调用
                    tool_calls = response.get('tool_calls')
                    print(f"[AI助手] 检查工具调用: tool_calls={tool_calls}, tool_executor={tool_executor is not None}", flush=True)
                    
                    if tool_calls and tool_executor:
                        print(f"[AI助手] 检测到工具调用: {len(tool_calls)} 个", flush=True)
                        for i, tc in enumerate(tool_calls):
                            print(f"[AI助手]   工具 {i+1}: {tc.get('function', {}).get('name', 'unknown')}", flush=True)
                            print(f"[AI助手]   参数: {tc.get('function', {}).get('arguments', '{}')[:200]}", flush=True)
                        
                        content = response.get('content', '')
                        
                        # 发送 AI 的思考过程
                        if content:
                            full_content += content + "\n"
                            newline = '\n'
                            yield f"data: {json.dumps({'type': 'chunk', 'content': content + newline}, ensure_ascii=False)}\n\n"
                        
                        # 添加助手消息（包含工具调用）
                        assistant_msg = {
                            'role': 'assistant',
                            'content': content,
                            'tool_calls': tool_calls
                        }
                        messages.append(assistant_msg)
                        
                        # 执行每个工具调用
                        for tool_call in tool_calls:
                            tool_name = tool_call['function']['name']
                            tool_args = tool_call.get('function', {}).get('arguments', '{}')
                            
                            print(f"[AI助手] 执行工具: {tool_name}", flush=True)
                            print(f"[AI助手] 参数: {tool_args[:300]}", flush=True)

                            # 发送工具调用通知
                            tool_msg = f'\n🔧 执行工具: {tool_name}\n'
                            yield f"data: {json.dumps({'type': 'chunk', 'content': tool_msg}, ensure_ascii=False)}\n\n"
                            full_content += tool_msg
                            
                            # 执行工具
                            tool_result = tool_executor.execute_tool_call(tool_call)
                            
                            print(f"[AI助手] 工具结果: {tool_result[:300]}...", flush=True)
                            
                            # 发送工具结果摘要
                            result_preview = tool_result[:200] + '...' if len(tool_result) > 200 else tool_result
                            result_msg = f'📋 结果: {result_preview}\n'
                            yield f"data: {json.dumps({'type': 'chunk', 'content': result_msg}, ensure_ascii=False)}\n\n"
                            full_content += result_msg
                            
                            # 添加工具结果消息
                            messages.append({
                                'role': 'tool',
                                'tool_call_id': tool_call['id'],
                                'name': tool_name,
                                'content': tool_result
                            })
                        
                        # 继续循环，让 AI 处理工具结果
                        print(f"[AI助手] 工具执行完成，继续下一轮", flush=True)
                        continue
                    else:
                        # 没有工具调用，返回最终响应
                        print(f"[AI助手] 没有工具调用，返回最终响应", flush=True)
                        content = response.get('content', '')
                        full_content += content
                        
                        # 分块发送
                        chunk_size = 100
                        for i in range(0, len(content), chunk_size):
                            chunk = content[i:i+chunk_size]
                            yield f"data: {json.dumps({'type': 'chunk', 'content': chunk}, ensure_ascii=False)}\n\n"
                        
                        # 保存完整响应到历史
                        _add_message(user_id, challenge_id, 'assistant', full_content)
                        yield f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"
                        print(f"[AI助手] 对话完成", flush=True)
                        break
                
                # 达到最大迭代次数
                if iteration >= max_iterations:
                    warning_msg = '\n⚠️ 达到最大工具调用次数限制'
                    yield f"data: {json.dumps({'type': 'chunk', 'content': warning_msg}, ensure_ascii=False)}\n\n"
                    _add_message(user_id, challenge_id, 'assistant', full_content)
                    yield f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"
                    
            except Exception as e:
                logger.error(f"AI 对话失败: {e}")
                traceback.print_exc()
                yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n"
        
        return Response(
            stream_with_context(generate()),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'X-Accel-Buffering': 'no',
                'Connection': 'keep-alive'
            }
        )
        
    except Exception as e:
        logger.error(f"处理对话请求失败: {e}")
        traceback.print_exc()
        return jsonify({'success': False, 'message': str(e)}), 500


@bp_assistant.route('/challenge/<int:challenge_id>/continue/history', methods=['GET'])
@login_required
def get_history(challenge_id: int):
    """获取对话历史"""
    try:
        user_id = g.user.id
        
        # 权限检查
        from app.models.database.operations import get_challenge_record
        challenge = get_challenge_record(challenge_id)
        if challenge and challenge.get('user_id') != user_id and g.user.role != 'admin':
            return jsonify({'success': False, 'message': '无权访问此题目'}), 403
        
        history = _get_history(user_id, challenge_id)
        
        return jsonify({
            'success': True,
            'history': history
        })
        
    except Exception as e:
        logger.error(f"获取对话历史失败: {e}")
        return jsonify({'success': False, 'message': str(e)}), 500


@bp_assistant.route('/challenge/<int:challenge_id>/continue/clear', methods=['POST'])
@login_required
def clear_history(challenge_id: int):
    """清空对话历史"""
    try:
        user_id = g.user.id
        
        # 权限检查
        from app.models.database.operations import get_challenge_record
        challenge = get_challenge_record(challenge_id)
        if challenge and challenge.get('user_id') != user_id and g.user.role != 'admin':
            return jsonify({'success': False, 'message': '无权访问此题目'}), 403
        
        _clear_history(user_id, challenge_id)
        
        return jsonify({
            'success': True,
            'message': '对话历史已清空'
        })
        
    except Exception as e:
        logger.error(f"清空对话历史失败: {e}")
        return jsonify({'success': False, 'message': str(e)}), 500