"""
优化后的3个核心蛋白质工具
从原来的8个工具优化为3个整合工具，减少复杂度，提升用户体验
"""

import json
from typing import Any

from fastmcp import Context

from .utils import (
    calculate_dssp,
    download_file,
    extract_sequence_from_pdb,
    format_error_response,
    format_success_response,
    get_supported_formats,
    make_http_request,
    validate_pdb_id,
)

RCSB_GRAPHQL_URL = "https://data.rcsb.org/graphql"
RCSB_REST_URL = "https://data.rcsb.org/rest/v1/core"
RCSB_SEARCH_URL = "https://search.rcsb.org/rcsbsearch/v2/query"
RCSB_DOWNLOAD_URL = "https://files.rcsb.org/download"


def _validate_pdb_exists(pdb_id: str) -> bool:
    """验证PDB ID是否存在于RCSB数据库"""
    try:
        url = f"{RCSB_REST_URL}/entry/{pdb_id}"
        response = make_http_request(url)
        return response is not None
    except Exception:
        return False


def _get_entry_info(pdb_id: str) -> dict[str, Any] | None:
    """获取PDB条目的基本信息"""
    try:
        url = f"{RCSB_REST_URL}/entry/{pdb_id}"
        return make_http_request(url)
    except Exception:
        return None


def _search_rcsb_structures(keywords: str, max_results: int = 10) -> list[dict[str, Any]]:
    """使用RCSB搜索API查找蛋白质结构"""
    try:
        search_query = {
            "query": {
                "type": "terminal",
                "service": "text",
                "parameters": {
                    "attribute": "struct.title",
                    "operator": "contains_words",
                    "value": keywords,
                },
            },
            "return_type": "entry",
            "request_options": {
                "return_all_hits": False,
                "paginate": {"start": 0, "rows": max_results},
                "sort": [
                    {"sort_by": "rcsb_accession_info.initial_release_date", "direction": "desc"}
                ],
            },
        }

        response = make_http_request(RCSB_SEARCH_URL, method="POST", data=json.dumps(search_query))

        if response and "result_set" in response:
            results = []
            for item in response["result_set"][:max_results]:
                results.append(
                    {
                        "pdb_id": item.get("identifier", ""),
                        "title": item.get("cite", {}).get("title", ""),
                        "score": item.get("score", 0.0),
                    }
                )
            return results

        return []
    except Exception as e:
        print(f"搜索错误: {str(e)}")
        return []


