import os
import spaces
import urllib.parse
import requests
import gradio as gr
from transformers import pipeline

SEARCH_TEMPLATE = "https://en.wikipedia.org/w/api.php?action=opensearch&search=%s&limit=1&namespace=0&format=json"
CONTENT_TEMPLATE = "https://en.wikipedia.org/w/api.php?format=json&action=query&prop=extracts&exintro&explaintext&redirects=1&titles=%s"

summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")

@spaces.GPU
def search_wikipedia(query):
    query = urllib.parse.quote_plus(query)
    data = requests.get(SEARCH_TEMPLATE % query).json()
    if data and data[1]:
        page = urllib.parse.quote_plus(data[1][0])
        content = requests.get(CONTENT_TEMPLATE % page).json()
        content = list(content["query"]["pages"].values())[0]["extract"]
        summary = summarizer(content, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
        source = data[3][0]
        return summary, source
    else:
        return "No results found.", ""

def chatbot_response(message, history):
    summary, source = search_wikipedia(message)
    response = f"Here's a summary of the top Wikipedia result:\n\n{summary}"
    if source:
        response += f"\n\nSource: {source}"
    history.append((message, response))
    return history

demo = gr.Blocks()

with demo:
    gr.Markdown("# Wikipedia Chatbot")
    gr.Markdown("This chatbot queries the Wikipedia API and summarizes the top result.")

    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Enter your query")
    clear = gr.Button("Clear")

    msg.submit(chatbot_response, [msg, chatbot], chatbot).then(
        lambda: gr.update(value=""), None, [msg], queue=False
    )
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()