app/api/pdf_routes.py

import os
import shutil
import uuid
from fastapi import APIRouter, UploadFile, File, Form, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
from typing import Optional, List, Dict, Any

from app.utils.pdf_processor import PDFProcessor
from app.models.pdf_models import PDFResponse, DeleteDocumentRequest, DocumentsListResponse
from app.api.pdf_websocket import (
    send_pdf_upload_started, 
    send_pdf_upload_progress, 
    send_pdf_upload_completed,
    send_pdf_upload_failed,
    send_pdf_delete_started,
    send_pdf_delete_completed,
    send_pdf_delete_failed
)

# Khởi tạo router
router = APIRouter(
    prefix="/pdf",
    tags=["PDF Processing"],
)

# Thư mục lưu file tạm - sử dụng /tmp để tránh lỗi quyền truy cập
TEMP_UPLOAD_DIR = "/tmp/uploads/temp"
STORAGE_DIR = "/tmp/uploads/pdfs"

# Đảm bảo thư mục upload tồn tại
os.makedirs(TEMP_UPLOAD_DIR, exist_ok=True)
os.makedirs(STORAGE_DIR, exist_ok=True)

# Endpoint upload và xử lý PDF
@router.post("/upload", response_model=PDFResponse)
async def upload_pdf(
    file: UploadFile = File(...),
    namespace: str = Form("Default"),
    index_name: str = Form("testbot768"),
    title: Optional[str] = Form(None),
    description: Optional[str] = Form(None),
    user_id: Optional[str] = Form(None),
    background_tasks: BackgroundTasks = None
):
    """
    Upload và xử lý file PDF để tạo embeddings và lưu vào Pinecone
    
    - **file**: File PDF cần xử lý
    - **namespace**: Namespace trong Pinecone để lưu embeddings (mặc định: "Default")
    - **index_name**: Tên index Pinecone (mặc định: "testbot768")
    - **title**: Tiêu đề của tài liệu (tùy chọn)
    - **description**: Mô tả về tài liệu (tùy chọn)
    - **user_id**: ID của người dùng để cập nhật trạng thái qua WebSocket
    """
    try:
        # Kiểm tra file có phải PDF không
        if not file.filename.lower().endswith('.pdf'):
            raise HTTPException(status_code=400, detail="Chỉ chấp nhận file PDF")
        
        # Tạo file_id và lưu file tạm
        file_id = str(uuid.uuid4())
        temp_file_path = os.path.join(TEMP_UPLOAD_DIR, f"{file_id}.pdf")
        
        # Gửi thông báo bắt đầu xử lý qua WebSocket nếu có user_id
        if user_id:
            await send_pdf_upload_started(user_id, file.filename, file_id)
        
        # Lưu file
        with open(temp_file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        # Tạo metadata
        metadata = {
            "filename": file.filename,
            "content_type": file.content_type
        }
        
        if title:
            metadata["title"] = title
        if description:
            metadata["description"] = description
        
        # Gửi thông báo tiến độ qua WebSocket
        if user_id:
            await send_pdf_upload_progress(
                user_id, 
                file_id, 
                "file_preparation", 
                0.2, 
                "File saved, preparing for processing"
            )
            
        # Khởi tạo PDF processor
        processor = PDFProcessor(index_name=index_name, namespace=namespace)
        
        # Gửi thông báo bắt đầu embedding qua WebSocket
        if user_id:
            await send_pdf_upload_progress(
                user_id, 
                file_id, 
                "embedding_start", 
                0.4, 
                "Starting to process PDF and create embeddings"
            )
        
        # Xử lý PDF và tạo embeddings
        # Tạo callback function để xử lý cập nhật tiến độ
        async def progress_callback_wrapper(step, progress, message):
            if user_id:
                await send_progress_update(user_id, file_id, step, progress, message)
        
        # Xử lý PDF và tạo embeddings với callback đã được xử lý đúng cách
        result = await processor.process_pdf(
            file_path=temp_file_path,
            document_id=file_id,
            metadata=metadata,
            progress_callback=progress_callback_wrapper
        )
        
        # Nếu thành công, chuyển file vào storage
        if result.get('success'):
            storage_path = os.path.join(STORAGE_DIR, f"{file_id}.pdf")
            shutil.move(temp_file_path, storage_path)
            
            # Gửi thông báo hoàn thành qua WebSocket
            if user_id:
                await send_pdf_upload_completed(
                    user_id,
                    file_id,
                    file.filename,
                    result.get('chunks_processed', 0)
                )
        else:
            # Gửi thông báo lỗi qua WebSocket
            if user_id:
                await send_pdf_upload_failed(
                    user_id,
                    file_id,
                    file.filename,
                    result.get('error', 'Unknown error')
                )
            
        # Dọn dẹp: xóa file tạm nếu vẫn còn
        if os.path.exists(temp_file_path):
            os.remove(temp_file_path)
            
        return result
    except Exception as e:
        # Dọn dẹp nếu có lỗi
        if 'temp_file_path' in locals() and os.path.exists(temp_file_path):
            os.remove(temp_file_path)
            
        # Gửi thông báo lỗi qua WebSocket
        if 'user_id' in locals() and user_id and 'file_id' in locals():
            await send_pdf_upload_failed(
                user_id,
                file_id,
                file.filename,
                str(e)
            )
            
        return PDFResponse(
            success=False,
            error=str(e)
        )

# Function để gửi cập nhật tiến độ - được sử dụng trong callback
async def send_progress_update(user_id, document_id, step, progress, message):
    if user_id:
        await send_pdf_upload_progress(user_id, document_id, step, progress, message)

# Endpoint xóa tài liệu
@router.delete("/namespace", response_model=PDFResponse)
async def delete_namespace(
    namespace: str = "Default",
    index_name: str = "testbot768",
    user_id: Optional[str] = None
):
    """
    Xóa toàn bộ embeddings trong một namespace từ Pinecone (tương ứng xoá namespace)

    - **namespace**: Namespace trong Pinecone (mặc định: "Default")
    - **index_name**: Tên index Pinecone (mặc định: "testbot768")
    - **user_id**: ID của người dùng để cập nhật trạng thái qua WebSocket
    """
    try:
        # Gửi thông báo bắt đầu xóa qua WebSocket
        if user_id:
            await send_pdf_delete_started(user_id, namespace)
            
        processor = PDFProcessor(index_name=index_name, namespace=namespace)
        result = await processor.delete_namespace()
        
        # Gửi thông báo kết quả qua WebSocket
        if user_id:
            if result.get('success'):
                await send_pdf_delete_completed(user_id, namespace)
            else:
                await send_pdf_delete_failed(user_id, namespace, result.get('error', 'Unknown error'))
                
        return result
    except Exception as e:
        # Gửi thông báo lỗi qua WebSocket
        if user_id:
            await send_pdf_delete_failed(user_id, namespace, str(e))
            
        return PDFResponse(
            success=False,
            error=str(e)
        )

# Endpoint lấy danh sách tài liệu
@router.get("/documents", response_model=DocumentsListResponse)
async def get_documents(namespace: str = "Default", index_name: str = "testbot768"):
    """
    Lấy thông tin về tất cả tài liệu đã được embed
    
    - **namespace**: Namespace trong Pinecone (mặc định: "Default")
    - **index_name**: Tên index Pinecone (mặc định: "testbot768")
    """
    try:
        # Khởi tạo PDF processor
        processor = PDFProcessor(index_name=index_name, namespace=namespace)
        
        # Lấy danh sách documents
        result = await processor.list_documents()
        
        return result
    except Exception as e:
        return DocumentsListResponse(
            success=False,
            error=str(e)
        )