from __future__ import annotations

from fastapi import APIRouter

from app.modules.api.controllers.request_controller import RequestController
from app.modules.api.controllers.session_controller import SessionController
from app.modules.api.controllers.stream_controller import StreamController
from app.schemas.agent_api import (
    AgentRequestCreateRequest,
    AgentRequestQueuedResponse,
    AgentRequestStateResponse,
    BindRagSessionRequest,
    BindRagSessionResponse,
    CreateAgentSessionResponse,
    ResetAgentSessionResponse,
)


def build_public_router(
    sessions: SessionController,
    requests: RequestController,
    streams: StreamController,
) -> APIRouter:
    router = APIRouter(tags=["agent-api"])

    @router.post("/api/agent/sessions", response_model=CreateAgentSessionResponse)
    async def create_session() -> CreateAgentSessionResponse:
        return sessions.create_session()

    @router.post("/api/agent/sessions/{session_id}/rag", response_model=BindRagSessionResponse)
    async def bind_rag_session(session_id: str, request: BindRagSessionRequest) -> BindRagSessionResponse:
        return sessions.bind_rag_session(session_id, request)

    @router.post("/api/agent/sessions/{session_id}/reset", response_model=ResetAgentSessionResponse)
    async def reset_session(session_id: str) -> ResetAgentSessionResponse:
        return sessions.reset_session(session_id)

    @router.post("/api/agent/requests", response_model=AgentRequestQueuedResponse)
    async def create_request(request: AgentRequestCreateRequest) -> AgentRequestQueuedResponse:
        return await requests.create_request(request)

    @router.get("/api/agent/requests/{request_id}", response_model=AgentRequestStateResponse)
    async def get_request(request_id: str) -> AgentRequestStateResponse:
        return requests.get_request(request_id)

    @router.get("/api/agent/streams/{request_id}")
    async def stream(request_id: str):
        return await streams.stream(request_id)

    return router
