import os
import gradio as gr
from transformers import pipeline

# Read API key from environment
HF_TOKEN = os.getenv("HF_API_TOKEN")

qa = pipeline(
    "question-answering",
    model="distilbert-base-uncased-distilled-squad",
    token=HF_TOKEN
)

def answer(question, context):
    if not context.strip():
        return "⚠️ Please paste document content first."

    if not question.strip():
        return "⚠️ Please enter a question."

    result = qa(question=question, context=context)
    answer_text = result.get("answer", "").strip()

    if not answer_text:
        return "❌ Answer not found in the provided content."

    return answer_text

with gr.Blocks() as demo:
    gr.Markdown("## ⚡ Free AI Document Chatbot")

    context = gr.Textbox(
        label="📄 Paste document text",
        lines=10
    )

    question = gr.Textbox(
        label="❓ Ask a question"
    )

    output = gr.Textbox(label="✅ Answer")

    ask = gr.Button("Ask")

    ask.click(answer, [question, context], output)

demo.launch()