import gradio as gr
from transformers import pipeline
from wikidata_utils import get_wikidata_labels, get_multilingual_classifier

def classify_text(text, topic, language, top_k=5):
    # Get classifier for selected language
    classifier = get_multilingual_classifier(language)
    
    # Get Wikidata labels if topic provided
    candidate_labels = [topic] if not topic else get_wikidata_labels(topic, lang=language)
    
    # Classify text
    result = classifier(text, candidate_labels=candidate_labels)
    
    # Format results
    return {label: float(score) 
            for label, score in zip(result['labels'][:top_k], 
                                   result['scores'][:top_k])}

# Create interface
with gr.Blocks() as demo:
    gr.Markdown("# 🌐 Wikidata-Powered Classifier")
    gr.Markdown("Classify text using live knowledge from Wikidata")
    
    with gr.Row():
        text_input = gr.Textbox(label="Input Text", lines=3)
        topic_input = gr.Textbox(label="Topic (optional)", placeholder="e.g., renewable energy")
    
    with gr.Row():
        lang_dropdown = gr.Dropdown(
            ["en", "es", "fr", "de", "zh"],
            value="en",
            label="Language"
        )
        top_k_slider = gr.Slider(1, 10, value=5, label="Top Results")
    
    submit_btn = gr.Button("Classify")
    result_output = gr.Label(label="Classification Results")
    
    examples = gr.Examples(
        examples=[
            ["Tesla installed solar roofs on 50k homes last quarter", "renewable energy", "en"],
            ["La energía eólica crece un 15% en España", "energía renovable", "es"],
            ["Neue Perowskit-Solarzellen erreichen Rekordwirkungsgrad", "Solarenergie", "de"]
        ],
        inputs=[text_input, topic_input, lang_dropdown]
    )
    
    submit_btn.click(
        fn=classify_text,
        inputs=[text_input, topic_input, lang_dropdown, top_k_slider],
        outputs=result_output
    )

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