import streamlit as st
import google.generativeai as genai
from dotenv import load_dotenv
import os
import PIL
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate

load_dotenv()
os.getenv("langchain_google_genai")
os.environ['GOOGLE_API_KEY'] = 'AIzaSyA5cVv6I1HxH68CTiPGalPQHymtunvDxVY'
genai.configure(api_key="AIzaSyA5cVv6I1HxH68CTiPGalPQHymtunvDxVY")
# Function to extract text from PDF files

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

def get_pdf_text(pdf_docs):
    text = ""
    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

# Function to split text into chunks
def get_text_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
    chunks = text_splitter.split_text(text)
    return chunks

# Function to create a vector store from text chunks
def get_vector_store(text_chunks):
    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
    vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
    vector_store.save_local("faiss_index")

# Function to get the conversational chain
if "text_chunks" not in st.session_state:
    st.session_state.text_chunks = None

if "vector_store" not in st.session_state:
    st.session_state.vector_store = None

if "document_messages" not in st.session_state:
    st.session_state.document_messages = []


def get_conversational_chain():
    prompt_template = """
    Answer the question as detailed as possible, but only if it relates to lung diseases or conditions. If the question is unrelated to lung diseases, respond with "This question is not related to lung diseases, so I cannot provide an answer." If the answer is not in the provided context, just say, "answer is not available in the context", and do not provide a wrong answer.\n\n
    Context:\n {context}?\n
    Question: \n{question}\n

    Answer:
    """
    model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.8)
    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
    return chain

# Function to handle user input
# Function to handle user input
def user_input(user_question):
    model = genai.GenerativeModel('gemini-1.5-pro')
    chat = model.start_chat(history=[])
    response = chat.send_message(user_question)
    return response.text

# Streamlit UI setup
st.markdown("<h1 style='text-align: center;'>Chào mừng tới Medical Question Answering 🎈</h1>", unsafe_allow_html=True)


with st.expander("Instructions"):
    st.markdown("Truyền vào một câu hỏi liên quan đến y tế, chúng tôi sẽ giải đáp cho bạn.")
    st.markdown("Bạn có thể hỏi các câu liên quan đến triệu chứng, nguyên nhân và một số phương pháp điều trị.")





with st.sidebar:
    mode = st.selectbox("Chọn chức năng", ["Question with Images", "Question with Documents"])
    if mode == "Question with Images":
        uploaded_files = st.file_uploader("Choose medical images...", type=["jpg", "jpeg", "png", "dicom"], accept_multiple_files=True)
    elif mode == "Question with Documents":
        folder_path = "medicalDocuments"
        if st.session_state.text_chunks is None:
            pdf_docs = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".pdf")]
            raw_text = get_pdf_text(pdf_docs)
            st.session_state.text_chunks = get_text_chunks(raw_text)
            st.session_state.vector_store = get_vector_store(st.session_state.text_chunks)

# Initialize session state
if "messages" not in st.session_state:
    st.session_state.messages = []

if "image_messages" not in st.session_state:
    st.session_state.image_messages = []

if "max_messages" not in st.session_state:
    st.session_state.max_messages = 1000

# Handle "Question with Images" mode
col_1, col_2, col_3 = st.columns([8, 1, 8])
if mode == "Question with Images" and uploaded_files:
    with col_1:
        image = PIL.Image.open(uploaded_files[0])
        st.image(image, caption="Uploaded Image", use_column_width=True)
    with col_3:
        # Display past messages for Question with Images
        for message in st.session_state.image_messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        if prompt := st.chat_input("Ask a question about the image..."):
            st.session_state.image_messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)
            model = genai.GenerativeModel('gemini-1.5-flash')
            with st.chat_message("assistant"):
                try:
                    response = model.generate_content([prompt, image])
                    st.session_state.image_messages.append({"role": "assistant", "content": response.text})
                    st.markdown(response.text)
                except Exception as e:
                    st.session_state.max_messages = len(st.session_state.image_messages)
                    st.session_state.image_messages.append(
                        {"role": "assistant", "content": f"Oops! There was an error: {str(e)}"}
                    )
                    st.rerun()

if "document_messages" not in st.session_state:
    st.session_state.document_messages = []

# Handle "Question with Documents" mode
if mode == "Question with Documents":
    # Display past messages for Document-based conversation
    for message in st.session_state.document_messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if user_question := st.chat_input("Hỏi câu hỏi từ file PDF"):
        st.session_state.document_messages.append({"role": "user", "content": user_question})
        with st.chat_message("user"):
            st.markdown(user_question)

        # Generate the response
        with st.chat_message("assistant"):
            try:
                response = user_input(user_question)
                st.session_state.document_messages.append({"role": "assistant", "content": response})
                st.markdown(response)
            except Exception as e:
                st.session_state.document_messages.append(
                    {"role": "assistant", "content": f"Oops! There was an error: {str(e)}"}
                )
                st.rerun()

# Display past messages for non-image-based conversation
if mode != "Question with Images" and mode != "Question with Documents":
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if len(st.session_state.messages) < st.session_state.max_messages:
        if prompt := st.chat_input("Hôm nay bạn như thế nào?"):
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)
            model = genai.GenerativeModel(model_name="gemini-pro")
            with st.chat_message("assistant"):
                try:
                    prompt_parts = [prompt]
                    response = model.generate_content(prompt_parts)
                    st.session_state.messages.append({"role": "assistant", "content": response.text})
                    st.markdown(response.text)
                except Exception as e:
                    st.session_state.max_messages = len(st.session_state.messages)
                    st.session_state.messages.append(
                        {"role": "assistant", "content": f"Oops! There was an error: {str(e)}"}
                    )
                    st.rerun()