import gradio as gr
import requests
from rich.pretty import pprint
from langfuse import Langfuse
from langfuse import observe
import os
from dotenv import load_dotenv
load_dotenv()

@observe()
def respond(message, history):
    url = "http://ec2-52-202-99-138.compute-1.amazonaws.com/api/v1/agents/6534b3ac82fa11f0b0019acfa5a3a049/completions"
    headers = {
        "Authorization": os.getenv("RAGFLOW_API_KEY"),
        "Content-Type": "application/json"
    }
    payload = {
    "question": message,
    "stream": False,
    }
    response = requests.post(url, headers=headers, json=payload)
    if response.status_code == 200:
        agent_response = response.json()
        # pprint(agent_response)
        citations = []
        if agent_response["data"]["data"]["reference"]["chunks"]:
            for id, chunk in agent_response["data"]["data"]["reference"]["chunks"].items():
                citations.append({"cite_id":id, 
                                  "content": chunk['content'],
                                  "id": chunk["id"],
                                  "document_name": chunk["document_name"],
                                  "similarity": chunk["similarity"]})

          
        history.append({"role": "assistant", "content": agent_response["data"]["data"]["content"]})
        
        if citations:
            for cite in citations:
                history.append({
                    "role": "assistant",
                    "content": (f"{cite['content']}"),
                    "metadata": {"title": f"📚 Citation {cite['cite_id']}",
                                #  "id": cite["id"],
                                #  "parent_id": cite["document_name"],
                                 "log": f" (Match: {100*cite['similarity']:.2f}%)"
                                 }
                })
        return history
    else:
        return f"Error: {response.status_code} - {response.text}"

chatbot = gr.ChatInterface(
    respond,
    type="messages",
)

with gr.Blocks() as demo:
    # with gr.Sidebar():
    #     gr.LoginButton()
    gr.Markdown("# UNICEF EYSN Agent")
    chatbot.render()


if __name__ == "__main__":

    langfuse = Langfuse(
    secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
    public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
    host="https://cloud.langfuse.com"
    )

    demo.launch(debug=True)