from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional
from pydantic import BaseModel


class AgentMessage(BaseModel):
    role: str
    content: str


class ToolDefinition(BaseModel):
    name: str
    description: str
    parameters_schema: Dict[str, Any]


class ModelAdapter(ABC):
    """统一模型调用抽象基类"""

    @abstractmethod
    def generate(
        self,
        system_prompt: str,
        messages: List[AgentMessage],
        tools: Optional[List[ToolDefinition]] = None,
        temperature: float = 0.2,
    ) -> Dict[str, Any]:
        """统一返回结构: {"text": str, "tool_calls": List[dict]}"""
        pass

import anthropic


class ClaudeAdapter(ModelAdapter):

    def __init__(self, api_key: str, model: str = "claude-3-5-sonnet-20241022"):
        self.client = anthropic.Anthropic(api_key=api_key)
        self.model = model

    def generate(
        self,
        system_prompt: str,
        messages: List[AgentMessage],
        tools: Optional[List[ToolDefinition]] = None,
        temperature: float = 0.2,
    ) -> Dict[str, Any]:
        # 1. 协议转换：映射成 Anthropic 格式
        ant_messages = [{"role": m.role, "content": m.content} for m in messages]
        ant_tools = []
        if tools:
            for t in tools:
                ant_tools.append({
                    "name": t.name,
                    "description": t.description,
                    "input_schema": t.parameters_schema,
                })

        # 2. 原厂调用
        response = self.client.messages.create(
            model=self.model,
            max_tokens=2048,
            system=system_prompt,
            messages=ant_messages,
            tools=ant_tools if ant_tools else None,
            temperature=temperature,
        )

        # 3. 产物规范化回传
        tool_calls = []
        text_out = ""
        for block in response.content:
            if block.type == "tool_use":
                tool_calls.append({
                    "id": block.id,
                    "name": block.name,
                    "args": block.input,
                })
            elif block.type == "text":
                text_out += block.text

        return {
            "text": text_out,
            "tool_calls": tool_calls,
            "stop_reason": response.stop_reason,
        }