import gradio as gr
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)
from langchain.prompts import PromptTemplate
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from pydantic import BaseModel, Field
from enum import Enum
from langchain_openai import ChatOpenAI
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
import json
import os
from dotenv import load_dotenv
load_dotenv()

class IsAnswerable(Enum):
    YES = "YES - the given 'question' can be confidently answered using the given 'context'"
    NO = "NO - the given 'question' cannot be answered with the given 'context'"  

class AnswerStatus(BaseModel):
    status: IsAnswerable = Field(description="")
    answer: str = Field(description="answer the student's 'question' based solely on the given 'context'. The answer should be in HTML format without any font color, with equations represented in HTML <code> tags, irrespective of the syntax - <code> or latex-style syntax used for equations in the provided HTML.")

class FAQBot():
    def __init__(self):

        self.model = ChatOpenAI(
            model_name='gpt-3.5-turbo',
            openai_api_key=os.getenv("OPENAI_API_KEY"),
            openai_organization=os.getenv("OPENAI_ORGANIZATION"),
        )

        embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
        self.db = Chroma(persist_directory="./chroma/db", embedding_function=embedding_function, collection_name="course")
        self.db_faq = Chroma(persist_directory="./chroma/db_faq", embedding_function=embedding_function, collection_name="faq")
        self.qna_dict = json.load(open('./chroma/qna_dict'))
        self.course_db = json.load(open('./chroma/course_db'))

        self.parser = PydanticOutputParser(pydantic_object=AnswerStatus)
        self.fix_parser = OutputFixingParser.from_llm(parser=self.parser, llm=self.model, max_retries=3)
        self.prompt = PromptTemplate(
            template =  '''
                You're a helpful teaching assistant for a technical course on {course}. You will only answer student's 'question' based on the given 'context' of the course.\n
                The 'context' is a combination of two things - 1) Previous question and answers on the {course} that are similar to the student's question, and 2) some snippets of text from the course contents that are relevant to the student's 'question'.
                
                {format_instructions}\n
                
                ***
                'query' : {question}
                ***
                
                $$$
                'context' : {context}
                $$$
                I am reminding you again, you are a teaching assistant, do not add any facts into the answer that is not given in the 'context'. 
                The answer should be in HTML format without any font color, with equations represented in HTML <code> tags, irrespective of the syntax - <code> or latex-style syntax used for equations in the provided HTML.
            ''',
                input_variables=["question", "context", "course"],
                partial_variables={
                    "format_instructions": self.parser.get_format_instructions(),
                },
            )    
        self.search_conf_thresh = 1
        self.excuse_me_msg = '''<p>I dont think I know the answer for this, let me check with the professor.</p>'''
    
    def ask_question(self, question, verbose=False):
        
        retrieved_answers = ''
        ## search in faq
        if verbose:
            print('Search in FAQ')
        results_faq = self.db_faq.similarity_search_with_score(question, k=3)

        ### save only the high confidence search results
        if verbose:
            print('\tanswers retrieved')

        is_faq_title_printed = False
        for i, val in enumerate(results_faq):
            if verbose:
                print('\t\ttext: {}\n\t\tChapter: {}\n\t\tconf:{}\n'.format(val[0].page_content, val[0].metadata, val[1]))

            if val[1] < self.search_conf_thresh: 
                if not is_faq_title_printed:
                    retrieved_answers += '''Question and Answers from the past that are similar to the student's question\n-----------------\n'''
                    is_faq_title_printed = True
                # collect the corresponding answers of the qna pair for gpt 
                retrieved_answers += ' Question:{}\n Answer:{}\n'.format(val[0].page_content, self.qna_dict[val[0].page_content])    



        ## search in coursework
        if verbose:
            print('Search in coursework')            
        results = self.db.similarity_search_with_score(question, k=5)

        ### save only the high confidence search results
        if verbose:
            print('\tanswers retrieved')

        is_snippet_title_printed = False
        max_chapters = 3
        neighboring_sections = 2 # + or -
        chapter_cnt = 0
        seen_chapters = []
        for i, val in enumerate(results):
            if verbose:
                print('\t\ttext: {}\n\t\tChapter: {}\n\t\tSection: {}\n\t\tconf:{}\n'.format(val[0].page_content, val[0].metadata['source'], val[0].metadata['split'], val[1]))
                print(self.course_db[val[0].metadata['source']].keys())
            if val[1] < self.search_conf_thresh: 
                if not is_snippet_title_printed:
                    retrieved_answers += '''\n$$$$$$$$$$\nSnippets of text from the course that are relevant to the student's question\n-----------------\n'''
                    is_snippet_title_printed = True

                if val[0].metadata['source'] not in seen_chapters and chapter_cnt<max_chapters:
                    
                    html_str = self.course_db[val[0].metadata['source']][str((val[0].metadata['split']))]
                    extended_context = ''
                    for ind in range(val[0].metadata['split']-neighboring_sections, val[0].metadata['split']+neighboring_sections):
                        if str(ind) in self.course_db[val[0].metadata['source']]:
                            extended_context += '\n{}'.format(self.course_db[val[0].metadata['source']][str(ind)])
                    
                    retrieved_answers += '\n Relevant text snippet {}: {}\n\n '.format(chapter_cnt, extended_context) 
                    if verbose:
                        print('\n\t\tlength:({}, {})'.format(len(html_str), len(extended_context)))

                    seen_chapters.append(val[0].metadata['source'])
                    chapter_cnt += 1
                    if len(retrieved_answers)>2000:
                        if verbose:
                            print('retrieved_answers length greater than 2000 : {}'.format(len(retrieved_answers)))
                        break
                    



        
        #### if there is atleast one search result ask GPT to answer
        if len(retrieved_answers):
            # ask GPT to answer
            prompt_string = self.prompt.format_prompt(question=question, context=retrieved_answers, course = 'Distributed Algorithms').to_string()

            if verbose:
                print(prompt_string)
            response = self.model([
                HumanMessage(
                    prompt_string
                ) 
            ])

            if verbose:
                print('\t\t\tRaw GPT response: {}\n'.format(response))
                
            faq_response = None
            try:
                faq_response = self.parser.parse(response.content)
            except Exception as e:
                faq_response = self.fix_parser.parse(response.content)                

            if verbose:
                print('\t\t\tfinal response: {}\n'.format(faq_response))

            if faq_response != None and faq_response.status == IsAnswerable.YES:
                return faq_response.answer
            else:
                return self.excuse_me_msg
        else:
            return self.excuse_me_msg


fb = FAQBot()

demo = gr.ChatInterface(fb.ask_question)

if __name__ == "__main__":
    demo.launch()