#  @line_count 500
"""测试专家系统 - Streamlit Web应用"""
import streamlit as st
import json
import os
from pathlib import Path
from typing import List, Dict, Any, Optional
import yaml

# 导入自定义模块
from modules.json_parser import JSONParser
from modules.api_client import APIClient
from modules.prompt_manager import PromptManager
from modules.test_generator import TestGenerator
from modules.output_formatter import OutputFormatter
from process_doc_file import convert_docx_to_json


# 页面配置
st.set_page_config(
    page_title="测试专家系统",
    page_icon="🧪",
    layout="wide",
    initial_sidebar_state="expanded"
)


def init_session_state():
    """初始化会话状态"""
    if 'api_client' not in st.session_state:
        st.session_state.api_client = None
    if 'prompt_manager' not in st.session_state:
        st.session_state.prompt_manager = None
    if 'json_parser' not in st.session_state:
        st.session_state.json_parser = None
    if 'test_items' not in st.session_state:
        st.session_state.test_items = []
    if 'test_cases' not in st.session_state:
        st.session_state.test_cases = []
    if 'function_points' not in st.session_state:
        st.session_state.function_points = []
    if 'document_info' not in st.session_state:
        st.session_state.document_info = {}
    if 'standard_manager' not in st.session_state:
        st.session_state.standard_manager = None
    # 用户自定义Prompt支持
    if 'custom_prompts' not in st.session_state:
        st.session_state.custom_prompts = {}  # {func_point_index: custom_prompt}
    if 'custom_prompt_template' not in st.session_state:
        st.session_state.custom_prompt_template = None  # 批量模板
    if 'use_custom_prompt' not in st.session_state:
        st.session_state.use_custom_prompt = False  # 是否启用自定义prompt


def load_config():
    """加载配置"""
    try:
        if st.session_state.api_client is None:
            st.session_state.api_client = APIClient()
        if st.session_state.prompt_manager is None:
            st.session_state.prompt_manager = PromptManager()
    except Exception as e:
        st.error(f"配置加载失败: {str(e)}")


def show_home_page():
    """显示首页"""
    st.title("🧪 测试专家系统")
    st.markdown("---")
    
    st.markdown("""
    ### 欢迎使用测试专家系统
    
    本系统基于大模型AI技术，能够自动从软件使用说明文档中提取功能点，并生成全面的测试项和测试用例。
    
    #### 主要功能
    
    1. **文档解析**: 自动解析JSON格式的软件使用说明书，提取功能点和操作步骤
    2. **智能生成**: 基于大模型API，为每个功能点生成专业的测试项和测试用例
    3. **多格式输出**: 支持JSON、Markdown、Excel多种格式导出
    4. **Prompt优化**: 支持自定义和优化测试生成Prompt模板
    
    #### 使用流程
    
    1. **配置API**: 在"配置"页面设置大模型API密钥和选择模型
    2. **上传文档**: 在"生成测试"页面上传JSON格式的文档
    3. **生成测试**: 选择功能点，点击生成按钮
    4. **下载结果**: 预览生成结果，下载所需格式的文件
    
    #### 支持的API提供商
    
    - DeepSeek (默认)
    - 通义千问
    - OpenAI/ChatGPT
    
    ---
    
    **开始使用**: 请先前往"配置"页面设置API密钥
    """)


def show_config_page():
    """显示配置页面"""
    st.title("⚙️ 系统配置")
    st.markdown("---")
    
    load_config()
    
    # API配置
    st.header("📡 API配置")
    
    api_client = st.session_state.api_client
    providers = api_client.config.get('providers', {})
    current_provider = api_client.current_provider
    
    # 选择API提供商
    provider_options = list(providers.keys())
    selected_provider = st.selectbox(
        "选择API提供商",
        provider_options,
        index=provider_options.index(current_provider) if current_provider in provider_options else 0
    )
    
    if selected_provider != current_provider:
        api_client.set_provider(selected_provider)
        st.success(f"已切换到: {selected_provider}")
    
    # 配置API密钥
    provider_config = providers[selected_provider]
    st.subheader(f"{selected_provider.upper()} 配置")
    
    api_key = st.text_input(
        "API密钥",
        value=provider_config.get('api_key', '') or os.getenv(f'{selected_provider.upper()}_API_KEY', ''),
        type="password",
        help="输入API密钥，或通过环境变量设置"
    )
    
    if st.button("保存API密钥"):
        try:
            api_client.update_api_key(selected_provider, api_key)
            st.success("API密钥已保存")
        except Exception as e:
            st.error(f"保存失败: {str(e)}")
    
    # 显示其他配置信息
    with st.expander("查看配置详情"):
        st.json(provider_config)
    
    # 提示信息
    st.info("💡 提示：Prompt和测试规范配置将在上传文档后进行设置")


