import os
import chainlit as cl
from groq import Groq
import fitz  # PyMuPDF
import logging

# ✅ Logging for debug (appears in HF logs)
logging.basicConfig(level=logging.INFO)

# ✅ Initialize Groq client
groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# ✅ Global variable to store PDF content
pdf_context = ""

@cl.on_chat_start
async def on_chat_start():
    logging.info("Chat started.")
    await cl.Message(content="👋 Welcome! Upload a PDF or ask a question.").send()

@cl.on_message
async def on_message(message: cl.Message):
    global pdf_context

    logging.info(f"User message received: {message.content}")

    # ✅ Handle PDF upload
    if message.elements:
        for element in message.elements:
            if element.name.endswith(".pdf"):
                try:
                    doc = fitz.open(element.path)
                    text = ""
                    for page in doc:
                        text += page.get_text()
                    pdf_context = text[:4000]
                    await cl.Message(content="✅ PDF uploaded! Now ask a question.").send()
                    return
                except Exception as e:
                    logging.error(f"Error reading PDF: {e}")
                    await cl.Message(content=f"❌ Error reading PDF: {e}").send()
                    return

    # ✅ Validate user message
    prompt = message.content.strip()
    if not prompt:
        await cl.Message(content="⚠️ Please enter a question.").send()
        return

    # ✅ Append PDF context if available
    full_prompt = prompt
    if pdf_context:
        full_prompt = f"Use the following PDF content to answer the question:\n\n{pdf_context}\n\nQuestion: {prompt}"

    try:
        # ✅ Wrap sync Groq call in async function
        def generate_response():
            return groq_client.chat.completions.create(
                model="llama3-8b-8192",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": full_prompt}
                ]
            )

        response = await cl.run_sync(generate_response)

        answer = response.choices[0].message.content.strip()

        logging.info(f"LLM response: {answer[:100]}...")

        # ✅ Display response
        await cl.Message(content=answer).send()

    except Exception as e:
        logging.error(f"❌ Error during LLM response: {e}")
        await cl.Message(content=f"❌ Error: {e}").send()