import tkinter as tk
from tkinter import ttk
import requests

API_URL = "https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct"
headers = {"Authorization": "Bearer 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"}

API_URL2 = "https://api-inference.huggingface.co/models/valhalla/longformer-base-4096-finetuned-squadv1"
headers2 = {"Authorization": "Bearer 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"}

class App:
    def __init__(self, root):
        self.root = root
        root.title("AI Query UI")

        self.question_label = ttk.Label(root, text="Question:")
        self.question_entry = ttk.Entry(root, width=50)

        self.context_label = ttk.Label(root, text="Context:")
        self.context_entry = ttk.Entry(root, width=50)

        self.get_context_button = ttk.Button(root, text="Get Context", command=self.get_context)
        self.ask_ai_button = ttk.Button(root, text="Ask AI", command=self.ask_ai)

        self.answer_label = ttk.Label(root, text="Answer:")
        self.answer_entry = ttk.Entry(root, width=50, state="readonly")

        # Grid layout
        self.question_label.grid(row=0, column=0, padx=10, pady=5, sticky="w")
        self.question_entry.grid(row=0, column=1, padx=10, pady=5, columnspan=2)

        self.context_label.grid(row=1, column=0, padx=10, pady=5, sticky="w")
        self.context_entry.grid(row=1, column=1, padx=10, pady=5, columnspan=2)

        self.get_context_button.grid(row=2, column=0, padx=10, pady=5)
        self.ask_ai_button.grid(row=2, column=1, padx=10, pady=5)

        self.answer_label.grid(row=3, column=0, padx=10, pady=5, sticky="w")
        self.answer_entry.grid(row=3, column=1, padx=10, pady=5, columnspan=2)

    def get_context(self):
        question = self.question_entry.get()
        payload = {"question": question}
        response = requests.post(API_URL, headers=headers, json=payload)
        context = response.json().get("context", "")
        self.context_entry.delete(0, tk.END)
        self.context_entry.insert(0, context)

    def ask_ai(self):
        question = self.question_entry.get()
        context = self.context_entry.get()
        payload = {"question": question, "context": context}
        response = requests.post(API_URL2, headers=headers2, json=payload)
        answer = response.json().get("answer", "")
        self.answer_entry.config(state="normal")
        self.answer_entry.delete(0, tk.END)
        self.answer_entry.insert(0, answer)
        self.answer_entry.config(state="readonly")

if __name__ == "__main__":
    root = tk.Tk()
    app = App(root)
    root.mainloop()