#------------------------------------------------------------------------
# Import Modules
#------------------------------------------------------------------------

import streamlit as st
import openai
import random
import os

from pinecone import Pinecone
from langchain.chat_models import ChatOpenAI
from langsmith import Client
from langchain.smith import RunEvalConfig, run_on_dataset

#------------------------------------------------------------------------
# Load API Keys From the .env File, & OpenAI, Pinecone, and LangSmith Client
#------------------------------------------------------------------------

# Fetch the OpenAI API key from Streamlit secrets
os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
# Retrieve the OpenAI API Key from environment variable
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# Initialize OpenAI Service
openai.api_key = OPENAI_API_KEY

# Fetch Pinecone API key from Streamlit secrets
os.environ["PINECONE_API_KEY"] = st.secrets["PINECONE_API_KEY"]
# Retrieve the Pinecone API Key from environment variable
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
# Initialize Pinecone Service
# from pinecone import Pinecone
pc = Pinecone(api_key=PINECONE_API_KEY)

# Fetch LangSmith API key from Streamlit secrets
# os.environ["LANGCHAIN_API_KEY"] = st.secrets["LANGCHAIN_API_KEY"]
os.environ["LANGCHAIN_API_KEY"] = "ls__1819fb2979e44f0a9e410688d81c6390"
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ["LANGCHAIN_PROJECT"] = "Inkqa"
# Retrieve the LangSmith API Key from environment variable
LANGCHAIN_API_KEY = os.getenv("LANGCHAIN_API_KEY")
# Initialize LangSmith Service
client = Client(api_key=LANGCHAIN_API_KEY) #langsmith client

#------------------------------------------------------------------------
# Initialize 
#------------------------------------------------------------------------

# # Define the name of the Pinecone index
index_name = 'mimtssinkqa'

# Initialize the OpenAI embeddings object 
from langchain_openai import OpenAIEmbeddings
# embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
embeddings = OpenAIEmbeddings()

# LOAD VECTOR STORE FROM EXISTING INDEX
from langchain_community.vectorstores import Pinecone
vector_store = Pinecone.from_existing_index(index_name='mimtssinkqa', embedding=embeddings)

def ask_with_memory(vector_store, query, chat_history=[]):
    from langchain_openai import ChatOpenAI
    from langchain.chains import ConversationalRetrievalChain
    from langchain.memory import ConversationBufferMemory

    from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate

    # llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0.5, openai_api_key=OPENAI_API_KEY)
    llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0.5)
    
    retriever = vector_store.as_retriever(search_type='similarity', search_kwargs={'k': 3})
    
    memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)

    system_template = r'''
    Article Title: 'Intensifying Literacy Instruction: Essential Practices.' 
    Article Focus: The main focus of the article is reading and the secondary focus is writing. 
    Expertise: Assume the role of an expert literacy coach with in-depth knowledge of the Simple View of Reading, School-Wide Positive Behavioral Interventions and Supports (SWPBIS), and Social Emotional Learning (SEL).
    Audience: Tailor your response for teachers and administrators seeking to enhance literacy instruction within their educational settings.
    Response Requirements: Provide an answer utilizing the context provided. Unless specifically requested by the user, avoid mentioning the article's header.
    Cover all necessary details relevant to the question posed, drawing on your expertise in literacy instruction and the Simple View of Reading.
    Utilize paragraphs for detailed and descriptive explanations, and bullet points for highlighting key points or steps, ensuring the information is easily understood.
    Conclude with a recapitulation of main points, summarizing the essential takeaways from your response.
    ----------------
    Context: ```{context}```
    '''

    user_template = '''
    Question: ```{question}```
    Chat History: ```{chat_history}```
    '''

    messages= [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template(user_template)
    ]

    qa_prompt = ChatPromptTemplate.from_messages (messages)

    chain = ConversationalRetrievalChain.from_llm(llm=llm, retriever=retriever, memory=memory,chain_type='stuff', combine_docs_chain_kwargs={'prompt': qa_prompt}, verbose=False
    )

    result = chain.invoke({'question': query, 'chat_history': st.session_state['history']})
    # Append to chat history as a dictionary
    st.session_state['history'].append((query, result['answer']))
    
    return (result['answer'])

# Initialize chat history
if 'history' not in st.session_state:
    st.session_state['history'] = []
    
# # STREAMLIT APPLICATION SETUP WITH PASSWORD

# Define the correct password
# correct_password = "MiBLSi"

#Add the image with a specified width
image_width = 300  # Set the desired width in pixels
st.image('MTSS.ai_Logo.png', width=image_width)
st.subheader('Ink QA™ | Dynamic PDFs')

