api/router/book.py

import logging
import base64

from typing import Annotated, Optional
from api.function import data_ingestion, get_data, delete_data, update_data
from api.router.user import user_dependency
from fastapi import Form, APIRouter, File, UploadFile, Depends
from fastapi.responses import JSONResponse
from db.repository import get_db_conn
from db.get_data import GetDatabase
from db.models import Category, Metadata
from db.database import get_db
from langfuse.llama_index import LlamaIndexCallbackHandler
from config import MYSQL_CONFIG
from script.vector_db import IndexManager
from service.dto import MetadataRequest, MetadataResponse
from sqlalchemy.orm import Session
from sqlalchemy.future import select
from sqlalchemy.exc import SQLAlchemyError


router = APIRouter(tags=["Book"])

db_conn = get_db_conn(MYSQL_CONFIG)
get_database = GetDatabase(db_conn)
index_manager = IndexManager()
db_dependency = Annotated[Session, Depends(get_db)]


@router.post("/book")
async def upload_file(
    user: user_dependency,
    db: db_dependency,
    title: str = Form(...),
    author: str = Form(...),
    category_id: int = Form(...),
    year: int = Form(...),
    publisher: str = Form(...),
    file: UploadFile = File(...),
    thumbnail: Optional[UploadFile] = File(None),
):
    print(user.get("role_id"))
    # if user is None or user.get('role_id') != 1:
    if user is None:
        return JSONResponse(status_code=401, content="Authentication Failed")

    # Create a new Metadata object
    new_book = Metadata(
        title=title,
        author=author,
        category_id=category_id,
        year=year,
        publisher=publisher,
    )

    db.add(new_book)
    db.commit()
    logging.info("Database Inserted")

    try:
        # Query the category based on category_id
        category_query = select(Category.category).where(Category.id == category_id)
        result = db.execute(category_query)
        category = result.scalar_one_or_none()

        # Check if the category exists
        if category is None:
            return JSONResponse(status_code=404, content="Category not found")

    except SQLAlchemyError as db_exc:
        # Handle any database-related errors (e.g., connection issues, query issues)
        print(f"Database error: {db_exc}")
        return JSONResponse(status_code=500, content="Database error occurred")

    except Exception as e:
        # Catch any other general exceptions
        print(f"Error: {e}")
        return JSONResponse(
            status_code=500, content="An error occurred while processing your request"
        )

    try:
        # Assuming you have a Langfuse callback handler
        langfuse_callback_handler = LlamaIndexCallbackHandler()
        langfuse_callback_handler.set_trace_params(
            user_id="admin_book_uploaded",
        )

        # Construct the reference dictionary
        reference = {
            "title": title,
            "author": author,
            "category": category,
            "year": year,
            "publisher": publisher,
        }

        # Process the file and handle data ingestion
        response = await data_ingestion(category_id, reference, file)

    except Exception as e:
        # Handle any errors related to file processing or data ingestion
        print(f"File processing error: {e}")
        return JSONResponse(status_code=500, content="File processing error")

    # Return a successful response with the uploaded filename and response from data ingestion
    return {
        "filename": file.filename,
        "response": response,
        "info": "upload file successfully",
    }


@router.get("/book")
async def get_metadata(user: user_dependency, db: db_dependency):
    if user is None:
        return JSONResponse(status_code=401, content="Authentication Failed")
    try:
        # Join Metadata with Category to get the category name
        results = (
            db.query(
                Metadata.id,
                Metadata.title,
                Metadata.author,
                Category.category,  # Assuming this is the correct field for category name
                Category.id,
                Metadata.year,
                Metadata.publisher,
                Metadata.thumbnail,
            )
            .join(Category, Metadata.category_id == Category.id)
            .all()
        )

        # Transform results into MetadataResponse model with optional thumbnail handling
        return [
            MetadataResponse(
                id = id,
                title=title,
                author=author,
                category=category,
                category_id = category_id,
                year=year,
                publisher=publisher,
                thumbnail=(
                    thumbnail if thumbnail else None
                ),  # Ensure None if thumbnail is not present
            )
            for id, title, author, category, category_id, year, publisher, thumbnail in results
        ]
    except SQLAlchemyError as db_exc:
        print(f"Database error: {db_exc}")
        return JSONResponse(status_code=500, content="Database error occurred")