async def find_protein_structures(
    keywords: str | None = None,
    category: str | None = None,
    pdb_id: str | None = None,
    max_results: int = 10,
    ctx: Context | None = None,
) -> dict[str, Any]:
    """
    蛋白质结构发现工具 - 搜索、示例、验证的统一入口

    这是蛋白质研究的起点，帮助你发现和验证PDB结构。

    Args:
        keywords: 搜索关键词 (如: "hemoglobin", "kinase", "DNA")
        category: 预设类别 ("癌症靶点", "病毒蛋白", "酶类", "抗体", "膜蛋白", "核糖体")
        pdb_id: 直接验证或查看特定PDB ID (如: "1A3N")
        max_results: 搜索结果最大数量 (默认10，最大100)
        ctx: FastMCP Context，用于进度反馈和日志记录

    Returns:
        包含PDB结构列表、验证结果、示例数据的综合响应
    """
    try:
        if ctx:
            await ctx.info(f"🔍 开始蛋白质结构搜索...")
        # 限制max_results范围
        max_results = min(max(max_results, 1), 100)

        # 模式1: 验证特定PDB ID
        if pdb_id:
            if ctx:
                await ctx.info(f"📋 验证 PDB ID: {pdb_id}")

            if not validate_pdb_id(pdb_id):
                if ctx:
                    await ctx.error(f"❌ 无效的PDB ID格式: {pdb_id}")
                return format_error_response(
                    "无效的PDB ID格式",
                    f"期望格式: 4位字符 (首位数字，后三位可数字可字母)，实际: {pdb_id}",
                )

            if _validate_pdb_exists(pdb_id):
                if ctx:
                    await ctx.info(f"🔍 获取蛋白质详细信息...")
                entry_info = _get_entry_info(pdb_id)
                title = "未知标题"
                if entry_info and "struct" in entry_info:
                    title = entry_info["struct"].get("title", "未知标题")

                if ctx:
                    await ctx.info(f"✅ PDB ID {pdb_id} 验证成功")
                return format_success_response(
                    {
                        "mode": "validation",
                        "pdb_id": pdb_id,
                        "exists": True,
                        "title": title,
                        "validation_result": "PDB ID有效且存在于RCSB数据库中",
                    },
                    f"PDB ID {pdb_id} 验证成功，结构存在于RCSB数据库",
                )
            else:
                if ctx:
                    await ctx.error(f"❌ PDB ID {pdb_id} 不存在")
                return format_error_response(
                    "PDB ID不存在", f"PDB ID {pdb_id} 在RCSB数据库中未找到"
                )

        # 模式2: 分类示例
        elif category:
            if ctx:
                await ctx.info(f"📚 获取 {category} 类别的蛋白质示例...")

            # 预定义的分类示例数据库
            category_examples = {
                "癌症靶点": {
                    "EGFR": ["1M14", "4MNF"],
                    "KRAS": ["4OBE", "6OIM"],
                    "p53": ["1TUP", "2OCJ"],
                    "BCR-ABL": ["1IEP", "1XBB"],
                },
                "病毒蛋白": {
                    "SARS-CoV-2刺突蛋白": ["6VSB", "6Y2E"],
                    "HIV蛋白酶": ["1HHP", "3PHV"],
                    "流感病毒血凝素": ["1RU7", "4WE8"],
                },
                "酶类": {
                    "溶菌酶": ["1HEW", "1LZ1"],
                    "DNA聚合酶": ["1KLN", "3K0M"],
                    "激酶": ["1ATP", "2SRC"],
                },
                "抗体": {"单克隆抗体": ["1HZH", "1IGT"], "抗体片段": ["1FVC", "3HFM"]},
                "膜蛋白": {"G蛋白偶联受体": ["1U19", "3DQB"], "离子通道": ["1BL8", "2A9W"]},
                "核糖体": {"核糖体亚基": ["1FJG", "2VQE"], "核糖体相关因子": ["1EF1", "2AW4"]},
            }

            if category not in category_examples:
                available_categories = list(category_examples.keys())
                return format_error_response(
                    "不支持的类别", f"可用类别: {', '.join(available_categories)}"
                )

            examples = category_examples[category]
            if ctx:
                await ctx.info(f"✅ 找到 {len(examples)} 个 {category} 示例")
            return format_success_response(
                {
                    "mode": "category_examples",
                    "category": category,
                    "examples": examples,
                    "total_proteins": len(examples),
                },
                f"获取 {category} 类别的蛋白质结构示例",
            )

        # 模式3: 关键词搜索
        elif keywords:
            if ctx:
                await ctx.info(f"🔍 搜索关键词: {keywords}")

            results = _search_rcsb_structures(keywords, max_results)
            if results:
                if ctx:
                    await ctx.info(f"✅ 找到 {len(results)} 个匹配的结构")
                return format_success_response(
                    {
                        "mode": "search",
                        "keywords": keywords,
                        "results": results,
                        "total_found": len(results),
                    },
                    f"找到 {len(results)} 个匹配的结构",
                )
            else:
                if ctx:
                    await ctx.error(f"❌ 使用关键词 '{keywords}' 未找到匹配结果")
                return format_error_response(
                    "未找到匹配结果", f"使用关键词 '{keywords}' 未找到匹配的PDB结构"
                )

        # 模式4: 默认 - 返回热门示例
        else:
            if ctx:
                await ctx.info("📋 提供热门蛋白质结构示例...")

            popular_examples = {
                "经典蛋白质": {
                    "血红蛋白": ["1A3N", "2HHB"],
                    "胰岛素": ["1ZNJ", "4INS"],
                    "溶菌酶": ["1HEW", "1LZ1"],
                    "DNA聚合酶": ["1KLN", "3K0M"],
                },
                "热门研究领域": {
                    "癌症研究": ["1M14", "4MNF", "1TUP"],
                    "病毒研究": ["6VSB", "6Y2E", "1HHP"],
                    "神经科学": ["2BEG", "5O3L", "1XQ8"],
                },
            }

            if ctx:
                await ctx.info("✅ 返回热门示例和使用指南")
            return format_success_response(
                {
                    "mode": "default_examples",
                    "popular_examples": popular_examples,
                    "total_categories": len(popular_examples),
                    "usage_tips": {
                        "search": "使用 keywords 参数搜索特定蛋白质",
                        "category": "使用 category 参数获取分类示例",
                        "validate": "使用 pdb_id 参数验证特定结构",
                    },
                },
                "提供热门蛋白质结构示例和使用指南",
            )

    except Exception as e:
        if ctx:
            await ctx.error(f"❌ 搜索失败: {str(e)}")
        return format_error_response("工具执行错误", f"find_protein_structures 执行失败: {str(e)}")


