app/routers/admin.py

# app/routers/admin.py
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
from typing import List
import os
from app.database.database_query import DatabaseQuery
from app.services.vector_database_search import VectorDatabaseSearch
from app.middleware.auth import get_current_user
from pydantic import BaseModel

router = APIRouter()
vector_db = VectorDatabaseSearch()
query = DatabaseQuery()
TEMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'temp')
os.makedirs(TEMP_DIR, exist_ok=True)

class SearchQuery(BaseModel):
    query: str
    k: int = 5

@router.get('/books')
async def get_books(username: str = Depends(get_current_user)):
    try:
        book_info = vector_db.get_book_info()
        return {
            'status': 'success',
            'data': book_info
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post('/books', status_code=201)
async def add_books(files: List[UploadFile] = File(...), username: str = Depends(get_current_user)):
    try:
        pdf_paths = []
        for file in files:
            if file.filename.endswith('.pdf'):
                safe_filename = os.path.basename(file.filename)
                temp_path = os.path.join(TEMP_DIR, safe_filename)
                
                with open(temp_path, "wb") as buffer:
                    content = await file.read()
                    buffer.write(content)
                
                pdf_paths.append(temp_path)

        if not pdf_paths:
            raise HTTPException(status_code=400, detail="No valid PDF files provided")

        success_count = 0
        for pdf_path in pdf_paths:
            if vector_db.add_pdf(pdf_path):
                success_count += 1

        # Clean up temporary files
        for path in pdf_paths:
            try:
                if os.path.exists(path):
                    os.remove(path)
            except Exception:
                pass

        return {
            'status': 'success',
            'message': f'Successfully added {success_count} of {len(pdf_paths)} books'
        }

    except Exception as e:
        # Clean up temporary files in case of error
        for path in pdf_paths:
            try:
                if os.path.exists(path):
                    os.remove(path)
            except:
                pass
        
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=str(e))

@router.post('/search')
async def search_books(search_data: SearchQuery, username: str = Depends(get_current_user)):
    try:
        query_text = search_data.query
        k = search_data.k

        results = vector_db.search(
            query=query_text,
            top_k=k
        )

        return {
            'status': 'success',
            'data': results
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))