import gradio as gr
import requests
import re
from transformers import pipeline

# Loading the model
summarizer = pipeline(
    "text-generation",
    model="HuggingFaceTB/SmolLM2-135M-Instruct",
    device=-1
)

def clean_query(query):
    query = query.lower()
    query = re.sub(r"what is|who is|tell me about|explain|define", "", query)
    query = re.sub(r'[^\w\s]', '', query)
    return query.strip()

def get_wikipedia_data(query):
    try:
        cleaned = clean_query(query)
        if not cleaned:
            return None, []
            
        search_url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{cleaned.replace(' ', '_')}"
        headers = {'User-Agent': 'MiniWikiBot/1.0 (contact@example.com)'}
        resp = requests.get(search_url, headers=headers, timeout=5)
        
        if resp.status_code == 200:
            data = resp.json()
            return data.get("extract"), []
        
        suggest_url = f"https://en.wikipedia.org/w/api.php?action=opensearch&search={cleaned}&limit=3&namespace=0&format=json"
        s_resp = requests.get(suggest_url, headers=headers, timeout=5)
        suggestions = s_resp.json()[1] if s_resp.status_code == 200 else []
        
        return None, suggestions
    except:
        return None, []

def ai_summarize(text, progress=gr.Progress()):
    if not text or len(text) < 10 or text.startswith("❌"):
        return "⚠️ Please search for a valid topic first."
    
    try:
        progress(0.1, desc="Initializing AI...")
        
        prompt = f"<|im_start|>user\nSummarize the important info from this text:\n{text}\n<|im_end|>\n<|im_start|>assistant\n"
        
        progress(0.3, desc="Generating summary...")
        
        result = summarizer(
            prompt,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            return_full_text=False
        )
        
        progress(0.9, desc="Cleaning up...")
        summary = result[0]['generated_text'].strip()

        summary = re.split(r'<\|im_start\|>|<\|im_end\|>', summary)[0].strip()

        if not summary:
            return "The AI was unable to generate a summary. Try a different topic!"

        return summary
        
    except Exception as e:
        return f"❌ AI Error: {str(e)}"

def process_search(query):
    text, sugs = get_wikipedia_data(query)
    if text:
        # Return Wikipedia text, clear AI output, and show AI button
        return text, "", gr.update(visible=True)
    if sugs:
        sug_text = "❌ Not found. Did you mean:\n" + "\n".join([f"• {s}" for s in sugs])
        return sug_text, "", gr.update(visible=False)
    return "❌ No results found.", "", gr.update(visible=False)

with gr.Blocks(analytics_enabled=False) as demo:
    gr.Markdown("# 🤖 Mini Wikipedia AI")
    
    with gr.Row():
        input_query = gr.Textbox(label="Ask something", placeholder="e.g. Black Holes", scale=4)
        search_btn = gr.Button("Search", variant="primary", scale=1)

    wiki_output = gr.Textbox(label="Wikipedia Summary", interactive=False, lines=8)
    
    with gr.Row():
        ai_btn = gr.Button("✨ Summarize with AI", visible=False)
    
    ai_output = gr.Textbox(label="AI Summary Result", interactive=False, placeholder="AI summary will appear here...")

    # Logic for searching
    search_btn.click(process_search, inputs=[input_query], outputs=[wiki_output, ai_output, ai_btn])
    input_query.submit(process_search, inputs=[input_query], outputs=[wiki_output, ai_output, ai_btn])
    
    # Logic for AI Summarization
    ai_btn.click(
        fn=lambda: "", 
        outputs=[ai_output]
    ).then(
        fn=ai_summarize, 
        inputs=[wiki_output], 
        outputs=[ai_output]
    )

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(),
        show_error=True
    )