async def get_protein_data(
    pdb_id: str, data_types: list[str], chain_id: str | None = None, ctx: Context | None = None
) -> dict[str, Any]:
    """
    蛋白质综合数据工具 - 获取完整蛋白质信息包

    这个工具是蛋白质数据获取的核心，一次性获取你需要的所有信息。

    Args:
        pdb_id: PDB ID (例如: "5G53")
        data_types: 需要的数据类型列表
            - "basic": 基本信息 (标题、方法、分辨率等)
            - "sequence": 氨基酸序列信息
            - "structure": 二级结构分析
            - "all": 获取所有数据
        chain_id: 特定链ID (例如: "A"，可选)
        ctx: FastMCP Context，用于进度反馈和日志记录

    Returns:
        完整的蛋白质数据包，包含请求的所有数据类型
    """
    try:
        if ctx:
            await ctx.info(f"📊 开始获取蛋白质数据: {pdb_id}")
            await ctx.report_progress(0, 100, "初始化...")
        if not validate_pdb_id(pdb_id):
            return format_error_response(
                "无效的PDB ID格式",
                f"期望格式: 4位字符 (首位数字，后三位可数字可字母)，实际: {pdb_id}",
            )

        # 处理 "all" 参数
        if "all" in data_types:
            data_types = ["basic", "sequence", "structure"]

        # 验证PDB ID存在性
        if not _validate_pdb_exists(pdb_id):
            if ctx:
                await ctx.error(f"❌ PDB ID {pdb_id} 不存在")
            return format_error_response("PDB ID不存在", f"PDB ID {pdb_id} 在RCSB数据库中未找到")

        result_data = {}

        # 获取基本信息
        if "basic" in data_types:
            if ctx:
                await ctx.report_progress(25, 100, "获取基本信息...")
                await ctx.info(f"🔍 查询 {pdb_id} 基本信息...")

            entry_info = _get_entry_info(pdb_id)
            if entry_info:
                struct_data = entry_info.get("struct", {})
                result_data["basic"] = {
                    "pdb_id": pdb_id,
                    "title": struct_data.get("title", "未知标题"),
                    "method": struct_data.get("pdbx_descriptor", "未知方法"),
                    "resolution": struct_data.get("pdbx_resolution", None),
                    "deposition_date": struct_data.get("pdbx_deposit_date", "未知"),
                    "authors": [
                        author.get("name", "未知") for author in entry_info.get("audit_author", [])
                    ],
                }
                if ctx:
                    await ctx.info("✅ 基本信息获取完成")
            else:
                result_data["basic"] = {"error": "无法获取基本信息"}

        # 获取序列信息
        if "sequence" in data_types:
            try:
                if ctx:
                    await ctx.report_progress(50, 100, "提取序列信息...")
                    await ctx.info("🧬 下载并解析PDB文件...")

                # 下载PDB文件并提取序列
                pdb_url = f"{RCSB_DOWNLOAD_URL}/{pdb_id}.pdb"
                local_pdb_file = f"{pdb_id}.pdb"

                if download_file(pdb_url, local_pdb_file):
                    sequence_data = extract_sequence_from_pdb(local_pdb_file, chain_id)
                    if sequence_data:
                        result_data["sequence"] = {
                            "chain_id": sequence_data.get("chain_id", chain_id or "A"),
                            "sequence_1_letter": sequence_data.get("sequence_1_letter", ""),
                            "sequence_3_letter": sequence_data.get("sequence_3_letter", ""),
                            "length": sequence_data.get("length", 0),
                        }
                        if ctx:
                            await ctx.info(f"✅ 序列信息提取完成 (长度: {sequence_data.get('length', 0)})")
                    else:
                        result_data["sequence"] = {"error": "无法提取序列信息"}
                else:
                    result_data["sequence"] = {"error": "PDB文件下载失败"}
            except Exception as e:
                result_data["sequence"] = {"error": f"序列提取失败: {str(e)}"}

        # 获取二级结构信息
        if "structure" in data_types:
            if "sequence" in result_data and "sequence_1_letter" in result_data["sequence"]:
                sequence = result_data["sequence"]["sequence_1_letter"]
                try:
                    if ctx:
                        await ctx.report_progress(75, 100, "分析二级结构...")
                        await ctx.info("🔬 执行DSSP二级结构分析...")

                    secondary_structure = calculate_dssp(pdb_id, sequence)
                    result_data["structure"] = {
                        "dssp_prediction": secondary_structure,
                        "sequence_length": len(sequence),
                        "composition": {
                            "helix": secondary_structure.count("H"),
                            "strand": secondary_structure.count("E"),
                            "coil": secondary_structure.count("C"),
                        },
                    }
                    if ctx:
                        await ctx.info("✅ 二级结构分析完成")
                except Exception as e:
                    result_data["structure"] = {"error": f"二级结构分析失败: {str(e)}"}
            else:
                result_data["structure"] = {"error": "需要先获取序列信息"}

        # 计算成功率
        successful_types = [
            dt for dt in data_types if dt in result_data and "error" not in result_data[dt]
        ]
        success_rate = len(successful_types) / len(data_types) * 100 if data_types else 0

        if ctx:
            await ctx.report_progress(100, 100, "完成")
            await ctx.info(f"✅ 成功获取 {pdb_id} 的数据 ({success_rate:.0f}% 成功率)")

        return format_success_response(
            {
                "pdb_id": pdb_id,
                "requested_data_types": data_types,
                "data": result_data,
                "success_rate": success_rate,
                "chain_id": chain_id,
            },
            f"成功获取 {pdb_id} 的数据: {', '.join(successful_types)} ({success_rate:.0f}%)",
        )

    except Exception as e:
        if ctx:
            await ctx.error(f"❌ 数据获取失败: {str(e)}")
        return format_error_response("数据获取错误", f"get_protein_data 执行失败: {str(e)}")