def show_generate_page():
    """显示生成页面"""
    st.title("🚀 生成测试")
    st.markdown("---")
    
    load_config()
    
    # 上传文档（支持 JSON 和 Word）
    st.header("📄 上传文档")
    uploaded_file = st.file_uploader(
        "选择文档文件（支持 JSON 或 Word）",
        type=['json', 'docx'],
        help="可以直接上传Word使用说明书（.docx），系统会自动转换为JSON；也可以上传已转换好的JSON文件"
    )
    
    if uploaded_file is not None:
        try:
            # 保存上传的文件
            upload_dir = Path("uploads")
            upload_dir.mkdir(exist_ok=True)
            file_path = upload_dir / uploaded_file.name
            
            with open(file_path, 'wb') as f:
                f.write(uploaded_file.getbuffer())
            
            # 如果是Word文档，先转换为JSON
            if file_path.suffix.lower() == ".docx":
                json_path = upload_dir / f"{file_path.stem}.json"
                convert_docx_to_json(str(file_path), str(json_path))
                parser_target_path = json_path
            else:
                # 直接使用上传的JSON文件
                parser_target_path = file_path
            
            # 解析JSON
            parser = JSONParser(str(parser_target_path))
            st.session_state.json_parser = parser
            
            # 只有在上传新文档时才清空自定义prompts（检查路径是否变化）
            current_doc_path = str(parser_target_path)
            previous_doc_path = st.session_state.get('last_uploaded_doc_path', '')
            
            if current_doc_path != previous_doc_path:
                # 新文档，清空之前的自定义prompts
                st.session_state.custom_prompts.clear()
                st.session_state.custom_prompt_template = None
                st.session_state.use_custom_prompt = False
                st.session_state.last_uploaded_doc_path = current_doc_path
            
            # 获取文档信息
            doc_info = parser.get_document_info()
            st.session_state.document_info = doc_info
            
            # 显示文档信息
            st.success("文档上传成功！")
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("文档标题", doc_info.get('title', 'N/A'))
            with col2:
                st.metric("章节数量", doc_info.get('section_count', 0))
            with col3:
                st.metric("版本", doc_info.get('version', 'N/A'))
            
            # 提取功能点
            function_points = parser.extract_function_points()
            st.session_state.function_points = function_points
            
            # 显示功能点
            if function_points:
                # ========== 新增：Prompt/测试规范配置区域 ==========
                st.markdown("---")
                st.header("📝 Prompt/测试规范配置")
                
                # 1. 需求类型统计
                requirement_types = {}
                for fp in function_points:
                    req_type = fp.get('requirement_type', '未知')
                    requirement_types[req_type] = requirement_types.get(req_type, 0) + 1
                
                if requirement_types:
                    st.subheader("需求类型统计")
                    cols = st.columns(len(requirement_types))
                    for idx, (req_type, count) in enumerate(requirement_types.items()):
                        with cols[idx]:
                            st.metric(req_type, count)
                
                # 2. Prompt策略选择
                use_standards = st.checkbox(
                    "使用测试规范（推荐）",
                    value=True,
                    help="启用后，系统会根据需求类型自动选择适用的测试规范，生成符合行业标准的测试用例"
                )
                
                # 3. 测试规范选择详情（如果启用）
                if use_standards:
                    try:
                        from modules.test_standard_manager import TestStandardManager
                        
                        # 初始化测试规范管理器
                        if st.session_state.standard_manager is None:
                            st.session_state.standard_manager = TestStandardManager(
                                api_client=st.session_state.api_client
                            )
                        
                        standard_manager = st.session_state.standard_manager
                        
                        # 显示测试规范选择详情
                        with st.expander("📋 查看测试规范选择详情"):
                            for fp in function_points:
                                requirement = _convert_to_requirement(fp)
                                standards = standard_manager.get_applicable_standards(
                                    requirement, 
                                    use_ai=False  # 先不使用AI，避免额外调用
                                )
                                
                                req_type = requirement.get('requirement_type', '未知')
                                req_id = requirement.get('requirement_id', fp.get('function_name', ''))
                                
                                st.markdown(f"**{req_id}** ({req_type})")
                                if standards:
                                    st.write(f"选择的测试规范: {', '.join(standards)}")
                                else:
                                    st.write("⚠️ 未找到适用的测试规范")
                                st.markdown("---")
                        
                        # 4. Prompt预览与编辑
                        preview_prompt = st.checkbox("预览规范化Prompt")
                        if preview_prompt and function_points:
                            sample_fp = function_points[0]
                            requirement = _convert_to_requirement(sample_fp)
                            standards = standard_manager.get_applicable_standards(requirement, use_ai=False)
                            
                            if standards:
                                # 生成默认prompt
                                default_prompt = standard_manager.build_prompt(requirement, standards)
                                
                                # 获取当前索引（第一个功能点）
                                sample_idx = 0
                                
                                # 如果没有自定义prompt，使用默认值
                                if sample_idx not in st.session_state.custom_prompts:
                                    st.session_state.custom_prompts[sample_idx] = default_prompt
                                
                                # 可编辑的Prompt文本框
                                edited_prompt = st.text_area(
                                    "Prompt预览（可编辑，以第一个功能点为例）",
                                    value=st.session_state.custom_prompts[sample_idx],
                                    height=400,
                                    key="prompt_editor_preview",
                                    help="您可以直接修改Prompt内容，修改后的内容将在生成测试时使用"
                                )
                                
                                # 更新session state
                                st.session_state.custom_prompts[sample_idx] = edited_prompt
                                
                                # 操作按钮
                                col1, col2, col3 = st.columns(3)
                                with col1:
                                    if st.button("🔄 恢复默认", key="restore_default_prompt"):
                                        st.session_state.custom_prompts[sample_idx] = default_prompt
                                        st.session_state.use_custom_prompt = False
                                        st.success("已恢复为默认Prompt")
                                        st.rerun()
                                
                                with col2:
                                    if st.button("📋 应用到所有功能点", key="apply_to_all"):
                                        # 将当前编辑的prompt应用到所有功能点
                                        for idx in range(len(function_points)):
                                            st.session_state.custom_prompts[idx] = edited_prompt
                                        st.session_state.use_custom_prompt = True
                                        st.success(f"✅ 已应用到 {len(function_points)} 个功能点")
                                
                                with col3:
                                    if st.button("🗑️ 清空所有自定义", key="clear_all_custom"):
                                        st.session_state.custom_prompts.clear()
                                        st.session_state.use_custom_prompt = False
                                        st.success("已清空所有自定义Prompt")
                                        st.rerun()
                                
                                # 显示状态信息
                                if st.session_state.custom_prompts:
                                    custom_count = len(st.session_state.custom_prompts)
                                    st.info(f"ℹ️ 当前有 {custom_count} 个功能点使用自定义Prompt")
                            else:
                                st.warning("无法生成Prompt预览：未找到适用的测试规范")

                    except Exception as e:
                        st.warning(f"测试规范功能加载失败: {str(e)}，将使用传统Prompt模式")
                        use_standards = False
                
                # 5. 更新PromptManager配置
                if st.session_state.prompt_manager.use_standards != use_standards:
                    st.session_state.prompt_manager.use_standards = use_standards
                    if use_standards and st.session_state.standard_manager:
                        st.session_state.prompt_manager.standard_manager = st.session_state.standard_manager
                
                # ========== 原有功能点选择部分 ==========
                st.markdown("---")
                st.header("📋 功能点列表")
                # 选择功能点
                selected_indices = st.multiselect(
                    "选择要生成测试的功能点（留空则选择全部）",
                    options=list(range(len(function_points))),
                    format_func=lambda x: f"{function_points[x]['module_name']} - {function_points[x]['function_name']}"
                )
                
                if not selected_indices:
                    selected_function_points = function_points
                else:
                    selected_function_points = [function_points[i] for i in selected_indices]
                
                # 显示功能点详情
                with st.expander("查看功能点详情"):
                    for idx, fp in enumerate(selected_function_points):
                        st.markdown(f"### {idx + 1}. {fp['function_name']}")
                        st.markdown(f"**模块**: {fp['module_name']}")
                        st.markdown(f"**描述**: {fp.get('description', 'N/A')}")
                        if fp.get('operation_steps'):
                            st.markdown("**操作步骤**:")
                            for step in fp['operation_steps']:
                                st.markdown(f"- {step}")
                        st.markdown("---")
                
                # 生成选项
                st.markdown("---")
                st.header("⚙️ 生成选项")
                
                generation_mode = st.radio(
                    "生成模式",
                    ['batch', 'separate'],
                    format_func=lambda x: {
                        'batch': '批量生成（一次性生成测试项和测试用例，更快）',
                        'separate': '分步生成（先生成测试项，再生成测试用例，更灵活）'
                    }[x]
                )
                
                # 开始生成
                if st.button("🚀 开始生成测试", type="primary", use_container_width=True):
                    if not st.session_state.api_client.get_provider_config().get('api_key'):
                        st.error("请先在配置页面设置API密钥！")
                    else:
                        generate_tests(selected_function_points, generation_mode)
            else:
                st.warning("未能从文档中提取到功能点，请检查文档格式。")
        
        except Exception as e:
            st.error(f"处理文件失败: {str(e)}")
            st.exception(e)
    
    else:
        st.info("👆 请上传JSON格式的文档文件")


