import streamlit as st
from langchain_community.llms import LlamaCpp
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_core.callbacks import StreamingStdOutCallbackHandler
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain import PromptTemplate
from langchain.retrievers import TFIDFRetriever

def get_model(m_name):
    callbacks = [StreamingStdOutCallbackHandler()]
    print("creating llm started")
    llm = LlamaCpp(
        model_path=m_name,
        temperature=0.75,
        max_tokens=30,
        top_p=4,
        callback_manager=callbacks,
        verbose=True,  # Verbose is required to pass to the callback manager
    )
    print("creating llm ended")


    # retriever = TFIDFRetriever.from_texts(
    #     ["Finatial AI"])
    
    
    # template = """You are the Finiantial expert:
    # {history}
    # {context}
    # ### Instruction:
    # {question}
    
    # ### Input:
    
    
    # ### Response:
    # """
    
    # prompt1 = PromptTemplate(
    #     input_variables=["history", "context", "question"],
    #     template=template,
    # )
    
    # qa = RetrievalQA.from_chain_type(
    #     llm=llm,
    #     chain_type='stuff',
    #     retriever=retriever,
    #     verbose=False,
    #     chain_type_kwargs={
    #         "verbose": False,
    #         "prompt": prompt1,
    #         "memory": ConversationBufferMemory(
    #             memory_key="history",
    #             input_key="question"),
    #     }
    # )
    template = """You are the Finiantial expert:
    ### Instruction:
    {question}
    ### Input:
    ### Response:
    """
    
    prompt = PromptTemplate(template=template, input_variables=["question"])
    
    llm_chain_model = LLMChain(prompt=prompt, llm=llm)
    print("creating model created")
    return llm_chain_model


def main():

    """Build a streamlit layout"""
    # Wide mode
    st.set_page_config(layout="wide")
    
    llm_models = {
    
        "Base": "unsloth.Q5_K_M.gguf",
        "Cerebras": "cerebras_Llama3-DocChat-1.0-8B_Base_adapt_basic_model_16bit.gguf",
        "Bavest": "bavest_fin_llama_33b_adapt_basic_model_16bit.gguf",
        "Aliyasir": "aliyasir_Llama-3-8B-Instruct-Finance-RAG_adapt_basic_model_16bit.gguf",
        "Basic Adapt": "adapt-unsloth.Q5_K_M.gguf",
        "adapt llm": "AdaptLLM_finance-LLM-13B_adapt_basic_model_16bit.gguf",
        "Fibro" : "finbro-v0.1.0-llama-3-8B-instruct-1m.gguf",
    }
    # Designing the interface
    st.title("Financial LLM test")
    # For newline
    st.write("\n")
    # Instructions
    st.markdown("*Hint: you can select the LLM model and write your prompt")
    # Set the columns

    col1, col2 = st.columns(2)
    col1.subheader("Prompt Section")
    col2.subheader("Model Output")

    llm_qa = get_model(llm_models.get('Base'))
    # Model selection
    st.sidebar.title("Model selection")
    det_arch = st.sidebar.selectbox("LLM model", list(llm_models.keys()))

    # For newline
    st.sidebar.write("\n")


    if st.sidebar.button("Select LLM"):
        
        with st.spinner("Loading model..."):
            llm_qa = get_model(llm_models.get(det_arch))
            # load the model TODO

    text_input = ''
    with col1:
        text_input_temp = st.text_input(
            "Please, type your question and submit.",
            "Write Your Prompt",
            key="placeholder",
        )
        
        if st.button("Submit"):
            text_input = text_input_temp
            
    with col2:
        if text_input != '':
            with st.spinner("Analyzing..."):
                out_gen = llm_qa.run(question)  
                st.write("LLM Response: ", out_gen)
                text_input = ''
            
    

if __name__ == "__main__":
    main()





# def greet(question):
#     print(f"question is {question}")

#     out_gen = qa.run(question)  
#     print(f"out is {out_gen}")
#     return out_gen

# demo = gr.Interface(fn=greet, inputs="text", outputs="text")
# demo.launch(debug=True, share=True)