async def download_structure(
    pdb_id: str, file_format: str = "pdb", save_local: bool = False, ctx: Context | None = None
) -> dict[str, Any]:
    """
    结构文件工具 - 下载和管理蛋白质结构文件

    这个工具处理所有文件相关的操作，从下载到格式说明。

    Args:
        pdb_id: PDB ID (例如: "5G53")
        file_format: 文件格式
            - "pdb": 标准PDB格式 (推荐，人类可读)
            - "mmcif": 大分子晶体信息文件格式 (现代标准)
            - "cif": 晶体信息文件格式
            - "mmtf": 大分子传输格式 (二进制，速度快)
        save_local: 是否保存到本地文件 (默认False返回内容)
        ctx: FastMCP Context，用于进度反馈和日志记录

    Returns:
        文件内容或下载信息 + 格式说明和使用指南
    """
    try:
        if ctx:
            await ctx.info(f"📁 开始下载结构文件: {pdb_id}.{file_format}")
            await ctx.report_progress(0, 100, "初始化下载...")
        if not validate_pdb_id(pdb_id):
            return format_error_response(
                "无效的PDB ID格式",
                f"期望格式: 4位字符 (首位数字，后三位可数字可字母)，实际: {pdb_id}",
            )

        # 验证PDB ID存在性
        if not _validate_pdb_exists(pdb_id):
            if ctx:
                await ctx.error(f"❌ PDB ID {pdb_id} 不存在")
            return format_error_response("PDB ID不存在", f"PDB ID {pdb_id} 在RCSB数据库中未找到")

        # 验证文件格式
        supported_formats = get_supported_formats()
        if file_format not in supported_formats:
            if ctx:
                await ctx.error(f"❌ 不支持的文件格式: {file_format}")
            return format_error_response(
                "不支持的文件格式", f"支持格式: {', '.join(supported_formats)}"
            )

        # 构建下载URL
        download_url = f"{RCSB_DOWNLOAD_URL}/{pdb_id}.{file_format}"
        local_filename = f"{pdb_id}.{file_format}"

        if ctx:
            await ctx.report_progress(50, 100, f"下载 {file_format.upper()} 格式文件...")

        # 下载文件
        if save_local:
            if ctx:
                await ctx.info(f"💾 保存到本地: {local_filename}")

            success = download_file(download_url, local_filename)
            if success:
                result_data = {
                    "pdb_id": pdb_id,
                    "file_format": file_format,
                    "file_path": local_filename,
                    "download_method": "saved_local",
                    "file_size": None,  # 可以添加文件大小信息
                }
                if ctx:
                    await ctx.info(f"✅ 文件保存成功: {local_filename}")
            else:
                if ctx:
                    await ctx.error(f"❌ 文件下载失败: {local_filename}")
                return format_error_response(
                    "文件下载失败", f"无法下载 {pdb_id}.{file_format} 文件"
                )
        else:
            # 返回文件内容（对于小文件）
            try:
                if ctx:
                    await ctx.info("🌐 从远程获取文件内容预览...")

                import requests

                response = requests.get(download_url, timeout=30)
                if response.status_code == 200:
                    result_data = {
                        "pdb_id": pdb_id,
                        "file_format": file_format,
                        "file_path": download_url,
                        "download_method": "url_provided",
                        "file_content": (
                            response.text[:1000] + "..."
                            if len(response.text) > 1000
                            else response.text
                        ),
                        "content_preview": True,
                    }
                    if ctx:
                        await ctx.info(f"✅ 文件内容获取成功 (预览: {len(response.text)} 字符)")
                else:
                    if ctx:
                        await ctx.error(f"❌ HTTP错误: {response.status_code}")
                    return format_error_response(
                        "文件下载失败", f"HTTP {response.status_code}: 无法访问文件"
                    )
            except Exception as e:
                if ctx:
                    await ctx.error(f"❌ 网络错误: {str(e)}")
                return format_error_response("网络错误", f"下载失败: {str(e)}")

        # 添加格式信息
        format_info = {
            "pdb": {
                "name": "Protein Data Bank (PDB) 格式",
                "description": "经典的文本格式，人类可读",
                "recommended": True,
                "use_case": "一般用途，化学演示，小分子结构",
                "advantages": ["人类可读", "广泛支持", "适合编辑"],
            },
            "mmcif": {
                "name": "大分子晶体信息文件格式",
                "description": "现代的XML风格格式，更灵活",
                "recommended": True,
                "use_case": "复杂结构，大批量数据，现代软件",
                "advantages": ["更详细", "支持复杂数据", "现代标准"],
            },
            "cif": {
                "name": "晶体信息文件格式",
                "description": "标准化格式， mmcif的简化版",
                "recommended": False,
                "use_case": "基本晶体学数据",
                "advantages": ["标准化", "简洁"],
            },
            "mmtf": {
                "name": "大分子传输格式",
                "description": "二进制格式，压缩高效",
                "recommended": True,
                "use_case": "大批量传输，高性能应用",
                "advantages": ["文件小", "加载快", "压缩效率高"],
            },
        }

        result_data["format_info"] = format_info.get(
            file_format,
            {
                "name": f"{file_format.upper()} 格式",
                "description": "支持的文件格式",
                "recommended": False,
                "use_case": "通用格式",
                "advantages": ["标准支持"],
            },
        )

        if ctx:
            await ctx.report_progress(100, 100, "完成")
            await ctx.info(f"✅ 成功获取 {pdb_id} 的 {file_format} 格式文件")

        return format_success_response(
            result_data,
            f"成功获取 {pdb_id} 的 {file_format} 格式文件。{format_info.get(file_format, {}).get('description', '')}",
        )

    except Exception as e:
        if ctx:
            await ctx.error(f"❌ 文件操作失败: {str(e)}")
        return format_error_response("文件操作错误", f"download_structure 执行失败: {str(e)}")