def generate_tests(function_points: List[Dict[str, Any]], mode: str):
    """生成测试项和测试用例"""
    api_client = st.session_state.api_client
    prompt_manager = st.session_state.prompt_manager
    parser = st.session_state.json_parser
    
    if parser is None:
        st.error("请先上传文档")
        return
    
    # 先清空之前的自定义prompts
    prompt_manager.clear_custom_prompts()
    
    # 注入用户自定义的prompts到PromptManager
    all_function_points = st.session_state.function_points
    if st.session_state.custom_prompts:
        st.info(f"📝 开始注入 {len(st.session_state.custom_prompts)} 个自定义Prompt")
        
        for idx, fp in enumerate(all_function_points):
            if idx in st.session_state.custom_prompts:
                # 获取功能点ID（与TestGenerator._convert_to_requirement保持一致）
                func_id = fp.get('requirement_id', fp.get('function_name', ''))
                custom_prompt = st.session_state.custom_prompts[idx]
                
                # 调试信息
                st.write(f"  - 索引 {idx}: {fp.get('function_name', 'N/A')}")
                st.write(f"    - requirement_id: `{fp.get('requirement_id', 'N/A')}`")
                st.write(f"    - 映射到的ID: `{func_id}`")
                st.write(f"    - Prompt长度: {len(custom_prompt)} 字符")
                
                prompt_manager.set_custom_prompt(func_id, custom_prompt)
        
        st.success(f"✅ 已加载 {len(st.session_state.custom_prompts)} 个自定义Prompt到缓存")
        
        # 显示缓存内容
        with st.expander("🔍 查看Prompt缓存详情"):
            st.write("当前PromptManager缓存的自定义Prompt:")
            for key, value in prompt_manager.custom_prompts_cache.items():
                st.write(f"- **{key}**: {len(value)} 字符")
    
    # 创建生成器
    generator = TestGenerator(
        str(parser.json_path),
        api_client=api_client,
        prompt_manager=prompt_manager
    )
    
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    try:
        if mode == 'batch':
            # 批量生成
            status_text.text("开始批量生成...")
            result = generator.generate_batch(
                function_points=function_points,
                progress_callback=lambda current, total, msg: (
                    progress_bar.progress(current / total),
                    status_text.text(f"进度: {current}/{total} - {msg}")
                )
            )
            st.session_state.test_items = result['test_items']
            st.session_state.test_cases = result['test_cases']
        else:
            # 分步生成
            status_text.text("第一步：生成测试项...")
            test_items = generator.generate_test_items(
                function_points=function_points,
                progress_callback=lambda current, total, msg: (
                    progress_bar.progress(current / (total * 2)),
                    status_text.text(f"生成测试项: {current}/{total} - {msg}")
                )
            )
            st.session_state.test_items = test_items
            
            status_text.text("第二步：生成测试用例...")
            test_cases = generator.generate_test_cases(
                progress_callback=lambda current, total, msg: (
                    progress_bar.progress(0.5 + current / (total * 2)),
                    status_text.text(f"生成测试用例: {current}/{total} - {msg}")
                )
            )
            st.session_state.test_cases = test_cases
        
        progress_bar.progress(1.0)
        status_text.text("生成完成！")
        st.success(f"✅ 成功生成 {len(st.session_state.test_items)} 个测试项和 {len(st.session_state.test_cases)} 个测试用例")
        
        # 跳转到结果页面
        st.session_state.current_page = 'results'
        st.rerun()
    
    except Exception as e:
        st.error(f"生成失败: {str(e)}")
        st.exception(e)
        progress_bar.empty()
        status_text.empty()


