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

groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
pdf_text = ""

@cl.on_chat_start
async def on_chat_start():
    await cl.Message(content="📄 براہ کرم PDF فائل اپلوڈ کریں یا سوال پوچھیں۔").send()

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

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

    user_prompt = message.content.strip()
    if not user_prompt:
        await cl.Message(content="⚠️ Please ask a valid question.").send()
        return

    prompt = f"PDF Content:\n{pdf_text}\n\nQuestion: {user_prompt}" if pdf_text else user_prompt

    try:
        response = groq_client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ]
        )
        answer = response.choices[0].message.content.strip()
        await cl.Message(content=answer).send()
    except Exception as e:
        await cl.Message(content=f"❌ Error from LLM: {e}").send()