import streamlit as st

from langchain import HuggingFaceHub

from langchain.chains.question_answering import load_qa_chain

from langchain.document_loaders import UnstructuredURLLoader

import os

with st.sidebar:
    st.title('🌎 Summarize your webpage')
    st.markdown('''
    ## About
    This app is using:
    - [Streamlit](https://streamlit.io/)
    - [LangChain](https://python.langchain.com/)
    - [Flan Alpaca Large](https://huggingface.co/declare-lab/flan-alpaca-large) LLM model

    ## How it works
    - Load up a web URL
    - Send the request to the LLM using the *load_qa_chain* in langchain
    - Get the answer and from Flan Alpaca Large LLM (open source model on HuggingFace)
 
    ''')
    st.write('Made with 🤖 by [Cazimir Roman](https://cazimir.dev)')

def load_app():
    llm = HuggingFaceHub(repo_id="declare-lab/flan-alpaca-large", model_kwargs={"temperature":0, "max_length":512})

    col1, col2 = st.columns([0.8, 0.2])

    url = col1.text_input('Enter a webpage url here to summarize')
    col2.write("")
    col2.write("")
    summarize = col2.button("Summarize")

    if url:
        loader = UnstructuredURLLoader(urls=[url])
        data = loader.load()
        
        if summarize:
            with st.spinner("Summarizing..."):
                chain = load_qa_chain(llm=llm, chain_type="stuff")
                response = chain.run(input_documents=data, question="Summarize this article in one paragraph")
                st.success(response)

def main():

    st.header("Summarize your webpage")

    col1, col2 = st.columns([0.8, 0.2])
    
    container = col1.container()

    with container:
        hugging_face_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")  
        api_key = container.text_input("Enter your HuggingFace API token", type="password", value="" if hugging_face_token == None else hugging_face_token)
        st.markdown('''You can find your token [here](https://huggingface.co/settings/tokens)''')

        col2.write("")
        col2.write("")
        submit = col2.button("Submit")
        
        if hugging_face_token:
            load_app()

        # submit button is pressed
        if submit:
            # check if api key length correct
                if len(api_key) == 37:
                    os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
                    load_app()
                else:
                    st.error("Api key is not correct")

if __name__ == '__main__':
    main()