from sentence_transformers import SentenceTransformer
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import AIMessage, HumanMessage
from langchain_chroma import Chroma
import gradio as gr
from huggingface_hub import InferenceClient
import os
import shutil

shutil.rmtree("/root/.cache", ignore_errors=True)
shutil.rmtree("/tmp", ignore_errors=True)

# Load environment variables
CHROMA_PATH = "chroma"
KEY = os.getenv("token")

# Hugging Face API setup
repo_id = "Qwen/Qwen2.5-7B-Instruct"

PROMPT_TEMPLATE = """
Answer the question based on the context provided. If no relevant information is found, state so.

Context:
{context}

Question:
{question}

Answer:
// Do not include the statement "Based on the provided context" in your answer. Start directly with the answer.
// If the user mentions that he does not want to know about these courses, ask him what topic he wants to learn about in the answer.

If you encounter a text message which is not related to the context, the bot must respond with "I don't have relevant information to answer that. Kindly ask queries related to free data science and machine learning courses on Analytics Vidhya."
"""

# Initialize the local embedding model
embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')


class LocalEmbeddingFunction:
    def embed_documents(self, texts):
        # Generate embeddings for a list of texts
        embeddings = embedding_model.encode(texts)
        return embeddings.tolist() if hasattr(embeddings, 'tolist') else embeddings

    def embed_query(self, query):
        # Generate an embedding for a single query string
        query_embedding = embedding_model.encode(query)
        return query_embedding.tolist() if hasattr(query_embedding, 'tolist') else query_embedding


class LLM:

    def generate_response(self, prompt):
        client = InferenceClient(
            provider="together",
            api_key = KEY
        )
        completion = client.chat.completions.create(
            model=repo_id,
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            max_tokens=500,
            temperature=0.8
        )
        return str(completion.choices[0].message.content)


def get_embedding_function():
    return LocalEmbeddingFunction()

def get_chat_response(query, history):
    # Retrieve embeddings from the Chroma DB
    embedding_function = get_embedding_function()
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)

    # Search the DB
    results = db.similarity_search_with_score(query, k=5)
    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])

    # Create prompt with context and query
    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, question=query)

    # Generate response using the LLM class and Hugging Face model
    model = LLM()
    response_text = model.generate_response(prompt)

    # Update the history
    history.append(AIMessage(content = response_text))

    return response_text


# Gradio Interface
def predict(message, history):
    # Initialize history if not provided
    history_langchain_format = []

    for msg in history:

        if msg['role'] == "user":
            history_langchain_format.append(HumanMessage(content=msg['content']))

        elif msg['role'] == "assistant":
            history_langchain_format.append(AIMessage(content=msg['content']))

    history_langchain_format.append(HumanMessage(content=message))

    # Get response from the model
    response = get_chat_response(message, history_langchain_format)

    return response
    
# Define the introductory content
intro_content = """
# Course Recommendation Bot
This bot helps you find **free courses related to python, data science and machine learning** available on [Analytics Vidhya](https://www.analyticsvidhya.com/).
You can ask any questions related to these courses.

For example:
- "What are the free courses on Python?"
- "Do you have any courses on data science?"
"""

# Build the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown(intro_content)  # Display introductory content
    chatbot = gr.ChatInterface(predict, type="messages")  # Chat interface

demo.launch(share=False, ssr_mode=False)