import gradio as gr
from huggingface_hub import InferenceClient
from typing import List, Tuple
import fitz  # PyMuPDF

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Placeholder for the app's state
class MyApp:
    def __init__(self) -> None:
        self.documents = []
        self.load_pdf("THEDIA1.pdf")

    def load_pdf(self, file_path: str) -> None:
        """Extracts text from a PDF file and stores it in the app's documents."""
        doc = fitz.open(file_path)
        self.documents = []
        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text()
            self.documents.append({"page": page_num + 1, "content": text})
        print("PDF processed successfully!")

    def search_documents(self, query: str, k: int = 3) -> List[str]:
        """Searches for relevant documents containing the query string."""
        results = [doc["content"] for doc in self.documents if query.lower() in doc["content"].lower()]
        return results[:k] if results else ["No relevant documents found."]

app = MyApp()

def respond(
    message: str,
    history: List[Tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    system_message = "You are a concise knowledgeable DBT coach. Use relevant documents to guide users through DBT exercises and provide helpful information. You must response in less word with empathy. If needed ask one follow up question to continue next chat. Like- Would you like to know more about it?"
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    # RAG - Retrieve relevant documents
    retrieved_docs = app.search_documents(message)
    context = "\n".join(retrieved_docs)
    messages.append({"role": "system", "content": "Relevant documents: " + context})

    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens= 512,
        stream=True,
        temperature= 0.90,
        top_p= 0.8,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

demo = gr.Blocks()

with demo:
    gr.Markdown("🧘‍♀️ **Dialectical Behaviour Therapy**")
    gr.Markdown(
        "Disclaimer: This chatbot is based on a DBT exercise book that is publicly available. "
        "We are not medical practitioners, and the use of this chatbot is at your own responsibility."
    )
    
    chatbot = gr.ChatInterface(
        respond,
        examples=[
            ["I feel overwhelmed with work. Help me to feel relaxed"],
            ["Can you guide me through a quick meditation?"],
            ["How do I stop worrying about things I can't control?"],
            ["What are some DBT skills for managing anxiety?"],
            ["Can you guide to practice a mindfulness excercise?"],
            ["What is radical acceptance?"]
        ],
        title='DBT Coach 👩‍💼'
    )
if __name__ == "__main__":
    demo.launch()