from typing import Any
import uvicorn
import json
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Body
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import pandas as pd
from datasets import load_from_disk
from transformers import AutoTokenizer, TFAutoModel
from dotenv import load_dotenv

load_dotenv()

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = TFAutoModel.from_pretrained(model_ckpt, from_pt=True)

def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]
    
def get_embeddings(text_list):
    encoded_input = tokenizer(
        text_list, padding=True, truncation=True, return_tensors="tf"
    )
    encoded_input = {k: v for k, v in encoded_input.items()}
    model_output = model(**encoded_input)
    return cls_pooling(model_output)  

embeddings_dataset = load_from_disk("data")
embeddings_dataset.add_faiss_index(column="embeddings")

def recommendations(question): 
    question_embedding = get_embeddings([question]).numpy()
    scores, samples = embeddings_dataset.get_nearest_examples(
        "embeddings", question_embedding, k=5
    )
    samples_df = pd.DataFrame.from_dict(samples)
    samples_df["scores"] = scores
    samples_df.sort_values("scores", ascending=False, inplace=True,ignore_index=True)
    return samples_df[['drugName', 'review', 'scores']]

@app.post("/recommend")
async def upload_image(question: str = Body(...,embed=True)):
    custom_recommendation_result = recommendations(question)
    custom_recommendation_result = custom_recommendation_result.to_dict(orient='records')
    custom_recommendation_result = json.dumps(custom_recommendation_result)
    return JSONResponse({"data":json.loads(custom_recommendation_result)},status_code=200)

# @app.post("/drugs")
# async def upload_image(drug: str = Form(...)):
#     recommendation_result = recommendations(drug)
#     recommendation_result = recommendation_result.to_dict(orient='records')
#     recommendation_result = json.dumps(recommendation_result)
#     return JSONResponse({"data":json.loads(recommendation_result)},status_code=200)


if __name__ == '__main__':
    uvicorn.run(app, host="127.0.0.1", port=5000)