@router.put("/book/{metadata_id}")
async def update_metadata(
    user: user_dependency,
    db: db_dependency,
    metadata_id: int,
    title: str = Form(...),
    author: str = Form(...),
    category_id: int = Form(...),
    year: int = Form(...),
    publisher: str = Form(...),
    thumbnail: Optional[UploadFile] = File(None),
):
    if user is None or user.get("role_id") != 1:
        return JSONResponse(status_code=401, content="Authentication Failed")

    try:
        # old_metadata = await get_database.get_data_by_id(metadata_id)

        old_metadata = db.execute(
            select(Metadata).where(Metadata.id == metadata_id)
        ).scalar_one_or_none()

        if old_metadata is None:
            return JSONResponse(status_code=404, content="Metadata not found.")

        # Fetch old and new categories
        old_category = db.execute(
            select(Category.category).join(Metadata).where(Metadata.id == metadata_id)
        ).scalar_one_or_none()

        new_category = db.execute(
            select(Category.category).where(Category.id == category_id)
        ).scalar_one_or_none()

        if old_category is None or new_category is None:
            return JSONResponse(status_code=404, content="Category not found.")

        # Prepare the references
        old_reference = {
            "title": old_metadata.title,
            "author": old_metadata.author,
            "category": old_category,
            "year": old_metadata.year,
            "publisher": old_metadata.publisher,
        }

        new_reference = {
            "title": title,
            "author": author,
            "category": new_category,
            "year": year,
            "publisher": publisher,
        }

        index_manager.update_vector_database(old_reference, new_reference)

        # Update existing metadata entry
        metadata = db.query(Metadata).filter(Metadata.id == metadata_id).first()

        if not metadata:
            return JSONResponse(status_code=404, content="Metadata not found")

        if thumbnail is not None:
            thumbnail_data = await thumbnail.read()  # Read the file data
            metadata.thumbnail = base64.b64encode(thumbnail_data).decode("utf-8")

        # Update fields only if provided in the request
        metadata.title = title
        metadata.author = author
        metadata.category_id = category_id
        metadata.year = year
        metadata.publisher = publisher

        # Commit the changes to the database
        db.commit()
        db.refresh(metadata)  # Refresh to get the updated data

        category_query = select(Category.category).where(
            Category.id == metadata.category_id
        )
        result = db.execute(category_query)
        category = result.scalar_one_or_none()

        return MetadataResponse(
            id=metadata_id,
            title=metadata.title,
            author=metadata.author,
            category=category,  # Assuming category relationship is available
            category_id=metadata.category_id,
            year=metadata.year,
            publisher=metadata.publisher,
            thumbnail=metadata.thumbnail if metadata.thumbnail is not None else None,
        )

    except Exception as e:
        return JSONResponse(
            status_code=500, content=f"An error occurred while updating metadata as {e}"
        )


@router.delete("/book/{metadata_id}")
async def delete_metadata(user: user_dependency, db: db_dependency, metadata_id: int):
    if user is None or user.get("role_id") != 1:
        return JSONResponse(status_code=401, content="Authentication Failed")

    try:
        old_reference = await get_database.get_data_by_id(metadata_id)
        
        old_metadata = db.execute(
            select(Metadata).where(Metadata.id == metadata_id)
        ).scalar_one_or_none()
        
        index_manager.delete_vector_database(old_reference)

        await delete_data(metadata_id, db_conn)

        return {"Status": "delete successfully"}

    except Exception as e:
        print(e)
        return JSONResponse(
            status_code=500, content="An error occurred while delete metadata"
        )