import streamlit as st 
import llama_index 
import logging
import sys
import openai
import os 
import wikipedia
from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.indices.query.query_transform import HyDEQueryTransform
from llama_index.query_engine.transform_query_engine import TransformQueryEngine
# from llama_index.indices.vector_store import ChatGPTRetrievalPluginIndex
from llama_index.readers import ChatGPTRetrievalPluginReader

def get_wikipedia_document(topic):
    wiki_wiki = wikipedia.Wikipedia('en')

    page = wiki_wiki.page(topic)
    if page.exists():
        return page.text
    else:
        return None
    
def write_string_to_file(text, filename):
    with open(filename, 'w') as file:
        file.write(text)

def remove_scrollbar():
    # Apply custom CSS to remove the scrollbar
    st.markdown(
        """
        <style>
        .css-1l02zno {
            overflow: hidden !important;
        }
        </style>
        """,
        unsafe_allow_html=True
    )


def ui(): 
    st.title('Hyde for Wikipedia Pages')
    st.markdown("This is question Answering for Wikipedia made using LLama Index and Hyde")
    api_key = st.text_input('Enter your OpenAI key here: ')
    if api_key is not None and api_key != '':
        os.environ["OPENAI_API_KEY"] = api_key
        openai.api_key = os.environ["OPENAI_API_KEY"]
        logging.basicConfig(stream=sys.stdout, level=logging.INFO)
        logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

        topic_name= st.text_input('Enter your topic for Wikipedia: ')
        if topic_name is not None and topic_name!='':
            page_object = wikipedia.page(topic_name)
            content = page_object.content
            filename = f"./Data/{topic_name}.txt"
            write_string_to_file(content, filename)

            documents = SimpleDirectoryReader(f'./Data/{topic_name}.txt"').load_data()
            index = VectorStoreIndex.from_documents(documents)


            query_str= st.text_input('Enter your query for the document: ')
            if query_str is not None and query_str!='':
                query_engine = index.as_query_engine()
                response = query_engine.query(query_str)
                hyde = HyDEQueryTransform(include_original=True)
                hyde_query_engine = TransformQueryEngine(query_engine, hyde)

                response = hyde_query_engine.query(query_str)
                query_bundle = hyde(query_str)
                hyde_doc = query_bundle.embedding_strs[0]     
                remove_scrollbar()
                st.text(hyde_doc)
            


if __name__=="__main__":
    ui()