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

class WikipediaBot:
    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"

    def __init__(self):
        self.summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")

    def search_wikipedia(self, query):
        query = urllib.parse.quote_plus(query)
        data = requests.get(self.SEARCH_TEMPLATE % query).json()
        if data and data[1]:
            page = urllib.parse.quote_plus(data[1][0])
            content = requests.get(self.CONTENT_TEMPLATE % page).json()
            content = list(content["query"]["pages"].values())[0]["extract"]
            summary = self.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(self, message, history):
        summary, source = self.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

def create_chatbot():
    return WikipediaBot()

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")

    bot = create_chatbot()

    msg.submit(bot.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()