"""
Kimi API 集成示例
演示如何将MCP工具与Kimi大模型结合使用
"""

from openai import OpenAI
import json
from typing import List, Dict, Any
from mcp_client_example import MCPClient

# Kimi API配置
KIMI_API_KEY = "sk-TzB4VYJfCoXGcGrGMiewukVRzjuDsbVCkaZXi2LvkS8s60E5"
KIMI_BASE_URL = "https://api.moonshot.cn/v1"
KIMI_MODEL = "kimi-k2-turbo-preview"

# 初始化Kimi客户端
kimi_client = OpenAI(
    api_key=KIMI_API_KEY,
    base_url=KIMI_BASE_URL,
)

# 初始化MCP客户端
mcp_client = MCPClient()


def convert_mcp_tools_to_kimi_format() -> tuple[List[Dict], Dict]:
    """
    将MCP工具转换为Kimi API的tools格式

    Returns:
        tools: Kimi格式的工具列表
        tool_map: 工具名称到执行函数的映射
    """
    # 获取所有MCP工具
    mcp_tools_response = mcp_client.list_tools()
    mcp_tools = mcp_tools_response["tools"]

    # 转换为Kimi格式
    kimi_tools = []
    tool_map = {}

    for tool in mcp_tools:
        # Kimi工具格式
        kimi_tool = {
            "type": "function",
            "function": {
                "name": tool["name"],
                "description": tool["description"],
                "parameters": tool["parameters"]
            }
        }
        kimi_tools.append(kimi_tool)

        # 创建工具执行函数
        tool_name = tool["name"]
        tool_map[tool_name] = lambda args, name=tool_name: execute_mcp_tool(name, args)

    return kimi_tools, tool_map


def execute_mcp_tool(tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
    """
    执行MCP工具

    Args:
        tool_name: 工具名称
        arguments: 工具参数

    Returns:
        工具执行结果
    """
    print(f"[工具调用] {tool_name}")
    print(f"[参数] {json.dumps(arguments, ensure_ascii=False, indent=2)}")

    result = mcp_client.call_tool(tool_name, arguments)

    print(f"[结果] 成功: {result.get('success', False)}")

    return result


def chat_with_kimi(user_message: str, verbose: bool = True) -> str:
    """
    与Kimi进行对话，支持工具调用

    Args:
        user_message: 用户消息
        verbose: 是否打印详细信息

    Returns:
        Kimi的回复
    """
    # 获取Kimi格式的工具
    tools, tool_map = convert_mcp_tools_to_kimi_format()

    if verbose:
        print(f"\n{'='*60}")
        print(f"加载了 {len(tools)} 个工具")
        print(f"{'='*60}\n")

    # 初始化对话
    messages = [
        {
            "role": "system",
            "content": """你是一个专业的金融数据分析助手，由 Moonshot AI 提供支持。
你可以使用各种工具来帮助用户查询和分析金融数据，包括：
- 新闻搜索（全球新闻、中国新闻、医疗新闻）
- 公司研究（路演信息、研究报告）
- 概念板块分析
- 股票分析（涨停分析、财务数据、交易数据）
- 财务报表（资产负债表、现金流量表）

请根据用户的问题，选择合适的工具来获取信息，并提供专业的分析和建议。"""
        },
        {
            "role": "user",
            "content": user_message
        }
    ]

    if verbose:
        print(f"[用户]: {user_message}\n")

    # 对话循环，处理工具调用
    finish_reason = None
    iteration = 0
    max_iterations = 10  # 防止无限循环

    while finish_reason is None or finish_reason == "tool_calls":
        iteration += 1
        if iteration > max_iterations:
            print("[警告] 达到最大迭代次数")
            break

        if verbose and iteration > 1:
            print(f"\n[轮次 {iteration}]")

        # 调用Kimi API
        completion = kimi_client.chat.completions.create(
            model=KIMI_MODEL,
            messages=messages,
            temperature=0.6,  # Kimi推荐的temperature值
            tools=tools,
        )

        choice = completion.choices[0]
        finish_reason = choice.finish_reason

        if verbose:
            print(f"[Kimi] finish_reason: {finish_reason}")

        # 处理工具调用
        if finish_reason == "tool_calls":
            # 将Kimi的消息添加到上下文
            messages.append(choice.message)

            # 执行每个工具调用
            for tool_call in choice.message.tool_calls:
                tool_name = tool_call.function.name
                tool_arguments = json.loads(tool_call.function.arguments)

                # 执行工具
                tool_result = tool_map[tool_name](tool_arguments)

                # 将工具结果添加到消息中
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "name": tool_name,
                    "content": json.dumps(tool_result, ensure_ascii=False),
                })

            if verbose:
                print()  # 空行分隔

    # 返回最终回复
    final_response = choice.message.content

    if verbose:
        print(f"\n[Kimi]: {final_response}\n")
        print(f"{'='*60}")

    return final_response