def _convert_to_requirement(func_point: Dict[str, Any]) -> Dict[str, Any]:
    """
    将功能点转换为需求格式（用于测试规范选择）
    
    Args:
        func_point: 功能点字典
        
    Returns:
        需求字典
    """
    requirement = {
        'requirement_id': func_point.get('requirement_id', func_point.get('function_name', '')),
        'requirement_type': func_point.get('requirement_type', '功能需求'),
        'description': func_point.get('description', func_point.get('function_name', '')),
        'module_name': func_point.get('module_name', '')
    }
    
    # 如果有接口信息，添加进去
    if func_point.get('interface_info'):
        requirement['interface_info'] = func_point['interface_info']
    
    return requirement


def show_results_page():
    """显示结果页面"""
    st.title("📊 生成结果")
    st.markdown("---")
    
    test_items = st.session_state.test_items
    test_cases = st.session_state.test_cases
    document_info = st.session_state.document_info
    
    if not test_items and not test_cases:
        st.warning('还没有生成测试项和测试用例，请先前往"生成测试"页面')
        return
    
    # 统计信息
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.metric("测试项数量", len(test_items))
    with col2:
        st.metric("测试用例数量", len(test_cases))
    with col3:
        modules = set(item.get('module_name', '') for item in test_items)
        st.metric("涉及模块", len(modules))
    with col4:
        priorities = {}
        for case in test_cases:
            p = case.get('priority', '中')
            priorities[p] = priorities.get(p, 0) + 1
        high_priority = priorities.get('高', 0)
        st.metric("高优先级用例", high_priority)
    
    st.markdown("---")
    
    # 预览
    st.header("📋 预览")
    
    preview_format = st.selectbox("选择预览格式", ['JSON', 'Markdown'])
    
    if preview_format == 'JSON':
        formatter = OutputFormatter(test_items, test_cases, document_info)
        json_str = formatter.to_json()
        st.json(json.loads(json_str))
    else:
        formatter = OutputFormatter(test_items, test_cases, document_info)
        markdown_str = formatter.to_markdown()
        st.markdown(markdown_str)
    
    st.markdown("---")
    
    # 下载
    st.header("💾 下载结果")
    
    col1, col2, col3 = st.columns(3)
    
    formatter = OutputFormatter(test_items, test_cases, document_info)
    
    with col1:
        if st.button("下载JSON格式", use_container_width=True):
            output_path = "outputs/test_results.json"
            Path("outputs").mkdir(exist_ok=True)
            formatter.to_json(output_path)
            with open(output_path, 'rb') as f:
                st.download_button(
                    "⬇️ 下载JSON",
                    f.read(),
                    "test_results.json",
                    "application/json",
                    key="download_json"
                )
    
    with col2:
        if st.button("下载Markdown格式", use_container_width=True):
            output_path = "outputs/test_results.md"
            Path("outputs").mkdir(exist_ok=True)
            formatter.to_markdown(output_path)
            with open(output_path, 'rb') as f:
                st.download_button(
                    "⬇️ 下载Markdown",
                    f.read(),
                    "test_results.md",
                    "text/markdown",
                    key="download_md"
                )
    
    with col3:
        if st.button("下载Excel格式", use_container_width=True):
            try:
                output_path = "outputs/test_results.xlsx"
                Path("outputs").mkdir(exist_ok=True)
                formatter.to_excel(output_path)
                with open(output_path, 'rb') as f:
                    st.download_button(
                        "⬇️ 下载Excel",
                        f.read(),
                        "test_results.xlsx",
                        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                        key="download_excel"
                    )
            except ImportError:
                st.error("需要安装openpyxl: pip install openpyxl")


def main():
    """主函数"""
    init_session_state()
    
    # 侧边栏导航
    with st.sidebar:
        st.title("🧪 测试专家系统")
        st.markdown("---")
        
        page = st.radio(
            "导航",
            ['home', 'config', 'generate', 'results'],
            format_func=lambda x: {
                'home': '🏠 首页',
                'config': '⚙️ 配置',
                'generate': '🚀 生成测试',
                'results': '📊 结果'
            }[x]
        )
        
        st.markdown("---")
        st.markdown("### 使用提示")
        st.info("""
        1. 首次使用请先配置API密钥
        2. 上传JSON格式的文档
        3. 选择功能点并生成测试
        4. 下载生成的结果
        """)
    
    # 路由到对应页面
    if page == 'home':
        show_home_page()
    elif page == 'config':
        show_config_page()
    elif page == 'generate':
        show_generate_page()
    elif page == 'results':
        show_results_page()


if __name__ == "__main__":
    main()

