from agno.agent import Agent
from agno.models.base import Model
from agno.tools import tool
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from openai import OpenAI
import os

load_dotenv()

# 定义转账请求模型
class TransferRequest(BaseModel):
    payee_name: str = Field(..., description="收款人全名")
    account_number: str = Field(..., min_length=12, max_length=12, description="12位银行账号")
    amount: float = Field(..., gt=0, le=500, description="转账金额（1-500元）")
    phone: str = Field(None, min_length=11, max_length=11, description="收款人手机号（可选）")

# 注册转账工具
@tool
def bank_transfer(request: TransferRequest):
    """执行银行转账操作，包含以下校验规则：
    1. 账号必须是12位数字
    2. 金额必须介于1-500元
    3. 手机号可选但必须是11位
    """
    # 模拟失败场景
    if request.amount > 500:
        return {"status": "failed", "reason": "单笔转账金额不能超过500元"}
    if len(request.account_number) != 12 or not request.account_number.isdigit():
        return {"status": "failed", "reason": "账号格式错误（需要12位数字）"}
    
    # 模拟成功响应
    return {
        "status": "success",
        "transaction_id": f"TX{os.urandom(4).hex().upper()}",
        "amount": request.amount
    }

# 修复后的Qwen模型实现
class QwenModel(Model):
    def __init__(self):
        self.client = OpenAI(
            base_url=os.getenv("BAILIAN_API_BASE_URL"),
            api_key=os.getenv("BAILIAN_API_KEY")
        )
        
    def invoke(self, messages, tools=None, **kwargs):
        response = self.client.chat.completions.create(
            model="qwen3-32b",
            messages=messages,
            tools=tools,
            **kwargs
        )
        return self.parse_provider_response(response)
    
    def parse_provider_response(self, response):
        return {
            "content": response.choices[0].message.content,
            "tool_calls": response.choices[0].message.tool_calls
        }
    
    # 未实现的方法抛出明确错误
    def invoke_stream(self, messages, tools=None, **kwargs):
        raise NotImplementedError("流式调用暂不支持")
    
    def ainvoke(self, messages, tools=None, **kwargs):
        raise NotImplementedError("异步调用暂不支持")
    
    def ainvoke_stream(self, messages, tools=None, **kwargs):
        raise NotImplementedError("异步流调用暂不支持")
    
    def parse_provider_response_delta(self, delta):
        return {"content": delta.get("content", "")}

# 配置Agent
agent = Agent(
    model=QwenModel(),
    tools=[bank_transfer],
    instructions="""您是银行AI助手，需要：  # 修改参数名
1. 按顺序收集：姓名→账号→金额→手机号（可选）
2. 金额超过500时主动提醒限额规则
3. 自动校验输入格式
4. 转账失败时解释具体原因"""
)

if __name__ == "__main__":
    # 启动对话循环
    while True:
        user_input = input("User: ")
        response = agent.run(user_input)
        
        # 处理转账结果
        if response.tool_responses:
            for resp in response.tool_responses:
                if resp.tool_name == "bank_transfer":
                    if resp.result["status"] == "success":
                        print(f"\n转账成功！交易号：{resp.result['transaction_id']}")
                        exit()
                    else:
                        print(f"\n错误：{resp.result['reason']}，请修改信息后重试")
        
        print(f"AI: {response.content}")