def demo_simple_query():
    """演示1: 简单查询"""
    print("\n" + "="*60)
    print("演示1: 简单新闻查询")
    print("="*60)

    response = chat_with_kimi("帮我查找关于人工智能的最新新闻")
    return response


def demo_stock_analysis():
    """演示2: 股票分析"""
    print("\n" + "="*60)
    print("演示2: 股票财务分析")
    print("="*60)

    response = chat_with_kimi("帮我分析贵州茅台（600519）的财务状况")
    return response


def demo_concept_research():
    """演示3: 概念研究"""
    print("\n" + "="*60)
    print("演示3: 概念板块研究")
    print("="*60)

    response = chat_with_kimi("查找新能源汽车相关的概念板块，并告诉我涨幅最高的是哪些")
    return response


def demo_industry_comparison():
    """演示4: 行业对比"""
    print("\n" + "="*60)
    print("演示4: 行业内股票对比")
    print("="*60)

    response = chat_with_kimi("帮我找出半导体行业的龙头股票，并对比它们的财务指标")
    return response


def demo_comprehensive_analysis():
    """演示5: 综合分析"""
    print("\n" + "="*60)
    print("演示5: 综合分析")
    print("="*60)

    response = chat_with_kimi("""
    我想投资白酒行业，请帮我：
    1. 搜索白酒行业的主要上市公司
    2. 对比贵州茅台和五粮液的财务数据
    3. 查看最近的行业新闻
    4. 给出投资建议
    """)
    return response


def interactive_chat():
    """交互式对话"""
    print("\n" + "="*60)
    print("Kimi 金融助手 - 交互模式")
    print("="*60)
    print("提示：输入 'quit' 或 'exit' 退出")
    print("="*60 + "\n")

    while True:
        try:
            user_input = input("你: ").strip()

            if not user_input:
                continue

            if user_input.lower() in ['quit', 'exit', '退出']:
                print("\n再见！")
                break

            response = chat_with_kimi(user_input)

        except KeyboardInterrupt:
            print("\n\n再见！")
            break
        except Exception as e:
            print(f"\n[错误] {str(e)}\n")


def test_kimi_connection():
    """测试Kimi API连接"""
    print("\n" + "="*60)
    print("测试 Kimi API 连接")
    print("="*60 + "\n")

    try:
        # 简单的测试请求
        response = kimi_client.chat.completions.create(
            model=KIMI_MODEL,
            messages=[
                {"role": "user", "content": "你好，请介绍一下你自己"}
            ],
            temperature=0.6
        )

        print("[✓] 连接成功!")
        print(f"[✓] 模型: {KIMI_MODEL}")
        print(f"[✓] 回复: {response.choices[0].message.content}\n")

        return True
    except Exception as e:
        print(f"[✗] 连接失败: {str(e)}\n")
        return False


def show_available_tools():
    """显示所有可用工具"""
    print("\n" + "="*60)
    print("可用工具列表")
    print("="*60 + "\n")

    tools, _ = convert_mcp_tools_to_kimi_format()

    for i, tool in enumerate(tools, 1):
        func = tool["function"]
        print(f"{i}. {func['name']}")
        print(f"   描述: {func['description'][:80]}...")
        print()

    print(f"总计: {len(tools)} 个工具\n")


if __name__ == "__main__":
    import sys

    # 首先测试连接
    if not test_kimi_connection():
        print("请检查API Key和网络连接")
        sys.exit(1)

    # 显示可用工具
    show_available_tools()

    # 运行演示
    print("\n选择运行模式:")
    print("1. 简单查询演示")
    print("2. 股票分析演示")
    print("3. 概念研究演示")
    print("4. 行业对比演示")
    print("5. 综合分析演示")
    print("6. 交互式对话")
    print("7. 运行所有演示")

    try:
        choice = input("\n请选择 (1-7): ").strip()

        if choice == "1":
            demo_simple_query()
        elif choice == "2":
            demo_stock_analysis()
        elif choice == "3":
            demo_concept_research()
        elif choice == "4":
            demo_industry_comparison()
        elif choice == "5":
            demo_comprehensive_analysis()
        elif choice == "6":
            interactive_chat()
        elif choice == "7":
            demo_simple_query()
            demo_stock_analysis()
            demo_concept_research()
            demo_industry_comparison()
            demo_comprehensive_analysis()
        else:
            print("无效选择")

    except KeyboardInterrupt:
        print("\n\n程序已退出")
    finally:
        mcp_client.close()