import streamlit as st
import logging
import os
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import CSVLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.llms import HuggingFaceHub
import dotenv
import yaml
import os
import zipfile

zip_file = "faiss_index.zip"

with zipfile.ZipFile(zip_file, 'r') as zip_ref:
    zip_ref.extractall(".")  # Extract to the current directory

print("Unzipping completed successfully.")
dotenv.load_dotenv()

def load_config():
    with open("yaml-editor-online.yaml", "r") as f:
        config = yaml.safe_load(f)
    return config

hf_token = os.getenv("HUGGING")
config = load_config()
logging.basicConfig(level=logging.INFO)

embeddings_model = HuggingFaceEmbeddings(model_name=config["embedding_model"])

def create_vector_db():
    try:
        loader = CSVLoader(file_path="disease.csv", source_column="Disease Information")
        data = loader.load()
        vectordb = FAISS.from_documents(documents=data, embedding=embeddings_model)
        vectordb.save_local(config["vector_db_path"])
        logging.info("Vector database successfully created and saved.")
    except Exception as e:
        logging.error("Error creating vector database:", exc_info=e)

def get_qa_chain(query):
    try:
        if not os.path.exists(config["vector_db_path"]):
            logging.error(f"FAISS index path does not exist: {config['vector_db_path']}")
            return "Error: No data found."

        vectordb = FAISS.load_local(
            config["vector_db_path"], embeddings_model, allow_dangerous_deserialization=True
        )
        retriever = vectordb.as_retriever(score_threshold=config["score_threshold"])
        relevant_docs = retriever.get_relevant_documents(query)[:3]

        if not relevant_docs:
            return "No relevant information found."

        summarized_context = " ".join(doc.page_content for doc in relevant_docs)
        prompt_template = """
        Given the following health-related context and a question, generate a structured answer:
        
        QUESTION: {query}
        
        Ensure the response is easy to understand and medically accurate.
        """
        prompt = PromptTemplate(input_variables=["query"], template=prompt_template).format(query=query)

        llm = HuggingFaceHub(
            repo_id=config["model_name"],
            model_kwargs={
                "temperature": config["temperature"],
                "max_length": 200,
                "top_p": config["top_p"],
                "top_k": config["top_k"]
            },
            huggingfacehub_api_token=hf_token
        )

        response = llm(prompt)
        return response.strip()
    except Exception as e:
        logging.error("Error getting response:", exc_info=e)
        return "Sorry, there was an error processing your request."

def main():
    st.set_page_config(page_title="Health Disease Chatbot", page_icon="🩺", layout="centered")
    
    st.markdown(
        """
        <style>
        .stApp {
            background-color: #f0f2f6;
            color: #333;
            font-family: 'Arial', sans-serif;
        }
        .title {
            color: #2E7D32;
            text-align: center;
        }
        .query-input {
            border-radius: 10px;
            padding: 10px;
        }
        .response-box {
            background-color: #ffffff;
            padding: 15px;
            border-radius: 8px;
            box-shadow: 2px 2px 10px rgba(0,0,0,0.1);
        }
        </style>
        """,
        unsafe_allow_html=True
    )
    
    st.markdown("<h1 class='title'>🩺 Health Disease Chatbot</h1>", unsafe_allow_html=True)
    st.write("Enter a question related to health conditions, symptoms, or treatments.")
    
    query = st.text_input("Your health-related question:", key="query", help="Ask about diseases, symptoms, or treatments.")
    
    if st.button("Get Information"):
        if query:
            response = get_qa_chain(query)
            st.markdown(f"<div class='response-box'><b>Response:</b><br>{response}</div>", unsafe_allow_html=True)
        else:
            st.warning("Please enter a query to get a response.")

if __name__ == "__main__":
    if not os.path.exists(config["vector_db_path"]):
        logging.info(f"Vector database not found at {config['vector_db_path']}, creating it now.")
        create_vector_db()
    main()