from fastapi import APIRouter, Depends
from fastapi.responses import Response

from app.application.services.ai_agent_service import AiAgentService
from app.domain.entities.user_registration import UserRegistration
from app.presentation.auth_dependencies import get_current_user
from app.presentation.company_dependencies import get_validated_company_id_query
from app.presentation.dependencies import get_ai_agent_service
from app.presentation.schemas.ai_agent import AiQueryRequest, AiQueryResponse

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


@router.post(
    "/query",
    response_model=AiQueryResponse,
    summary="Ask the AI data agent and preview rows",
)
async def query_ai_agent(
    payload: AiQueryRequest,
    company_id: str = Depends(get_validated_company_id_query),
    current_user: UserRegistration = Depends(get_current_user),
    service: AiAgentService = Depends(get_ai_agent_service),
):
    result = await service.query(current_user, company_id, payload.query.strip())
    return AiQueryResponse.model_validate(result)


@router.post(
    "/export",
    summary="Export the AI data agent result as Excel",
)
async def export_ai_agent(
    payload: AiQueryRequest,
    company_id: str = Depends(get_validated_company_id_query),
    current_user: UserRegistration = Depends(get_current_user),
    service: AiAgentService = Depends(get_ai_agent_service),
):
    content, filename = await service.export_excel(
        current_user, company_id, payload.query.strip()
    )
    return Response(
        content=content,
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )