import os
import gradio as gr
from huggingface_hub import InferenceClient

# ----------------------------------------------------------------------
# 1️⃣  Force the UI language (prevents the “svelte‑i18n” error)
# ----------------------------------------------------------------------
gr.set_default_language("en")          # English UI – change if you need another locale

# ----------------------------------------------------------------------
# Helper to read a secret (with a safe fallback for local testing)
# ----------------------------------------------------------------------
def _secret(key: str, fallback: str = "") -> str:
    """Return the value of a secret or the supplied fallback."""
    return os.getenv(key, fallback)


# ----------------------------------------------------------------------
# 2️⃣  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,
    hf_token: gr.OAuthToken,
):
    """
    Generate a response using the HuggingFace Inference API.

    The system prompt is taken from the secret **prec_chat**.
    Users cannot edit it from the UI.
    """
    # Load the system prompt (fallback = generic assistant)
    system_message = _secret("prec_chat", "You are a helpful assistant.")

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

    # 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

    # 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


# ----------------------------------------------------------------------
# 3️⃣  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)",
        ),
    ],
)

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


# ----------------------------------------------------------------------
# 5️⃣  Launch – protect the UI with the credentials from secrets.
# ----------------------------------------------------------------------
if __name__ == "__main__":
    # Pull the allowed credentials from secrets (raise early if they are 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)."
        )

    demo.launch(
        auth=(allowed_user, allowed_pass),   # <-- Gradio's built‑in basic auth
        # In a remote environment (HF Spaces, Docker, cloud VM) you need a shareable link:
        share=True,                          # <-- remove if you run locally and can reach http://0.0.0.0:7860
        # Optional – makes the server listen on all interfaces (useful in containers)
        server_name="0.0.0.0",
        # Optional – you can set a custom title, favicon, etc.
        # title="Secure Chatbot",
    )