api.py

"""
FastAPI application — REST API chuẩn production cho Study Group Assistant.

Khởi động:
    uvicorn src.api:app --host 0.0.0.0 --port 8000 --reload
"""

import logging
import os
import uvicorn
import tempfile
import uuid
from asyncio import get_running_loop
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from datetime import datetime, timezone

from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from src.core import final_answer
from src.qdrant_store import get_custom_prompt, save_custom_prompt
from src.redis_client import redis_client

logger = logging.getLogger(__name__)

_executor = ThreadPoolExecutor()


# ── Lifespan ──────────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Study Group Assistant API starting up.")
    yield
    _executor.shutdown(wait=False)
    logger.info("Study Group Assistant API shut down.")


# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Study Group Assistant API",
    description=(
        "AI agent giúp nhóm học tập tóm tắt hội thoại, "
        "tra cứu lịch trình và quản lý ghi nhớ."
    ),
    version="1.0.0",
    lifespan=lifespan,
    docs_url="/docs",
    redoc_url="/redoc",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Middlewares ───────────────────────────────────────────────────────────────

@app.middleware("http")
async def attach_request_id(request: Request, call_next):
    """Gắn X-Request-ID vào mỗi request để dễ trace log."""
    request_id = str(uuid.uuid4())
    request.state.request_id = request_id
    response = await call_next(request)
    response.headers["X-Request-ID"] = request_id
    return response


@app.middleware("http")
async def log_requests(request: Request, call_next):
    """Log method, path và status code của mỗi request."""
    response = await call_next(request)
    logger.info(
        "%s %s → %d [rid=%s]",
        request.method,
        request.url.path,
        response.status_code,
        getattr(request.state, "request_id", "-"),
    )
    return response


# ── Pydantic models ───────────────────────────────────────────────────────────

class ChatRequest(BaseModel):
    conversation_id: str = Field(..., description="ID cuộc hội thoại DM")
    sender_id: str = Field(..., description="ID hoặc tên người gửi")
    query: str = Field(..., description="Câu hỏi hoặc yêu cầu")

    model_config = {
        "json_schema_extra": {
            "example": {
                "conversation_id": "98996225-512c-4491-96a2-bc71552328ca",
                "sender_id": "@Hoang",
                "query": "Tóm tắt cuộc trò chuyện hôm nay",
            }
        }
    }


class ChatResponse(BaseModel):
    answer: str = Field(..., description="Câu trả lời từ agent")
    processing_time: str = Field(..., description="Thời gian xử lý, ví dụ '1.23s'")
    conversation_id: str
    sender_id: str


class HealthComponent(BaseModel):
    status: str = Field(..., description="'ok' | 'degraded' | 'down'")
    detail: str = ""


class HealthResponse(BaseModel):
    status: str = Field(..., description="'ok' | 'degraded'")
    timestamp: str
    components: dict[str, HealthComponent]


class ErrorDetail(BaseModel):
    error: str
    detail: str = ""
    request_id: str = ""


class CustomPromptRequest(BaseModel):
    user_id: str = Field(..., description="ID người dùng")
    prompt: str  = Field(..., description="Nội dung custom prompt")

    model_config = {
        "json_schema_extra": {
            "example": {
                "user_id": "@Hoang",
                "prompt": "Luôn trả lời ngắn gọn trong 3 câu. Dùng bullet point khi liệt kê.",
            }
        }
    }


class CustomPromptResponse(BaseModel):
    success: bool
    user_id: str
    prompt:  str


# ── Helper ────────────────────────────────────────────────────────────────────

def _request_id(request: Request) -> str:
    return getattr(request.state, "request_id", "")


def _utcnow() -> str:
    return datetime.now(timezone.utc).isoformat()


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/", include_in_schema=False)
async def root():
    return {
        "service": "Study Group Assistant API",
        "version": "1.0.0",
        "docs": "/docs",
        "health": "/health",
    }


