import os
import torch
import streamlit as st
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.memory import ConversationBufferMemory
from langchain import PromptTemplate, LLMChain
from langchain.llms import HuggingFacePipeline
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from dotenv import load_dotenv
from htmlTemplates import css
import warnings

# Suppress GPTNeoXSdpaAttention deprecation warnings
warnings.filterwarnings(
    "ignore", 
    message="The `GPTNeoXSdpaAttention` class is deprecated", 
    category=UserWarning
)

# Load environment variables
load_dotenv()

# Dolly-v2-3b model pipeline
@st.cache_resource
def load_pipeline():
    model_name = "databricks/dolly-v2-3b"
    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

    # Update _attn_implementation
    model.config._attn_implementation = "triton"  # Or another supported implementation

    return pipeline(
        model=model,
        tokenizer=tokenizer,
        torch_dtype=torch.float32,  # Use float32 for CPU
        device_map="cpu",  # Force CPU usage
        return_full_text=True
    )

# Initialize Dolly pipeline
generate_text = load_pipeline()

# Create a HuggingFace pipeline wrapper for LangChain
hf_pipeline = HuggingFacePipeline(pipeline=generate_text)

# Template for instruction-only prompts
prompt = PromptTemplate(
    input_variables=["instruction"],
    template="{instruction}"
)

# Template for prompts with context
prompt_with_context = PromptTemplate(
    input_variables=["instruction", "context"],
    template="{instruction}\n\nInput:\n{context}"
)

# Create LLM chains
llm_chain = LLMChain(llm=hf_pipeline, prompt=prompt)
llm_context_chain = LLMChain(llm=hf_pipeline, prompt=prompt_with_context)

# Extracting text from .txt files
def get_text_files_content(folder):
    text = ""
    for filename in os.listdir(folder):
        if filename.endswith('.txt'):
            with open(os.path.join(folder, filename), 'r', encoding='utf-8') as file:
                text += file.read() + "\n"
    return text

# Converting text to chunks
def get_chunks(raw_text):
    text_splitter = CharacterTextSplitter(
        separator="\n",
        chunk_size=2000,
        chunk_overlap=500,
        length_function=len
    )
    chunks = text_splitter.split_text(raw_text)
    return chunks

# Using Hugging Face embeddings model and FAISS to create vectorstore
def get_vectorstore(chunks):
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        model_kwargs={'device': 'cpu'}  # Ensure embeddings use CPU
    )
    vectorstore = FAISS.from_texts(texts=chunks, embedding=embeddings)
    return vectorstore

# Generating response from user queries
def handle_question(question, vectorstore=None):
    if vectorstore:
        documents = vectorstore.similarity_search(question, k=3)
        context = "\n".join([doc.page_content for doc in documents])
        if context:
            result_with_context = llm_context_chain.run(instruction=question, context=context)
            return result_with_context
    return llm_chain.run(instruction=question)

def main():
    st.set_page_config(page_title="Chat with Notes and AI", page_icon=":books:", layout="wide")
    st.write(css, unsafe_allow_html=True)

    # Initialize session state
    if "vectorstore" not in st.session_state:
        st.session_state.vectorstore = None

    st.header("Chat with Notes and AI :books:")

    # Subject selection dropdown
    subjects = [
        "A Trumped World", "Agri Tax in Punjab", "Assad's Fall in Syria", "Elusive National Unity", "Europe and Trump 2.0",
        "Going Down with Democracy", "Indonesia's Pancasila Philosophy", "Pakistan in Choppy Waters",
        "Pakistan's Semiconductor Ambitions", "Preserving Pakistan's Cultural Heritage", "Tackling Informal Economy",
        "Technical Education in Pakistan", "The Case for Solidarity Levies", "The Decline of the Sole Superpower",
        "The Power of Big Oil", "Trump 2.0 and Pakistan's Emerging Foreign Policy", "Trump and the World 2.0",
        "Trump vs BRICS", "US-China Trade War", "War on Humanity", "Women's Suppression in Afghanistan"
    ]
    data_folder = "data"
    subject_folders = {subject: os.path.join(data_folder, subject.replace(' ', '_')) for subject in subjects}
    selected_subject = st.sidebar.selectbox("Select a Subject:", subjects)

    st.sidebar.info(f"You have selected: {selected_subject}")

    # Process data folder for question answering
    subject_folder_path = subject_folders[selected_subject]
    if os.path.exists(subject_folder_path):
        raw_text = get_text_files_content(subject_folder_path)
        if raw_text:
            text_chunks = get_chunks(raw_text)
            vectorstore = get_vectorstore(text_chunks)
            st.session_state.vectorstore = vectorstore
        else:
            st.error("No content found for the selected subject.")
    else:
        st.error(f"Folder not found for {selected_subject}.")

    # Chat interface
    question = st.text_input("Ask a question about your selected subject:")
    if question:
        if st.session_state.vectorstore:
            response = handle_question(question, st.session_state.vectorstore)
            st.subheader("Response:")
            st.write(response)
        else:
            st.warning("Please load the content for the selected subject before asking a question.")

if __name__ == '__main__':
    main()