import os
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from qwen_agent.agents import Assistant

# Load environment variables from .env file
load_dotenv()

# Get API key from environment variables
api_key = os.getenv("API_KEY")
if not api_key:
    raise ValueError("API_KEY environment variable not set.")

app = FastAPI()

class MainAgent:
    def __init__(self):
        # Initialize Qwen-Agent here
        # For demonstration, we'll use a simple AssistantAgent
        self.agent = Assistant(
            llm=dict(model='gemini-2.5-flash-preview-05-20', model_server='https://gmn.ainet.ltd/v1', api_key=api_key), # Use OpenAI compatible API
            name='assistant',
            description='A helpful assistant agent.'
        )

    def process_query(self, query: str): # 更改回 def
        # Process the query using the agent
        # Wrap the query into the expected messages format for qwen-agent
        messages = [{'role': 'user', 'content': query}]
        # print('messages:', messages)
        final_response_content = "" # 用于存储最终的响应内容
        for response_item in self.agent.run(messages=messages):
            # print('response_item:', response_item) # Added for debugging
            # qwen-agent's run method yields a list of dicts, or sometimes a single dict
            if isinstance(response_item, list):
                for chunk in response_item:
                    if isinstance(chunk, dict) and 'content' in chunk:
                        # 每次都累加内容
                        final_response_content = chunk['content']
            elif isinstance(response_item, dict) and 'content' in response_item:
                # 每次都累加内容
                final_response_content = response_item['content']
            # 暂时不处理其他类型的 response_item，例如 tool_calls，因为当前目标是获取文本响应
        return final_response_content # 返回最终的完整内容

main_agent = MainAgent()

@app.get("/")
async def read_root():
    return {"message": "Hello, FastAPI!"}

@app.post("/agent_chat")
async def agent_chat(query: dict):
    user_query = query.get("query")
    print(f"Received query: {user_query}")
    if not user_query:
        return JSONResponse(content={"error": "Query parameter is missing"}, status_code=400)
    
    agent_response = main_agent.process_query(user_query)
    return {"response": agent_response}