import gradio as gr
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_groq import ChatGroq
import os

def initial_llm():
    llm = ChatGroq(
        temperature=0,
        groq_api_key="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",
        model_name="llama-3.3-70b-versatile"
    )
    return llm

def create_db():
    pdf_path = "The_GALE_ENCYCLOPEDIA_of_MEDICINE_SECOND.pdf"
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"🚨 PDF file not found at: {pdf_path}")
    
    loader = PyPDFLoader(pdf_path)
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    texts = text_splitter.split_documents(docs)
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    db_path = "./chroma_db"
    vector_db = Chroma.from_documents(texts, embeddings, persist_directory=db_path)
    print("✅ ChromaDB created and medical data saved!")
    return vector_db

def setup_qachain(vector_db, llm):
    retriever = vector_db.as_retriever()
    prompt_template = """You are DiagnoBot, an empathetic medical assistant focused on providing evidence-based health information. Your responses should:
1. Be clear, concise, and use simple language with medical terms in parentheses when needed
2. Always emphasize the importance of consulting healthcare professionals
3. Mark urgent symptoms with ⚠️
4. Provide practical lifestyle and preventive care recommendations
5. Include reliable sources when possible
6. Stay within your scope: no diagnoses, prescriptions, or definitive medical advice
7. Maintain HIPAA compliance and medical ethics
8. Show empathy while remaining professional
9. Keep the answers, diet plans India oriented, cause most of the users will be Indians.
10. Provide mental health related guidelines as well, citing proper references.
11. If users ask questions in bengali, respond them in bengali. If they ask questions in english, respond them in english.
Context from medical resources:
{context}
Question: {question}
Response (structured with clear headings and bullet points when appropriate):"""
    
    PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
    
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        chain_type_kwargs={"prompt": PROMPT}
    )
    return qa_chain

def chatbot_response(message, chat_history):
    if not message.strip():
        return "", chat_history
    
    try:
        response = qa_chain.invoke({"query": message})
        chat_history.append((message, response["result"]))
        return "", chat_history
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        chat_history.append((message, error_message))
        return "", chat_history

# Example questions
example_questions = [
    "Headache and fever for the past two days",
    "I have slept enough yet I am having a bad headache accompanied by sensitivity to light",
    "Chest pain and shortness of breath after minimal exertion",
    "Persistent fatigue and dizziness, especially when standing up quickly",
    "Abdominal pain in the lower right side and nausea that worsens after eating"
]

# Initialize database and chain
db_path = "./chroma_db"
if not os.path.exists(db_path):
    vector_db = create_db()
else:
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vector_db = Chroma(persist_directory=db_path, embedding_function=embeddings)

llm = initial_llm()
qa_chain = setup_qachain(vector_db, llm)

def launch_chatbot():
    with gr.Blocks() as demo:
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                gr.Image("logo.png", show_label=False, container=False)
        
        gr.Markdown("""
        DiagnoBot is a side project of **EarlyMed**—an initiative dedicated to empowering you with early health insights. 
        Leveraging AI for early detection, our mission is simple: *"Early Detection, Smarter Decision."*
        """)
        
        chatbot = gr.Chatbot(
            show_copy_button=True,
            height=400
        )
        
        with gr.Row():
            txt = gr.Textbox(
                placeholder="Describe your symptoms or ask a health question...",
                scale=4
            )
            send_btn = gr.Button("Send", scale=1)
        
        gr.Markdown("### Common Symptom Examples")
        with gr.Row():
            for question in example_questions:
                gr.Button(question).click(
                    lambda q: q,
                    inputs=[gr.Textbox(value=question, visible=False)],
                    outputs=[txt]
                ).then(
                    chatbot_response,
                    inputs=[txt, chatbot],
                    outputs=[txt, chatbot]
                )
        
        gr.Markdown("""
        **Important Disclaimer:** Our DiagnoBot provides general health information and preliminary insights based on described symptoms. 
        It should NOT be used for emergency situations or as a substitute for professional medical advice. 
        The information provided is not a diagnosis. Always consult a qualified healthcare provider for personal health concerns. 
        If you're experiencing severe symptoms, please seek immediate medical attention.
        """)
        
        txt.submit(chatbot_response, inputs=[txt, chatbot], outputs=[txt, chatbot])
        send_btn.click(chatbot_response, inputs=[txt, chatbot], outputs=[txt, chatbot])

    demo.launch(share=True)

if __name__ == "__main__":
    launch_chatbot()