# Using Markdown for formatted text
st.markdown("""
Resource: **Intensifying Literacy Instruction: Essential Practices**
""", unsafe_allow_html=True)

with st.sidebar:
    # Password input field
    # password = st.text_input("Enter Password:", type="password")
    
    st.image('mimtss.png', width=200)
    st.image('Literacy_Cover.png', width=200)
    st.link_button("View | Download", "https://mimtsstac.org/sites/default/files/session-documents/Intensifying%20Literacy%20Instruction%20-%20Essential%20Practices%20%28NATIONAL%29.pdf")
        
    Audio_Header_text = """
    **Tune into Dr. St. Martin's introduction**"""
    st.markdown(Audio_Header_text)
        
    # Path or URL to the audio file
    audio_file_path = 'Audio_Introduction_Literacy.m4a'
    # Display the audio player widget
    st.audio(audio_file_path, format='audio/mp4', start_time=0)
        
    # Citation text with Markdown formatting
    citation_Content_text = """
    **Citation**  
    St. Martin, K., Vaughn, S., Troia, G., Fien, & H., Coyne, M. (2023). *Intensifying literacy instruction: Essential practices, Version 2.0*. Lansing, MI: MiMTSS Technical Assistance Center, Michigan Department of Education.
        
    **Table of Contents**  
    * **Introduction**: pg. 1  
    * **Intensifying Literacy Instruction: Essential Practices**: pg. 4  
    * **Purpose**: pg. 4  
    * **Practice 1**: Knowledge and Use of a Learning Progression for Developing Skilled Readers and Writers: pg. 6
    * **Practice 2**: Design and Use of an Intervention Platform as the Foundation for Effective Intervention: pg. 13
    * **Practice 3**: On-going Data-Based Decision Making for Providing and Intensifying Interventions: pg. 16
    * **Practice 4**: Adaptations to Increase the Instructional Intensity of the Intervention: pg. 20
    * **Practice 5**: Infrastructures to Support Students with Significant and Persistent Literacy Needs: pg. 24
    * **Motivation and Engagement**: pg. 28
    * **Considerations for Understanding How Students' Learning and Behavior are Enhanced**: pg. 28
    * **Summary**: pg. 29
    * **Endnotes**: pg. 30
    * **Acknowledgment**: pg. 39
    """
    st.markdown(citation_Content_text)

# if password == correct_password:
# Define a list of possible placeholder texts
placeholders = [
    'Example: Summarize the article in 200 words or less',
    'Example: What are the essential practices?',
    'Example: I am a teacher, why is this resource important?',
    'Example: How can this resource support my instruction in reading and writing?',
    'Example: Does this resource align with the learning progression for developing skilled readers and writers?',
    'Example: How does this resource address the needs of students scoring below the 20th percentile?',
    'Example: Are there assessment tools included in this resource to monitor student progress?',
    'Example: Does this resource provide guidance on data collection and analysis for monitoring student outcomes?',
    "Example: How can this resource be used to support students' social-emotional development?",
    "Example: How does this resource align with the district's literacy goals and objectives?",
    'Example: What research and evidence support the effectiveness of this resource?',
    'Example: Does this resource provide guidance on implementation fidelity'
]

# Select a random placeholder from the list
if 'placeholder' not in st.session_state:
    st.session_state.placeholder = random.choice(placeholders)


# CLEAR THE TEXT BOX
with st.form("Question",clear_on_submit=True):
    q = st.text_input(label='Ask a Question | Send a Prompt', placeholder=st.session_state.placeholder, value='', )
    submitted = st.form_submit_button("Submit")
    
    st.divider()
    
    if submitted:
        with st.spinner('Thinking...'):
            answer = ask_with_memory(vector_store, q, st.session_state.history)
        
        # st.write(q)
        st.write(f"**{q}**")
        
        import time
        import random

        def stream_answer():
            for word in answer.split(" "):
                yield word + " "
                # time.sleep(0.02)
                time.sleep(random.uniform(0.03, 0.08))
        
        st.write(stream_answer)
        
        # Display the response in a text area
        # st.text_area('Response: ', value=answer, height=400, key="response_text_area")
        # OR to display as Markdown (interprets Markdown formatting)
        # st.markdown(answer)
            
        st.success('Powered by MTSS GPT. AI can make mistakes. Consider checking important information.')
        
        st.divider()

        # # Prepare chat history text for display
        history_text = "\n\n".join(f"Q: {entry[0]}\nA: {entry[1]}" for entry in reversed(st.session_state.history))

        # Display chat history
        st.text_area('Chat History', value=history_text, height=800)