import os
import json
import gradio as gr
import pandas as pd

from openai import OpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS

API_KEY = os.getenv("YOUR_KEY")
PASSWORD = os.getenv("PASSWORD")
MODEL = "gpt-4o"

client = OpenAI(api_key = API_KEY)
embeddings = OpenAIEmbeddings(model = "text-embedding-3-large", api_key = API_KEY)
yt_chunks = FAISS.load_local("vector-large", embeddings, allow_dangerous_deserialization = True)
df = pd.read_csv("data/ko-youtube-trans-U10k.csv")

def find_docs(message):
    
    finding_docs = yt_chunks.similarity_search(message, k = 5)
    
    indices = [doc.metadata['row'] for doc in finding_docs]
    
    retrievers = [json.loads(df.loc[idx].to_json(force_ascii = False)) for idx in indices]
    
    return retrievers

def predict(message, history):
    
    openai_input = list()
    
    retriever = find_docs(message)
    
    system_prompt = """- You are an AI chat bot that recommends YouTube content to users as an assistant.\n- You were created and powered by 'bigster (빅스터)', an AI & bigdata expert company.\n- Recommend YouTube content to users based on what's in “retriever”.\n- If the user's question is not related to content recommendations, please display a message declining to answer.\n- You must recommend at least 3 YouTube content items to the user based on the information in the 'retriever'. Be sure to explicitly include 'url' & 'videoChannelName' & 'videoName' information in your response. Also, for each featured piece of content, summarize what's in the 'transcription' and present it to the user. Use the following Markdown format to create hyperlinks: '[videoName](url)'\n\n retriever:\n{retriever}"""
    
    for human, assistant in history:
        openai_input.append({"role": "user", "content": human})
        openai_input.append({"role": "assistant", "content": assistant})
        
    openai_input = [item for item in openai_input if item['role'] != "system"]
    
    openai_input.append({"role": "system", "content": system_prompt.format(retriever = retriever)})
    openai_input.append({"role": "user", "content": message})
    
    response = client.chat.completions.create(
        model = MODEL,
        messages = openai_input,
        temperature = 1.0,
        stream = True
    )

    partial_message = ""
    for chunk in response:
        if chunk.choices[0].delta.content is not None:
              partial_message = partial_message + chunk.choices[0].delta.content
              yield partial_message

    print(openai_input)
    
gr.ChatInterface(
    predict, 
    title = "YOUTUBE REC",
    theme = gr.themes.Soft(primary_hue = "purple"),
    examples = [
        "파이썬 프로그래밍 언어를 독학하기 위한 영상을 추천해줘.",
    ]
).launch(share = True, auth = ("user", PASSWORD))