import os
import gradio as gr
from huggingface_hub import InferenceClient

# ----------------------------------------------------------------------
# Helper to read a secret (fallback is useful when you run locally)
# ----------------------------------------------------------------------
def _secret(key: str, fallback: str = "") -> str:
    """Return the value of a secret or the supplied fallback."""
    return os.getenv(key, fallback)


# ----------------------------------------------------------------------
# Core chat logic – system prompt comes from the secret `prec_chat`
# ----------------------------------------------------------------------
def respond(
    message: str,
    history: list[dict[str, str]],
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    """
    Generate a response using the HuggingFace Inference API.

    * System prompt = secret `prec_chat`
    * HF inference token = secret `HF_TOKEN`
    """
    # 1️⃣  Load the system prompt (fallback = generic assistant)
    system_message = _secret("prec_chat", "You are a helpful assistant.")

    # 2️⃣  Load the HF inference token
    hf_token = _secret("HF_TOKEN")
    if not hf_token:
        raise RuntimeError(
            "HF_TOKEN not found in secrets. Add it to secrets.toml (or via the Space UI)."
        )

    # 3️⃣  Initialise the HF inference client
    client = InferenceClient(token=hf_token, model="openai/gpt-oss-20b")

    # 4️⃣  Build the message list for the chat‑completion endpoint
    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)                     # previous conversation turns
    messages.append({"role": "user", "content": message})  # current user query

    # 5️⃣  Stream the response back to the UI
    response = ""
    for chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        choices = chunk.choices
        token = ""
        if choices and choices[0].delta.content:
            token = choices[0].delta.content
        response += token
        yield response


# ----------------------------------------------------------------------
# UI – the system‑prompt textbox has been removed.
# ----------------------------------------------------------------------
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        # Only generation parameters are exposed now.
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top‑p (nucleus sampling)",
        ),
    ],
)

# ----------------------------------------------------------------------
# Assemble the Blocks layout (no LoginButton – we use basic auth)
# ----------------------------------------------------------------------
with gr.Blocks() as demo:
    chatbot.render()


# ----------------------------------------------------------------------
# Launch – protect the UI with the credentials from secrets.
# ----------------------------------------------------------------------
if __name__ == "__main__":
    # ------------------------------------------------------------------
    # 1️⃣  Pull the allowed credentials from secrets (fail fast if missing)
    # ------------------------------------------------------------------
    allowed_user = _secret("CHAT_USER")
    allowed_pass = _secret("CHAT_PASS")

    if not allowed_user or not allowed_pass:
        raise RuntimeError(
            "Authentication credentials not found in secrets. "
            "Add CHAT_USER and CHAT_PASS to secrets.toml (or via the HF Spaces UI)."
        )

    # ------------------------------------------------------------------
    # 2️⃣  Launch
    # ------------------------------------------------------------------
    demo.launch(
        auth=(allowed_user, allowed_pass),   # <-- Gradio's built‑in basic auth
        ssr_mode=False,                     # <-- avoids the i18n locale error
        # In a Space we **must not** set share=True (Spaces already give a public URL)
        # If you run locally and want a shareable link, add share=True here.
        server_name="0.0.0.0",              # listen on all interfaces (needed in containers)
        # Optional: give the app a nice title
        # title="Secure Chatbot",
    )