from fastapi import APIRouter, status, Depends, BackgroundTasks, HTTPException
from fastapi.responses import JSONResponse
from src.utils.logger import logger
from src.agents.role_play.func import create_agents
from pydantic import BaseModel, Field
from typing import List, Dict, Any
from src.agents.role_play.scenarios import get_scenarios, get_scenario_by_id

router = APIRouter(prefix="/ai", tags=["AI"])


class RoleplayRequest(BaseModel):
    query: str = Field(..., description="User's query for the AI agent")
    session_id: str = Field(
        ..., description="Session ID for tracking user interactions"
    )
    scenario: Dict[str, Any] = Field(..., description="The scenario for the roleplay")


@router.get("/scenarios", status_code=status.HTTP_200_OK)
async def list_scenarios():
    return JSONResponse(content=get_scenarios())


@router.post("/roleplay", status_code=status.HTTP_200_OK)
async def roleplay(request: RoleplayRequest):
    scenario = request.scenario
    if not scenario:
        raise HTTPException(status_code=400, detail="Scenario not provided")
    print("alo", scenario)
    print(create_agents(scenario).get_state({"configurable": {"thread_id": request.session_id}}))
    response = await create_agents(scenario).ainvoke(
        {
            "messages": [request.query],
        },
        {"configurable": {"thread_id": request.session_id}},
    )

    return JSONResponse(content=response["messages"][-1].content)