def register_all_tools(mcp) -> None:
    """
    注册3个核心整合工具到FastMCP服务器

    优化后的工具设计：
    1. find_protein_structures - 蛋白质结构发现工具
    2. get_protein_data - 蛋白质综合数据工具
    3. download_structure - 结构文件工具

    Args:
        mcp: FastMCP服务器实例
    """

    # 工具1: 蛋白质结构发现工具 - 整合搜索、示例、验证功能
    @mcp.tool()
    async def find_protein_structures_tool(
        keywords: str | None = None,
        category: str | None = None,
        pdb_id: str | None = None,
        max_results: int = 10,
        ctx: Context | None = None,
    ) -> dict[str, Any]:
        """
        蛋白质结构发现工具 - 搜索、示例、验证的统一入口

        这是蛋白质研究的起点，帮助你发现和验证PDB结构。

        Args:
            keywords: 搜索关键词 (如: "hemoglobin", "kinase", "DNA")
            category: 预设类别 ("癌症靶点", "病毒蛋白", "酶类", "抗体", "膜蛋白", "核糖体")
            pdb_id: 直接验证或查看特定PDB ID (如: "1A3N")
            max_results: 搜索结果最大数量 (默认10，最大100)
            ctx: FastMCP Context，用于进度反馈和日志记录

        Returns:
            包含PDB结构列表、验证结果、示例数据的综合响应

        Examples:
            # 搜索血红蛋白相关结构
            find_protein_structures(keywords="hemoglobin")

            # 获取癌症靶点示例
            find_protein_structures(category="癌症靶点")

            # 验证PDB ID
            find_protein_structures(pdb_id="1A3N")
        """
        return await find_protein_structures(keywords, category, pdb_id, max_results, ctx)

    # 工具2: 蛋白质综合数据工具 - 一次获取所有蛋白质信息
    @mcp.tool()
    async def get_protein_data_tool(
        pdb_id: str,
        data_types: list[str] | None = None,
        chain_id: str | None = None,
        ctx: Context | None = None,
    ) -> dict[str, Any]:
        """
        蛋白质综合数据工具 - 获取完整蛋白质信息包

        这个工具是蛋白质数据获取的核心，一次性获取你需要的所有信息。

        Args:
            pdb_id: PDB ID (例如: "5G53")
            data_types: 需要的数据类型列表
                - "basic": 基本信息 (标题、方法、分辨率等)
                - "sequence": 氨基酸序列信息
                - "structure": 二级结构分析
                - "all": 获取所有数据
            chain_id: 特定链ID (例如: "A"，可选)
            ctx: FastMCP Context，用于进度反馈和日志记录

        Returns:
            完整的蛋白质数据包，包含请求的所有数据类型

        Examples:
            # 获取所有数据
            get_protein_data("5G53", ["all"])

            # 只获取基本信息和序列
            get_protein_data("1A3N", ["basic", "sequence"])

            # 获取特定链的数据
            get_protein_data("2HHB", ["all"], "A")
        """
        # 如果没有指定数据类型，默认获取基本数据
        if data_types is None:
            data_types = ["basic", "sequence", "structure"]
        return await get_protein_data(pdb_id, data_types, chain_id, ctx)

    # 工具3: 结构文件工具 - 下载和管理蛋白质结构文件
    @mcp.tool()
    async def download_structure_tool(
        pdb_id: str,
        file_format: str = "pdb",
        save_local: bool = False,
        ctx: Context | None = None,
    ) -> dict[str, Any]:
        """
        结构文件工具 - 下载和管理蛋白质结构文件

        这个工具处理所有文件相关的操作，从下载到格式说明。

        Args:
            pdb_id: PDB ID (例如: "5G53")
            file_format: 文件格式
                - "pdb": 标准PDB格式 (推荐，人类可读)
                - "mmcif": 大分子晶体信息文件格式 (现代标准)
                - "cif": 晶体信息文件格式
                - "mmtf": 大分子传输格式 (二进制，速度快)
            save_local: 是否保存到本地文件 (默认False返回内容)
            ctx: FastMCP Context，用于进度反馈和日志记录

        Returns:
            文件内容或下载信息 + 格式说明和使用指南

        Examples:
            # 获取PDB文件内容
            download_structure("1A3N")

            # 下载mmCIF格式并保存到本地
            download_structure("2HHB", "mmcif", True)

            # 获取快速MMTF格式
            download_structure("6VSB", "mmtf")
        """
        return await download_structure(pdb_id, file_format, save_local, ctx)
