import streamlit as st
import os
from dotenv import load_dotenv
import time
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_groq import ChatGroq
from langchain.chains import RetrievalQA
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import LLMChain

# Set persistent storage path
PERSISTENT_DIR = "vector_db"

def initialize_vector_db():
    # Check if vector database already exists in persistent storage
    if os.path.exists(PERSISTENT_DIR) and os.listdir(PERSISTENT_DIR):
        embeddings = HuggingFaceEmbeddings()
        vector_db = Chroma(persist_directory=PERSISTENT_DIR, embedding_function=embeddings)
        return None, vector_db
    
    base_dir = os.path.dirname(os.path.abspath(__file__))
    pdf_files = [f for f in os.listdir(base_dir) if f.endswith('.pdf')]
    loaders = [PyPDFLoader(os.path.join(base_dir, fn)) for fn in pdf_files]
    
    documents = []
    for loader in loaders:
        documents.extend(loader.load())
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
        separators=["\n\n", "\n", " ", ""]
    )
    texts = text_splitter.split_documents(documents)
    
    embeddings = HuggingFaceEmbeddings()
    vector_db = Chroma.from_documents(
        texts, 
        embeddings, 
        persist_directory=PERSISTENT_DIR
    )
    vector_db.persist()
    return documents, vector_db

# System instructions for the LLM
system_prompt = """You are an expert Agro-Homeopathy doctor. When providing remedies:
1. Always specify medicine potency as 6c unless the uploaded text mentions some other value explicitly
3. Provide comprehensive diagnosis and treatment plans
4. Base recommendations on homeopathic principles
"""

api_key1 = os.getenv("api_key")

start_time = time.time()
st.set_page_config(page_title="Dr. Radha: The Agro-Homeopath", page_icon="🚀", layout="wide")

st.markdown("""
<style>
    #the-title {
        text-align: center;
    }
</style>
""", unsafe_allow_html=True)

st.title("📚 Ask Dr. Radha - World's First AI based Agro-Homeopathy Doctor")

# Add information request message
st.markdown("""
Please provide complete details about your plant issue including:
- Detailed description of plant symptoms
- Current weather conditions
- Current Temperature
""")

human_image = "human.png"
robot_image = "bot.png"

# Set up Groq API with temperature 0.7
llm = ChatGroq(
    api_key=api_key1, 
    max_tokens=None, 
    timeout=None, 
    max_retries=2, 
    temperature=0.7,
    model="llama-3.1-70b-versatile"
)

embeddings = HuggingFaceEmbeddings()
end_time = time.time()
print(f"Setting up Groq LLM & Embeddings took {end_time - start_time:.4f} seconds")

# Initialize session state
if "documents" not in st.session_state:
    st.session_state["documents"] = None
if "vector_db" not in st.session_state:  
    st.session_state["vector_db"] = None
if "query" not in st.session_state:
    st.session_state["query"] = ""
    
start_time = time.time()
if st.session_state["documents"] is None or st.session_state["vector_db"] is None:
    with st.spinner("Loading data..."):
        documents, vector_db = initialize_vector_db()
        st.session_state["documents"] = documents  
        st.session_state["vector_db"] = vector_db
else:
    documents = st.session_state["documents"]
    vector_db = st.session_state["vector_db"]

end_time = time.time()
print(f"Loading and processing PDFs & vector database took {end_time - start_time:.4f} seconds")

start_time = time.time()
retriever = vector_db.as_retriever()

prompt_template = """You are an expert Agro-Homeopathy doctor. Analyze the following context and question to provide a clear, structured response.

Context: {context}

Question: {question}

Provide your response in the following format:
1. Diagnosis: Analyze the plant symptoms and conditions described
2. Treatment: Recommend specific homeopathic medicine(s) with exact potency
3. Guidelines: List 3 key recommendations for addressing the issue

Remember to maintain a professional, clear tone and ensure all medicine recommendations include specific potency.

Answer:"""

# Create the QA chain with correct variables
qa = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs={
        "prompt": PromptTemplate(
            template=prompt_template,
            input_variables=["context", "question"]
        )
    }
)

# Create a separate LLMChain for fallback
fallback_template = """As an expert Agro-Homeopathy doctor, provide a structured response to the following question:

Question: {question}

Format your response as follows:
1. Diagnosis: Analyze the described plant condition
2. Treatment: Recommend specific homeopathic medicine(s) with exact potency
3. Guidelines: Provide 3 key recommendations

Maintain a professional tone and ensure all medicine recommendations include specific potency.

Answer:"""

fallback_prompt = PromptTemplate(template=fallback_template, input_variables=["question"])
fallback_chain = LLMChain(llm=llm, prompt=fallback_prompt)

chat_container = st.container()

with st.form(key='query_form'):
    query = st.text_input("Ask your question:", value="")
    submit_button = st.form_submit_button(label='Submit')

end_time = time.time()
print(f"Setting up retrieval chain took {end_time - start_time:.4f} seconds")
start_time = time.time()

if submit_button and query:
    with st.spinner("Generating response..."):
        result = qa({"query": query})  
        if result['result'].strip() == "":
            # If no result from PDF, use fallback chain
            fallback_result = fallback_chain.run(query)
            response = fallback_result
        else:
            response = result['result']

    col1, col2 = st.columns([1, 10])
    with col1:
        st.image(human_image, width=80)
    with col2:
        st.markdown(f"{query}")
    col1, col2 = st.columns([1, 10])
    with col1:
        st.image(robot_image, width=80)
    with col2:
        st.markdown(f"{response}")
    
    st.markdown("---")
    
    st.session_state["query"] = ""

end_time = time.time()
print(f"Actual query took {end_time - start_time:.4f} seconds")