from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound, TranscriptsDisabled, CouldNotRetrieveTranscript, TranslationLanguageNotAvailable, NotTranslatable
from fastapi import FastAPI
from pydantic import BaseModel
from functools import lru_cache
import pandas as pd
import faiss
from sentence_transformers import SentenceTransformer
import logging

logging.basicConfig(level=logging.INFO)

app = FastAPI()

class SearchRequest(BaseModel):
    query: str
    top_k: int = 3

@lru_cache(maxsize=1)
def load_model():
    return SentenceTransformer('model/BGE-Medico-Model')

@lru_cache(maxsize=1)
def load_index():
    return faiss.read_index('model/device_index.faiss')

@lru_cache(maxsize=1)
def load_data():
    return pd.read_csv('model/Final_devices.csv')

@app.post("/MedInsSch")
def search_device(request: SearchRequest):
    df = load_data()
    model = load_model()
    index = load_index()
    
    cols = ['name', 'definition', 'instructions']
    query = request.query
    top_k = request.top_k
    
    if query.isdigit():
        results = df[cols][df["code"] == int(query)]
        if not results.empty:
            return {"results": results.to_dict(orient="records")}
    
    query_vec = model.encode([query])
    D, I = index.search(query_vec, top_k)
    results = [df.iloc[idx][cols].to_dict() for idx in I[0]]
    
    return {"results": results}


@app.get("/")
def root():
    return {"message": "MedInstruct, Youtube Transcript Search API is running"}