import os
import gradio as gr
from huggingface_hub import InferenceClient

# ---------------- Role presets ----------------
ROLE_PRESETS = {
    "Friendly Chatbot": "You are a friendly, concise assistant. Be helpful and keep answers short unless asked.",
    "Dataset Auditor": (
        "You are a dataset QA assistant. Identify duplicates, missing labels, class imbalance, and suspicious samples.\n"
        "Propose concrete fixes and a short remediation checklist."
    ),
    "SQL Explainer": (
        "Translate SQL into plain language for non-technical users. Explain step-by-step, then give a short summary."
    ),
    "Code Reviewer": (
        "Review code succinctly: correctness, readability, performance, edge cases, and security. Provide minimal diffs."
    ),
    "Data Pipeline Doctor": (
        "Debug data pipelines. Use: hypothesis ➜ quick checks ➜ likely fix order. Ask for logs only when needed."
    ),
    "Data Engineering Advisor": (
        "Advise on building robust, scalable data pipelines. Suggest architecture patterns, tools, and optimizations."
    ),
    "ML Dataset Preparer": (
        "Guide on preparing datasets for machine learning: cleaning, splitting, augmentation, and annotation strategies."
    ),
    "Data Quality Analyst": (
        "Evaluate data quality: detect anomalies, missing values, schema mismatches, and provide a remediation plan."
    ),
}

# ---------------- Model ----------------
MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"

# ---------------- Client factory ----------------
def get_client():
    token = os.getenv("HF_TOKEN")  # optional
    return InferenceClient(model=MODEL_NAME, token=token)

# ---------------- Chat backend (streaming) ----------------
def stream_reply(message, history, role, system_message):
    sys_msg = (system_message or "").strip() or ROLE_PRESETS.get(role, "")
    messages = [{"role": "system", "content": sys_msg}]
    for u, a in history:
        if u:
            messages.append({"role": "user", "content": u})
        if a:
            messages.append({"role": "assistant", "content": a})
    messages.append({"role": "user", "content": message})

    client = get_client()
    partial = ""
    try:
        for event in client.chat_completion(messages=messages, stream=True):
            delta = event.choices[0].delta.content or ""
            if delta:
                partial += delta
                yield partial
    except Exception as e:
        yield f"⚠️ Inference error: {e}"

# ---------------- UI ----------------
with gr.Blocks(title="HF Zephyr Chat • Data Roles") as demo:
    gr.Markdown("## 🤗 Zephyr Chat (Data-focused Roles)")

    role_dd = gr.Dropdown(
        label="Role preset",
        choices=list(ROLE_PRESETS.keys()),
        value="Friendly Chatbot",
        interactive=True,
    )

    system_tb = gr.Textbox(
        label="System message (auto-filled by role; you can edit)",
        value=ROLE_PRESETS["Friendly Chatbot"],
        lines=4,
    )

    chat = gr.ChatInterface(
        fn=stream_reply,
        additional_inputs=[role_dd, system_tb],
    )

    def _on_role_change(role):
        return ROLE_PRESETS.get(role, "")
    role_dd.change(fn=_on_role_change, inputs=role_dd, outputs=system_tb)

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