@app.get(
    "/health",
    response_model=HealthResponse,
    summary="Kiểm tra trạng thái hệ thống",
    tags=["System"],
)
async def health():
    redis_ok = redis_client.ping()
    return HealthResponse(
        status="ok" if redis_ok else "degraded",
        timestamp=_utcnow(),
        components={
            "redis": HealthComponent(
                status="ok" if redis_ok else "down",
                detail="Connected" if redis_ok else "Connection failed — using local fallback",
            ),
            "agent": HealthComponent(status="ok"),
        },
    )


@app.post(
    "/api/v1/chat",
    response_model=ChatResponse,
    status_code=status.HTTP_200_OK,
    summary="Gửi câu hỏi đến Agent",
    tags=["Agent"],
    responses={
        422: {"model": ErrorDetail, "description": "Tham số không hợp lệ"},
        500: {"model": ErrorDetail, "description": "Lỗi xử lý nội bộ"},
        503: {"model": ErrorDetail, "description": "Dịch vụ tạm thời không khả dụng"},
    },
)
async def chat(request: Request, body: ChatRequest):
    """
    Gửi query đến agent, nhận câu trả lời và thời gian xử lý.

    Agent sẽ tự động:
    - Phân loại yêu cầu (trả lời trực tiếp hoặc tra cứu hội thoại)
    - Gọi các tool phù hợp (tóm tắt, lịch trình, ghi nhớ, web...)
    - Tổng hợp kết quả thành câu trả lời tự nhiên
    """
    loop = get_running_loop()
    try:
        answer, elapsed = await loop.run_in_executor(
            _executor,
            lambda: final_answer(body.conversation_id, body.sender_id, body.query),
        )
    except ValueError as e:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=str(e),
        )
    except Exception as e:
        logger.exception(
            "Unhandled error in POST /api/v1/chat [rid=%s]", _request_id(request)
        )
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Lỗi xử lý nội bộ. Vui lòng thử lại.",
        )

    return ChatResponse(
        answer=answer,
        processing_time=elapsed,
        conversation_id=body.conversation_id,
        sender_id=body.sender_id,
    )


@app.post(
    "/api/v1/chat_with_pdf",
    response_model=ChatResponse,
    status_code=status.HTTP_200_OK,
    summary="Gửi câu hỏi kèm file PDF đến Agent",
    tags=["Agent"],
    responses={
        400: {"model": ErrorDetail, "description": "File không phải PDF"},
        422: {"model": ErrorDetail, "description": "Tham số không hợp lệ"},
        500: {"model": ErrorDetail, "description": "Lỗi xử lý nội bộ"},
    },
)
async def chat_with_pdf(
    request: Request,
    conversation_id: str = Form(..., description="ID cuộc hội thoại DM"),
    sender_id: str = Form(..., description="ID hoặc tên người gửi"),
    query: str = Form(..., description="Câu hỏi hoặc yêu cầu về nội dung PDF"),
    file: UploadFile = File(..., description="File PDF cần xử lý"),
):
    if not file.filename.lower().endswith(".pdf"):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Chỉ chấp nhận file PDF.",
        )

    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
            tmp.write(await file.read())
            tmp_path = tmp.name

        loop = get_running_loop()
        answer, elapsed = await loop.run_in_executor(
            _executor,
            lambda: final_answer(conversation_id, sender_id, query, pdf_path=tmp_path),
        )
    except ValueError as e:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=str(e),
        )
    except Exception:
        logger.exception(
            "Unhandled error in POST /api/v1/chat_with_pdf [rid=%s]", _request_id(request)
        )
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Lỗi xử lý nội bộ. Vui lòng thử lại.",
        )
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)

    return ChatResponse(
        answer=answer,
        processing_time=elapsed,
        conversation_id=conversation_id,
        sender_id=sender_id,
    )


@app.post(
    "/api/v1/custom_prompt",
    response_model=CustomPromptResponse,
    status_code=status.HTTP_200_OK,
    summary="Lưu custom prompt cho người dùng",
    tags=["Agent"],
    responses={
        500: {"model": ErrorDetail, "description": "Lỗi lưu prompt"},
    },
)
async def set_custom_prompt(request: Request, body: CustomPromptRequest):
    """
    Lưu hoặc cập nhật custom prompt của người dùng lên Qdrant.
    Prompt này sẽ được tự động inject vào system prompt khi user đó gửi query.
    """
    loop = get_running_loop()
    ok = await loop.run_in_executor(
        _executor,
        lambda: save_custom_prompt(body.user_id, body.prompt),
    )
    if not ok:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Không thể lưu custom prompt. Kiểm tra cấu hình QDRANT_URL.",
        )
    return CustomPromptResponse(success=True, user_id=body.user_id, prompt=body.prompt)


