from retrival import retrival
from indexing import indexing
import os
from dotenv import load_dotenv
import gradio as gr
from groq import Groq
import uuid
load_dotenv()

client=Groq(api_key=os.getenv("GROQ_API_KEY"))
current_file = None

def handle_file(file, session_id, progress=gr.Progress()):
    global current_file

    if file is not None:
        if current_file != file.name:
            status = indexing(file, session_id)
            current_file = file.name
            return status

    return "Already indexed."


def mainwork(query,history,session_id):
    
    if history is None:
        history=[]

    contextchunks=retrival(query,session_id)
    context="\n\n".join([
        f"[Source {i}: {doc.metadata.get('source','?')}, page {doc.metadata.get('page','?')}]\n{doc.page_content}"
        for i, doc in enumerate(contextchunks, 1)
    ])

    messages = [
        {
            "role": "system",
            "content": """
            You are DocBuddy, a helpful document assistant.

            Use the retrieved document context as the primary source of truth.

            Use the conversation history to understand follow-up questions,
            references, and pronouns such as:
            - it
            - this
            - that
            - these
            - these two
            - the previous concept

            When answering:
            1. Prefer information from the retrieved documents.
            2. Combine information from multiple retrieved chunks when needed.
            3. Explain concepts clearly and naturally.
            4. If the answer cannot be supported by the retrieved documents, say:
            'I don't have enough information in the uploaded documents to answer that.'
            5. Do not invent facts or citations.
        
            Always end with:

            Sources:
            - source name and page number
            """
        }]
    for msg in history:
        messages.append({
            "role":msg["role"],
            "content":msg["content"]
        })
    messages.append({"role": "user",   "content": f"Context:\n{context}\n\nQuestion: {query}"})

    response=client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        stream=True,
        temperature=0
    )
    
    history.append(
        {
            "role":"user",
            "content":query
        })
    history.append(
        {
            "role":"assistant",
            "content":""
        }
    )

    for chunk in response:
        content=chunk.choices[0].delta.content or ""
        history[-1]["content"]+=content
        yield history


with gr.Blocks() as demo:
    session_id=gr.State(value=lambda:str(uuid.uuid4()))
    with gr.Tab("💬 Chat"):
        with gr.Column(scale=3):
            chatbot=gr.Chatbot(height=440,label="Conversation")
            textbox=gr.Textbox(placeholder="Ask anything here.....",show_label=False)
            button=gr.Button("Send")
            textbox.submit(
                mainwork,
                inputs=[textbox,chatbot,session_id],
                outputs=[chatbot]
                ).then(lambda:(""),outputs=textbox)
            button.click(
                mainwork,
                inputs=[textbox,chatbot,session_id],
                outputs=[chatbot]
            ).then(lambda: (""),outputs=[textbox])
    
    with gr.Tab("Documents"):
        gr.Markdown("Give your documents here")
        file=gr.File()
        indx_status=gr.Textbox("Indexing status",interactive=False)
        file.change(
            handle_file,
            inputs=[file,session_id],
            outputs=[indx_status]
        ).then(lambda: None,outputs=[file])

demo.launch()