#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Claude Agent SDK Demo - 简洁版

演示 claude_agent_sdk 的核心功能：定义本地工具并让 agent 调用。
"""
import asyncio
from typing import Any, Optional, List
import httpx
import json

from claude_agent_sdk import (
    tool,
    create_sdk_mcp_server,
    ClaudeSDKClient,
    ClaudeAgentOptions,
    AssistantMessage,
    ResultMessage,
    ToolUseBlock,
    ToolResultBlock,
)

from claude_agent_sdk.types import (
    HookMatcher,
    PermissionResultAllow,
    PermissionResultDeny,
    ToolPermissionContext,
)

BASE_URL = "http://aigc-channel.aiddit.com/aigc/channel"
DEFAULT_TIMEOUT = 60.0

# ============================================================================
# 工具定义
# ============================================================================

@tool(
    "search_post",
    "搜索帖子",
    {
        "type": "object",
        "properties": {
            "keyword": {"type": "string", "description": "搜索关键词"}
        },
        "required": ["keyword"]
    }
)
async def search_post(args: dict[str, Any]) -> dict[str, Any]:
    try:
        keyword = args.get("keyword")
        payload = {
            "type": "zhihu",
            "keyword": keyword,
            "cursor": "0",
            "max_count": 2,
            "content_type": "图文",
        }

        print(f"search_post,payload:{payload}")

        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
            resp = await client.post(
                f"{BASE_URL}/data",
                json=payload,
                headers={"Content-Type": "application/json"},
            )
            resp.raise_for_status()
            data = resp.json()
        posts = data.get("data") or []

        posts_json_str = json.dumps(posts, ensure_ascii=False, indent=2)

        print(f"search_post,result posts_json_str: {posts_json_str}")

        return {"content": [{"type": "text", "text": posts_json_str}]}
    except Exception as e:
        return {"content": [{"type": "text", "text": f"错误: {str(e)}"}]}



async def _auto_approve_tool(
        tool_name: str, input_data: dict, context: ToolPermissionContext
) -> PermissionResultAllow | PermissionResultDeny:
    if tool_name == "AskUserQuestion":
        questions = input_data.get("questions", [])
        answers = {}
        for q in questions:
            question_text = q.get("question", "")
            options = q.get("options", [])
            if options:
                answers[question_text] = options[0].get("label", "")
            else:
                answers[question_text] = ""
        print(f"[auto_approve] AskUserQuestion 自动选择: {answers}")
        return PermissionResultAllow(updated_input={**input_data, "answers": answers})

    return PermissionResultAllow(updated_input=input_data)

# ============================================================================
# Agent 封装
# ============================================================================

class SimpleAgent:
    def __init__(self, model: str = "claude-sonnet-4-6"):
        self.model = model
        self.server = create_sdk_mcp_server(
            name="demo-tools",
            version="1.0.0",
            tools=[search_post],
        )

    async def run(self, query: str, verbose: bool = True):
        """运行 agent"""
        if verbose:
            print(f"\n{'='*60}")
            print(f"查询: {query}")
            print(f"{'='*60}\n")

        options = ClaudeAgentOptions(
            system_prompt="你是一个内容收集专家，善于利用 search_post 工具搜索帖子内容，并对搜索结果进行格式化处理，理解并输出每个帖子的主要内容。\n "
                          "注意：你需要提取帖子中的每张图片url并分析出图片中的关键信息，综合帖子其他信息返回帖子的主要内容",
            model=self.model,
            mcp_servers={"demo-tools": self.server},
            allowed_tools=["mcp__demo-tools__search_post"],
            disallowed_tools=[
                "Bash", "Read", "Write", "Edit", "MultiEdit",
                "Glob", "Grep", "WebSearch", "WebFetch",
                "TodoRead", "TodoWrite",
            ],
            permission_mode="bypassPermissions",
            max_turns=10,
            effort="low",
            can_use_tool=_auto_approve_tool,
        )

        tool_calls = []
        response = ""

        async with ClaudeSDKClient(options=options) as client:
            await client.query(query)

            async for msg in client.receive_response():
                if isinstance(msg, AssistantMessage):
                    for block in msg.content:
                        if isinstance(block, ToolUseBlock):
                            tool_calls.append(f"{block.name}({block.input})")
                            if verbose:
                                print(f"🔧 工具调用: {block.name}({block.input})")
                        elif isinstance(block, ToolResultBlock):
                            if verbose and hasattr(block, 'content'):
                                for c in block.content:
                                    if hasattr(c, 'text'):
                                        print(f"📥 工具结果: {c.text}")
                        elif hasattr(block, 'text') and block.text:
                            response = block.text
                            if verbose:
                                print(f"💬 回复: {block.text}")
                elif isinstance(msg, ResultMessage):
                    if verbose:
                        print(f"\n📊 统计: 成本=${msg.total_cost_usd:.4f}, 轮次={msg.num_turns}")

        return {"response": response, "tool_calls": tool_calls}


# ============================================================================
# 使用示例
# ============================================================================

async def main():
    agent = SimpleAgent()

    await agent.run("搜索 “北京秋天” 并返回帖子主要内容")

    # await search_post(keyword="柴犬")


if __name__ == "__main__":
    asyncio.run(main())