@app.get(
    "/api/v1/custom_prompt/{user_id}",
    response_model=CustomPromptResponse,
    status_code=status.HTTP_200_OK,
    summary="Lấy custom prompt của người dùng",
    tags=["Agent"],
    responses={
        404: {"model": ErrorDetail, "description": "Không tìm thấy prompt"},
    },
)
async def get_user_custom_prompt(user_id: str, request: Request):
    loop = get_running_loop()
    prompt = await loop.run_in_executor(
        _executor,
        lambda: get_custom_prompt(user_id),
    )
    if prompt is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Không tìm thấy custom prompt cho user '{user_id}'.",
        )
    return CustomPromptResponse(success=True, user_id=user_id, prompt=prompt)


_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}


@app.post(
    "/api/v1/chat_with_image",
    response_model=ChatResponse,
    status_code=status.HTTP_200_OK,
    summary="Gửi câu hỏi kèm ảnh đến Agent",
    tags=["Agent"],
    responses={
        400: {"model": ErrorDetail, "description": "File không phải ảnh hợp lệ"},
        422: {"model": ErrorDetail, "description": "Tham số không hợp lệ"},
        500: {"model": ErrorDetail, "description": "Lỗi xử lý nội bộ"},
    },
)
async def chat_with_image(
    request: Request,
    conversation_id: str = Form(..., description="ID cuộc hội thoại DM"),
    sender_id: str = Form(..., description="ID hoặc tên người gửi"),
    query: str = Form(..., description="Câu hỏi hoặc yêu cầu về nội dung ảnh"),
    file: UploadFile = File(..., description="File ảnh cần xử lý"),
):
    ext = os.path.splitext(file.filename.lower())[1]
    if ext not in _IMAGE_EXTENSIONS:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Chỉ chấp nhận ảnh: {', '.join(_IMAGE_EXTENSIONS)}.",
        )

    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp:
            tmp.write(await file.read())
            tmp_path = tmp.name

        loop = get_running_loop()
        answer, elapsed = await loop.run_in_executor(
            _executor,
            lambda: final_answer(conversation_id, sender_id, query, image_path=tmp_path),
        )
    except ValueError as e:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=str(e),
        )
    except Exception:
        logger.exception(
            "Unhandled error in POST /api/v1/chat_with_image [rid=%s]", _request_id(request)
        )
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Lỗi xử lý nội bộ. Vui lòng thử lại.",
        )
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)

    return ChatResponse(
        answer=answer,
        processing_time=elapsed,
        conversation_id=conversation_id,
        sender_id=sender_id,
    )


# ── Exception handlers ────────────────────────────────────────────────────────

@app.exception_handler(404)
async def not_found(request: Request, exc):
    return JSONResponse(
        status_code=404,
        content=ErrorDetail(
            error="Not Found",
            detail=f"Endpoint '{request.url.path}' không tồn tại.",
            request_id=_request_id(request),
        ).model_dump(),
    )


@app.exception_handler(405)
async def method_not_allowed(request: Request, exc):
    return JSONResponse(
        status_code=405,
        content=ErrorDetail(
            error="Method Not Allowed",
            detail=f"Method '{request.method}' không được hỗ trợ tại '{request.url.path}'.",
            request_id=_request_id(request),
        ).model_dump(),
    )


@app.exception_handler(500)
async def internal_error(request: Request, exc):
    return JSONResponse(
        status_code=500,
        content=ErrorDetail(
            error="Internal Server Error",
            detail="Đã xảy ra lỗi không mong muốn.",
            request_id=_request_id(request),
        ).model_dump(),
    )


if __name__ == "__main__":
    uvicorn.run("src.api:app", host="127.0.0.1", port=8000, reload=True)