import initialize
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from langchain.chains import VectorDBQA
from langchain_community.llms import OpenAI
from langchain_core.prompts import PromptTemplate
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain.chains import LLMChain
from langchain_google_genai import GoogleGenerativeAI
from langchain_google_genai import ChatGoogleGenerativeAI

import gradio as gr
import requests
import os

from langchain_ollama import OllamaLLM


import sys
sys.path.append('../..')

# For Google Colab
'''
from google.colab import userdata
OPENAI_API_KEY = userdata.get('OPENAI_API_KEY')
hf_token = userdata.get('hf_token')
GEMINI_API_KEY = userdata.get('GEMINI_API_KEY')
# For Desktop
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # Read local .env file
OPENAI_API_KEY = os.environ['OPENAI_API_KEY']
hf_token = os.environ['hf_token']
GEMINI_API_KEY = os.environ['GEMINI_API_KEY']
'''

# For Hugging Face
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
hf_token = os.environ.get('hf_token')
GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY')
fs_token = os.environ.get('fs_token')

llm_name = "gpt-3.5-turbo-0125"


vectordb = initialize.initialize()

chat_history_doc = []
chat_history_IS = []

# For getting source documents
def get_file(source_documents):
  files = set()
  for doc in source_documents:
    file = os.path.basename(doc.metadata['source'])
    files.add(file)
    # Print unique filenames
  return list(set(files))


def chat_query_doc(question, chat_history_doc):
    
    query_old = f"""Provide an elaborate, precise and pointwise reply to the question: {question}.
                    Also, Please consider the provided chat history: {chat_history_doc}.
                    Ensure that your current response is detailed, accurate, and addresses each aspect of the question thoroughly.
                    If the context of the question doesn't align with your last reply, please provide your response in a fresh manner.
                    If don't get the answer, feel free to reply from your own knowledge."""
    
    
    # query = f"""You'll be asked with a User Query. If the Query is related to Electrical Domain, Provide a precise and point-wise reply to the query: {question} \
    #            based on provided context only. Ensure that your reply addresses each aspect of the query thoroughly. """

    query = f""" Provide a precise and point-wise reply to the query: {question} \
                based on provided context only. Ensure that your reply addresses each aspect of the query thoroughly, and highlight the important points using text formatting 
                in your reply."""



    #llm = ChatOpenAI(model = llm_name, temperature = 0.1, api_key = OPENAI_API_KEY)
    #llm = GoogleGenerativeAI(model = "gemini-pro", google_api_key = GEMINI_API_KEY)                                   ###
    #llm = ChatGoogleGenerativeAI(model = "gemini-1.0-pro", google_api_key = GEMINI_API_KEY, temperature = 0)
    llm = OllamaLLM(model="unsloth/Llama-3.2-3B")

    # Conversation Retrival Chain with Memory
    #memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    retriever = vectordb.as_retriever()
    qa = ConversationalRetrievalChain.from_llm(llm, retriever = retriever, return_source_documents = True)

    # Replace input() with question variable for Gradio
    result = qa({"question": query, "chat_history" : chat_history_doc})

    # Update the history with the latest question and response
    # history.append({"user": question, "bot": result["answer"]})
    # chat_history_doc.append((query, result["answer"]))

    source_docs = result["source_documents"]
    file_names = get_file(source_docs)
    #file_name = os.path.basename(source_docs[0].metadata['source'])
    file_name = ', '.join([f"{x}" for x in file_names[:3]])

    # print("History : ", history)
    # print("\n Chat_his : ", chat_history)

    return result["answer"] + "\n\nSources : " + file_name


def chat_query_IS(question, chat_history_IS):
    
    #llm = ChatOpenAI(model = llm_name, temperature = 0.1, api_key = OPENAI_API_KEY)
    #llm = ChatGoogleGenerativeAI(model="gemini-pro", google_api_key=GEMINI_API_KEY)                                   ###
    llm = OllamaLLM(model="unsloth/Llama-3.2-3B")

    system_old = f""" Provide an elaborate, detailed and pointwise reply about the Topic, as per relevant IS/IEEE/BIS Standard.
              Also, at the end of your reply, quote the Relevant Standard Referred. Topic : {question}
              """
    system = f""" Provide a reply poetically precise as william shakespeare for the Topic : {question}
              """

    
    result = llm.invoke(system_old)
    

    # Update the history with the latest question and response
    # history.append({"user": question, "bot": result.content})
    # chat_history_IS.append((system_old, result.content))

    
    return result.content    


iface_doc = gr.ChatInterface(
    fn=chat_query_doc,
    title="""Standard TS of POWERGRID""",
    concurrency_limit = None,
    examples = ["What should be the GIB height outside the GIS hall ?" ,
                 "STATCOM Station Ratings" , 
                 "Specifications of XLPE POWER Cables."], 
             #    "Specification for Ethernet Switches in SAS."] ,
    theme=gr.themes.Base(),
    fill_height = True,
    delete_cache = (300,360),
    css = "CSS/chat_style.css",

)

iface_IS = gr.ChatInterface(
    fn = chat_query_IS,
    title = """Indian / International Standards""",
    concurrency_limit = None,
    examples = ["Type Tests for HV Switchgears." ,
                "Measurement of acoustic noise level of Transformers & Reactors" , 
                "Technical Requirement for 765kV class Transformer", 
                "Specification of Distance Relays"] ,
    theme=gr.themes.Base(),
    fill_height = True,
    delete_cache = (300,360),
    css = "CSS/chat_style.css",
)

Title= "# Conversational BOT for Model-TS & Indian / International Standards"

Description = """
### Welcome to the Language Model (SS-Engg-Dept.)! 👋

This model is trained on **Model Technical Specifications** of the SS-Engg. Dept. and leverages the power of **ChatGPT** to answer your queries based on:

*   Relevant TS, GTR & Specific Requirements 📑
*   International/Indian Standards 🌎🇮🇳

**Tips for Effective Use:**

*   Use elaborate questions for more accurate responses. 🤔
*   Clear the chat if you don't receive a reply. 🔄
*   Include **Specific Keywords** in your query for precise results. 🎯 
"""              
 

with gr.Blocks(css="CSS/style.css", fill_height=True) as demo:
  
    # history = gr.State([])  # Initialize the state component
   
    with gr.Column():
    
        with gr.Row():
          with gr.Column(scale=1):
            gr.Image("Images/Chatbot.png", width = 110, show_download_button = False, show_label = False, show_share_button = False, elem_id = "Logo")
          with gr.Column(scale=3):
            gr.Markdown(Title)
          with gr.Column(scale=1):
            gr.Image("Images/PG Logo.png", width = 200, show_download_button = False, show_label = False, show_share_button = False, elem_id = "PG_Logo")  
    
        with gr.Row():
          gr.Markdown(Description)
    
        with gr.Row(equal_height=True):
          with gr.Column(elem_classes = ["chat_container"]):
            iface_doc.render()
          with gr.Column(elem_classes = ["chat_container"]):
            iface_IS.render()  
       

if __name__ == "__main__":
    demo.launch(debug=True, share=True)