"""Gradio chat interface for talkie-1930-13b-it."""

from __future__ import annotations

import os

import gradio as gr

from talkie import Message, Talkie


MODEL_NAME = "talkie-1930-13b-it"

print(f"Loading {MODEL_NAME}…", flush=True)
model = Talkie(MODEL_NAME)
print("Model loaded.", flush=True)


def chat_fn(message: str, history: list[dict], temperature=0.7, max_tokens=400, top_p=0.0):
    safe_temperature = float(temperature) if temperature is not None else 0.7
    safe_max_tokens = int(max_tokens) if max_tokens is not None else 400
    safe_top_p_value = top_p if top_p is not None else 0.0
    effective_top_p = float(safe_top_p_value) if safe_top_p_value > 0 else None

    prior_messages = [
        Message(role=turn["role"], content=turn["content"])
        for turn in history
        if turn.get("role") in {"user", "assistant", "system"}
    ]
    prior_messages.append(Message(role="user", content=message))

    partial_response = ""
    for token in model.chat_stream(
        prior_messages,
        temperature=safe_temperature,
        max_tokens=safe_max_tokens,
        top_p=effective_top_p,
    ):
        partial_response += token
        yield partial_response


description = (
    "A 13B instruction-tuned vintage language model trained on pre-1931 "
    "English-language text. See the "
    "[model card](https://huggingface.co/talkie-lm/talkie-1930-13b-it) "
    "and [project page](https://talkie-lm.com/) for details."
)


demo = gr.ChatInterface(
    fn=chat_fn,
    title="talkie-1930-13b-it",
    description=description,
    additional_inputs=[
        gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature"),
        gr.Slider(16, 1024, value=400, step=16, label="Max tokens"),
        gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Top-p (0 disables)"),
    ],
    examples=[
        ["Write an essay predicting what life will be like in the year 1960."],
        ["What were the causes of the French Revolution?"],
        ["Compose a polite letter declining a dinner invitation."],
    ],
)


AUTH_MESSAGE = (
    "Username: guest    |    Password: ask the owner. "
    "(The username field will accept any value — only the password is checked.)"
)


def _check_password(_username: str, password: str) -> bool:
    expected_password = os.environ.get("AUTH_PASSWORD")
    if not expected_password:
        return False
    return password == expected_password


def _require_auth_callable():
    if not os.environ.get("AUTH_PASSWORD"):
        raise RuntimeError(
            "AUTH_PASSWORD must be set as a Space secret. "
            "Refusing to launch without authentication on a public Space."
        )
    return _check_password


if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        auth=_require_auth_callable(),
        auth_message=AUTH_MESSAGE,
    )