from typing import Optional

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from gbase_agent import init_agent_service

app = FastAPI(title="Database Assistant API", version="1.0.0")

# Initialize agent globally at startup
bot = init_agent_service()


class QueryRequest(BaseModel):
    question: str
    file_url: Optional[str] = None


class QueryResponse(BaseModel):
    answer: str


@app.post("/query", response_model=QueryResponse)
async def query_database(request: QueryRequest):
    """
    Process a database query using the assistant agent.

    Args:
        request: QueryRequest containing the query and optional file URL

    Returns:
        QueryResponse containing the assistant's response
    """
    try:
        messages = []

        if request.file_url:
            messages.append(
                {
                    "role": "user",
                    "content": [{"text":"使用sqlite数据库，用日语回答下面问题："+request.question}, {"file": request.file_url}],
                }
            )
        else:
            messages.append({"role": "user", "content": request.question})

        responses = []
        for response in bot.run(messages):
            responses.append(response)

        if responses:
            final_response = responses[-1][-1]
            return QueryResponse(answer=final_response["content"])
        else:
            raise HTTPException(status_code=500, detail="No response from agent")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")


@app.get("/")
async def root():
    """Health check endpoint"""
    return {"message": "Database Assistant API is running"}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)