import gradio as gr
from huggingface_hub import InferenceClient
import requests

def wiki_info(frage: str) -> str:
    search_url = "https://de.wikipedia.org/w/api.php"
    search_params = {
        "action": "query",
        "list": "search",
        "srsearch": frage,
        "format": "json"
    }

    r = requests.get(search_url, params=search_params)
    data = r.json()

    suchergebnisse = data.get("query", {}).get("search", [])
    if not suchergebnisse:
        return None

    pageid = suchergebnisse[0]["pageid"]

    # 2. Volltext laden
    page_url = "https://de.wikipedia.org/w/api.php"
    page_params = {
        "action": "query",
        "prop": "extracts",
        "pageids": pageid,
        "explaintext": 1,
        "format": "json"
    }

    r2 = requests.get(page_url, params=page_params)
    page_data = r2.json()

    pages = page_data.get("query", {}).get("pages", {})
    page = pages.get(str(pageid), {})

    text = page.get("extract")
    if not text:
        return None

    saetze = re.split(r'(?<=[.!?]) +', text)
    gekuerzt = " ".join(saetze[:10])

    return gekuerzt

def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
    use_wiki # Checkbox Input
):
    client = InferenceClient(token=hf_token.token, model="meta-llama/Llama-3.2-1B-Instruct")

    # --- HIER PASSIERT DAS IN-CONTEXT LEARNING ---
    context_text = ""
    if use_wiki:
        wiki_content = get_wikipedia_summary(message)
        if wiki_content:
            context_text = f"\n\nEXTERNER KONTEXT (WIKIPEDIA): {wiki_content}\n"
            gr.Info(f"Kontext gefunden: {wiki_content[:50]}...") # Kleines UI Feedback
        else:
            gr.Info("Kein Wikipedia-Artikel gefunden.")

    # Der Prompt zwingt das Modell, den Kontext zu nutzen
    final_system_prompt = (
        f"{system_message} "
        f"Wenn 'EXTERNER KONTEXT' bereitgestellt wird, nutze dieses Wissen, um die Frage zu beantworten. "
        f"Verlasse dich mehr auf den Kontext als auf dein eigenes Wissen."
        f"{context_text}"
    )

    messages = [{"role": "system", "content": final_system_prompt}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    response = ""
    for msg in client.chat_completion(
        messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p,
    ):
        token = msg.choices[0].delta.content
        if token:
            response += token
            yield response

# --- GUI ---
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Der Wikipedia-gestützte Assistent")
    gr.Markdown("Stelle eine Frage. Wenn du die Checkbox aktivierst, suche ich live nach Fakten!")
    
    with gr.Row():
        wiki_checkbox = gr.Checkbox(label="Nutze Wikipedia-Wissen (RAG)", value=True)

    chatbot = gr.ChatInterface(
        respond,
        additional_inputs=[
            gr.Textbox(value="Du bist ein hilfreicher Assistent der Dinge genau und exakt erklärt.", label="System"),
            gr.Slider(1, 1024, 512, label="Max Tokens"),
            gr.Slider(0.1, 2.0, 0.7, label="Temp"),
            gr.Slider(0.1, 1.0, 0.95, label="Top-p"),
            wiki_checkbox
        ]
    )
    with gr.Sidebar():
        gr.LoginButton()

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