import streamlit as st
import pinecone
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Pinecone, Chroma
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
import tiktoken
import random

# Fetch the OpenAI API key from Streamlit secrets
openai_api_key = st.secrets["openai_api_key"]

# Fetch Pinecone API key and environment from Streamlit secrets
pinecone_api_key = st.secrets["pinecone_api_key"]
pinecone_environment = st.secrets["pinecone_environment"]

# Initialize Pinecone
pinecone.init(api_key=pinecone_api_key, environment=pinecone_environment)

# Define the name of the Pinecone index
index_name = 'mi-resource-qa'

# Initialize the OpenAI embeddings object with the hardcoded API key
embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)

# Define functions
def insert_or_fetch_embeddings(index_name):
    if index_name in pinecone.list_indexes():
        vector_store = Pinecone.from_existing_index(index_name, embeddings)
        return vector_store
    else:
        raise ValueError(f"Index {index_name} does not exist. Please create it before fetching.")

# Initialize or fetch Pinecone vector store
vector_store = insert_or_fetch_embeddings(index_name)

# calculate embedding cost using tiktoken
def calculate_embedding_cost(text):
    import tiktoken
    enc = tiktoken.encoding_for_model('text-embedding-ada-002')
    total_tokens = len(enc.encode(text))
    # print(f'Total Tokens: {total_tokens}')
    # print(f'Embedding Cost in USD: {total_tokens / 1000 * 0.0004:.6f}')
    return total_tokens, total_tokens / 1000 * 0.0004


def ask_with_memory(vector_store, query, chat_history=[]):
    from langchain.chains import ConversationalRetrievalChain
    from langchain.chat_models import ChatOpenAI
    
    llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=1, openai_api_key=openai_api_key)
    
    # The retriever is created with metadata filter directly in search_kwargs
    # retriever = vector_store.as_retriever(search_type='similarity', search_kwargs={'k': 3, 'filter': {'source': {'$eq': 'https://mimtsstac.org/sites/default/files/session-documents/Intensifying%20Literacy%20Instruction%20-%20Essential%20Practices%20%28NATIONAL%29.pdf'}}})
    retriever = vector_store.as_retriever(search_type='similarity', search_kwargs={'k': 3, 'filter': {'source':'https://mimtsstac.org/sites/default/files/session-documents/Intensifying%20Literacy%20Instruction%20-%20Essential%20Practices%20%28NATIONAL%29.pdf'}})
    
    chain= ConversationalRetrievalChain.from_llm(llm, retriever)
    result = chain({'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)
    
q = st.text_input(label='Ask a question or make a request ', value='', placeholder=st.session_state.placeholder)
# q = st.text_input(label='Ask a question or make a request ', value='')

k = 3  # Set k to 3

# # Initialize chat history if not present
# if 'history' not in st.session_state:
#     st.session_state.history = []
    
if q:
    with st.spinner('Thinking...'):
        answer = ask_with_memory(vector_store, q, st.session_state.history)
    
    # Display the response in a text area
    st.text_area('Response: ', value=answer, height=400, key="response_text_area")
    
    st.success('Powered by MTSS GPT. AI can make mistakes. Consider checking important information.')

    # # Prepare chat history text for display
    # history_text = "\n\n".join(f"Q: {entry[0]}\nA: {entry[1]}" for entry in st.session_state.history)
    # Prepare chat history text for display in reverse order
    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)