import os
import functools
from typing import List

from environs import Env
from fastapi import FastAPI, HTTPException
from fastapi_utils.api_model import APIModel

from local_embedding import LocalEmbeddings

env = Env()
env.read_env(".env")

MODEL_PATH = env.str("MODEL_PATH", "/data/embedding_models")
if not MODEL_PATH.endswith("/"):
    MODEL_PATH += "/"
MODEL_NAMES = ["pingchat-embedding"]


@functools.lru_cache
def get_model(model_name: str):
    if model_name not in MODEL_NAMES:
        raise Exception(f"{model_name} is not supported currently")
    return LocalEmbeddings(os.path.join(MODEL_PATH, model_name))


class EmbeddingRequest(APIModel):
    text: str
    prefix: str = "[QRY]"
    suffix: str = ""


class BatchEmbeddingRequest(APIModel):
    texts: List[str]
    prefix: str = "[DOC]"
    suffix: str = ""


class EmbeddingModel(APIModel):
    name: str
    path: str
    embedding_dimension: int
    description: str = ""


app = FastAPI()


# List models endpoint
@app.get("/embedding-models")
def list_models() -> List[EmbeddingModel]:
    models = []
    for name in MODEL_NAMES:
        model = get_model(name)
        if not model:
            continue
        models.append(
            EmbeddingModel(
                name=name,
                path=MODEL_PATH + name,
                embedding_dimension=model.embedding.get_sentence_embedding_dimension(),
                description=model.description
            )
        )
    return models


# Embedding endpoint
@app.post("/embedding-models/{model_name}")
def embedding(model_name: str, request: EmbeddingRequest) -> List[float]:
    texts = [request.text]
    model = get_model(model_name)

    try:
        return model.embed(texts=texts, prefix=request.prefix, suffix=request.suffix)[0]
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e)) from e


@app.post("/embedding-models/{model_name}/batch")
def batch_embedding(
    model_name: str, request: BatchEmbeddingRequest
) -> List[List[float]]:
    model = get_model(model_name)

    try:
        return model.embed(
            texts=request.texts, prefix=request.prefix, suffix=request.suffix
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e)) from e


@app.get("/ping")
def ping():
    return {"status": "success"}