from fastapi import FastAPI, Depends, HTTPException
from pydantic import BaseModel
from fastapi_health import health
from PIL import Image
import logging
import sys
from io import BytesIO
import base64

from process_img import Image_Processor
from vector_emb import EmbeddingModels

logger = logging.getLogger(__name__)

logging.basicConfig(
    level=logging.getLevelName("INFO"),
    handlers=[logging.StreamHandler(sys.stdout)],
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logging.info('Logging module started')

def get_session():
    return True

def is_database_online(session: bool = Depends(get_session)):
    return session

app = FastAPI()
app.add_api_route("/healthz", health([is_database_online]))



model = EmbeddingModels()
img_Processor = Image_Processor()

class ImageBase64(BaseModel):
    base64_string: str

class TextInput(BaseModel):
    text: str

@app.post("/design-dense/")
async def embed_image(data: ImageBase64):
    base64_string = data.base64_string
    image_data = base64.b64decode(base64_string)
    image = Image.open(BytesIO(image_data))
    final_image = img_Processor.get_processed_img(image)
    embeddings = model.get_single_image_embedding(final_image)
    return embeddings


@app.post("/sparse/")
async def embed_text(item: TextInput):
    try:
        logging.info(f'Received text for embedding: {item.text}')
        embeddings = model.get_single_sparse_text_embedding(item.text)
        logging.info('Embedding process completed')
        return embeddings
    except Exception as e:
        logging.error(f'Error during embedding process: {e}')
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/design-sparse/")
async def embed_text(item: TextInput):
    try:
        logging.info(f'Received text for embedding: {item.text}')
        embeddings = model.get_single_sparse_text_embedding(item.text)
        embeddings = model.normalize_sparse_vector_values(embeddings)
        logging.info('Embedding process completed')
        return embeddings
    except Exception as e:
        logging.error(f'Error during embedding process: {e}')
        raise HTTPException(status_code=500, detail=str(e))