"""Gradio chat client for a Hugging Face Inference Endpoint."""

import itertools
import json
import logging
import os
import sys
import time
from secrets import compare_digest
from typing import NamedTuple

import gradio as gr
import requests

# Root stays at WARNING so LOG_LEVEL doesn't also enable httpx/urllib3.
logging.basicConfig(
    level=logging.WARNING,
    format="%(asctime)s %(levelname)-7s %(name)s %(message)s",
    stream=sys.stdout,
)

logger = logging.getLogger("gemma-endpoint")
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO").upper())

_turn_ids = itertools.count(1)


class EndpointWaking(Exception):
    """Raised before any tokens, so retrying the request is safe."""


REQUEST_TIMEOUT_SECONDS = 300
MAX_TOKENS = 4096  # reasoning models spend most of a smaller budget thinking

# A cold endpoint rejects outright instead of holding the request.
WAKE_STATUSES = (502, 503, 504)
AUTH_STATUSES = (401, 403)

WAKE_TIMEOUT_SECONDS = int(os.environ.get("WAKE_TIMEOUT_SECONDS", "300"))
WAKE_RETRY_SECONDS = 5
WAKE_TICK_SECONDS = 1  # shorter than the retry so Stop stays responsive

# `status: pending` renders a spinner. It needs text: reasoning_tags drops
# empty messages, spinner included.
WAITING_METADATA = {"title": "Waiting for the endpoint…", "status": "pending"}
WAITING_TEXT = "If the endpoint was asleep, the first reply can take a few minutes."

INPUT_PLACEHOLDER = "Message the model"

# llama.cpp sends thinking as reasoning_content; these tags make it collapsible.
THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"

# The default 512-byte reads deliver the whole reply at once.
STREAM_CHUNK_SIZE = 1

GEMMA_MODEL = "google/gemma-4-12B-it-qat-q4_0-gguf"
QWEN_MODEL = "unsloth/Qwen3.5-9B-GGUF"
DEEPSEEK_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"

MODEL_LABELS = {
    GEMMA_MODEL: "Gemma 4 12B",
    QWEN_MODEL: "Qwen 3.5 9B",
    DEEPSEEK_MODEL: "DeepSeek R1 14B",
}


def _required_env(name):
    value = os.environ.get(name, "").strip()

    if not value:
        raise RuntimeError(
            f"{name} is not set. Set it as a Space secret (or export it "
            f"locally) before starting the app."
        )

    return value


class Endpoint(NamedTuple):
    url: str
    token: str
    token_env: str  # secret name, for logs and errors; never the token itself


def _endpoint(url_env, token_env):
    """Per-org token, falling back to HF_TOKEN."""
    token = os.environ.get(token_env, "").strip()

    return Endpoint(
        url=_required_env(url_env).rstrip("/"),
        token=token or HF_TOKEN,
        token_env=token_env if token else "HF_TOKEN",
    )


HF_TOKEN = _required_env("HF_TOKEN")
ENDPOINT_URL = _required_env("ENDPOINT_URL").rstrip("/")
LOGIN_PASSWORD = os.environ.get("LOGIN_PASSWORD", "")

MODELS = {
    GEMMA_MODEL: Endpoint(ENDPOINT_URL, HF_TOKEN, "HF_TOKEN"),
    QWEN_MODEL: _endpoint("QWEN_ENDPOINT_URL", "QWEN_HF_TOKEN"),
    DEEPSEEK_MODEL: _endpoint("DEEPSEEK_ENDPOINT_URL", "DEEPSEEK_HF_TOKEN"),
}

AVAILABLE_MODELS = list(MODELS)
DEFAULT_MODEL = GEMMA_MODEL


def _log_config():
    for model, endpoint in MODELS.items():
        logger.info(
            "configured model=%s url=%s token=%s",
            model,
            endpoint.url,
            endpoint.token_env,
        )

    logger.info(
        "configured auth=%s timeout=%ss max_tokens=%s",
        "on" if LOGIN_PASSWORD else "OFF",
        REQUEST_TIMEOUT_SECONDS,
        MAX_TOKENS,
    )


_log_config()


def endpoint_for(model_name):
    return MODELS.get(model_name) or Endpoint(ENDPOINT_URL, HF_TOKEN, "HF_TOKEN")


def message_text(content):
    """Gradio content is a string or a list of typed parts; keep the text."""
    if isinstance(content, str):
        return content

    if isinstance(content, (list, tuple)):
        return "".join(
            part.get("text") or ""
            for part in content
            if isinstance(part, dict) and part.get("type") == "text"
        )

    return ""


def normalize_history(history):
    """Plain text turns, without spinners or other metadata messages."""
    turns = []

    for item in history or []:
        if not isinstance(item, dict) or item.get("role") not in ("user", "assistant"):
            continue

        if (item.get("metadata") or {}).get("title"):
            continue

        text = message_text(item.get("content"))

        if text:
            turns.append({"role": item["role"], "content": text})

    return turns


def build_api_messages(history, message):
    return normalize_history(history) + [{"role": "user", "content": message}]


def stream_completion(messages, model_name, turn="-"):
    """Yield SSE chunks. reasoning_content is wrapped in think tags."""
    endpoint = endpoint_for(model_name)
    prompt_chars = sum(len(m["content"]) for m in messages)
    logger.info(
        "turn=%s POST %s/v1/chat/completions model=%s token=%s messages=%d "
        "prompt_chars=%d",
        turn,
        endpoint.url,
        model_name,
        endpoint.token_env,
        len(messages),
        prompt_chars,
    )
    logger.debug("turn=%s prompt=%s", turn, messages)

    started = time.monotonic()

    with requests.post(
        f"{endpoint.url}/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {endpoint.token}",
            "Content-Type": "application/json",
        },
        json={
            "model": model_name,
            "messages": messages,
            "max_tokens": MAX_TOKENS,
            "stream": True,
        },
        timeout=REQUEST_TIMEOUT_SECONDS,
        stream=True,
    ) as response:
        logger.info(
            "turn=%s response status=%s after %.2fs",
            turn,
            response.status_code,
            time.monotonic() - started,
        )

        if response.status_code in WAKE_STATUSES:
            raise EndpointWaking(response.status_code)

        if response.status_code in AUTH_STATUSES:
            raise gr.Error(
                f"{MODEL_LABELS.get(model_name, model_name)} rejected the token "
                f"(HTTP {response.status_code}). Check that {endpoint.token_env} "
                f"has access to the organisation that owns this endpoint."
            )

        response.raise_for_status()

        # requests guesses ISO-8859-1 when the charset is missing.
        response.encoding = "utf-8"

        chunks = 0
        thinking = False

        for line in response.iter_lines(
            chunk_size=STREAM_CHUNK_SIZE, decode_unicode=True
        ):
            if not line or not line.startswith("data:"):
                continue

            payload = line[len("data:") :].strip()

            if payload == "[DONE]":
                logger.info(
                    "turn=%s stream done chunks=%d in %.2fs",
                    turn,
                    chunks,
                    time.monotonic() - started,
                )
                break

            choices = json.loads(payload).get("choices") or [{}]
            delta = choices[0].get("delta") or {}
            reasoning = delta.get("reasoning_content")
            content = delta.get("content")
            piece = ""

            if reasoning:
                if not thinking:
                    piece += THINK_OPEN
                    thinking = True

                piece += reasoning

            if content:
                if thinking:
                    piece += THINK_CLOSE
                    thinking = False

                piece += content

            if piece:
                chunks += 1

                if chunks == 1:
                    logger.info(
                        "turn=%s first token after %.2fs",
                        turn,
                        time.monotonic() - started,
                    )

                yield piece

        else:
            logger.warning(
                "turn=%s stream ended without [DONE] after %d chunks", turn, chunks
            )

        if thinking:
            yield THINK_CLOSE


def waiting_turn(title=None):
    metadata = dict(WAITING_METADATA)

    if title:
        metadata["title"] = title

    return {"role": "assistant", "content": WAITING_TEXT, "metadata": metadata}


def stream_turn(messages, model_name, turn="-"):
    """Yield ("waking", seconds) while the endpoint is cold, then ("token", text)."""
    started = time.monotonic()
    deadline = started + WAKE_TIMEOUT_SECONDS
    attempt = 0

    while True:
        attempt += 1

        try:
            for chunk in stream_completion(messages, model_name, turn=turn):
                yield "token", chunk

            return

        except EndpointWaking as waking:
            waited = time.monotonic() - started

            if time.monotonic() >= deadline:
                raise gr.Error(
                    f"The endpoint did not come up within "
                    f"{WAKE_TIMEOUT_SECONDS}s (last status {waking}). Check that "
                    f"it is Running and not scaled to zero."
                ) from None

            logger.info(
                "turn=%s endpoint waking status=%s attempt=%d waited=%.0fs",
                turn,
                waking,
                attempt,
                waited,
            )

            for _ in range(WAKE_RETRY_SECONDS // WAKE_TICK_SECONDS):
                if time.monotonic() >= deadline:
                    break

                yield "waking", int(time.monotonic() - started)

                time.sleep(WAKE_TICK_SECONDS)


def respond(message, history, model_name):
    history = list(history or [])
    message = (message or "").strip()

    if not message:
        logger.debug("ignoring an empty submission")
        yield history, *idle_controls()
        return

    turn = next(_turn_ids)
    started = time.monotonic()

    api_messages = build_api_messages(history, message)
    logger.info(
        "turn=%s start model=%s history_turns=%d message_chars=%d",
        turn,
        model_name or DEFAULT_MODEL,
        len(api_messages) - 1,
        len(message),
    )

    history.append({"role": "user", "content": message})
    history.append(waiting_turn())

    yield list(history), *busy_controls()

    reply = ""

    try:
        for kind, value in stream_turn(
            api_messages, model_name or DEFAULT_MODEL, turn=turn
        ):
            if kind == "waking":
                history[-1] = waiting_turn(f"Waking up the endpoint… {value}s")
            else:
                reply += value
                history[-1] = {"role": "assistant", "content": reply}

            yield list(history), *busy_controls()

        if not reply:
            raise gr.Error("The endpoint returned an empty reply.")

    except Exception as error:
        if isinstance(error, gr.Error):
            detail = str(error)

        elif isinstance(error, requests.exceptions.Timeout):
            detail = "The endpoint timed out. Check that it is in the Running state."

        elif isinstance(error, requests.exceptions.RequestException):
            detail = f"Request to the endpoint failed: {_error_detail(error)}"

        elif isinstance(error, (KeyError, IndexError, ValueError)):
            detail = f"Unexpected response from the endpoint: {error}"

        else:
            logger.exception("turn=%s unexpected failure", turn)
            raise

        logger.error(
            "turn=%s failed after %.2fs with %d chars streamed: %s",
            turn,
            time.monotonic() - started,
            len(reply),
            detail,
        )

        if reply:
            yield list(history), *idle_controls()
        else:
            del history[-2:]
            yield list(history), *idle_controls(message)

        raise gr.Error(detail, print_exception=False) from None

    logger.info(
        "turn=%s complete reply_chars=%d in %.2fs",
        turn,
        len(reply),
        time.monotonic() - started,
    )
    logger.debug("turn=%s reply=%s", turn, reply)

    yield list(history), *idle_controls()


def busy_controls(value=""):
    return (
        gr.update(value=value, interactive=False, placeholder="Replying…"),
        gr.update(interactive=False, visible=False),
        gr.update(interactive=True, visible=True),
    )


def idle_controls(value=""):
    return (
        gr.update(value=value, interactive=True, placeholder=INPUT_PLACEHOLDER),
        gr.update(interactive=True, visible=True),
        gr.update(interactive=False, visible=False),
    )


def cancel_stream(history):
    logger.info("stopped by the user")

    return normalize_history(history), *idle_controls()


def new_chat():
    logger.info("new chat")

    return [], *idle_controls()


def model_caption(model_name):
    return f"`{model_name or DEFAULT_MODEL}`"


def example_text(event: gr.SelectData):
    return event.value["text"]


def _error_detail(error):
    if error.response is None:
        return str(error)

    try:
        return error.response.text or str(error)

    except requests.exceptions.RequestException:
        return str(error)


def verify_login(username, password):
    allowed = compare_digest(password or "", LOGIN_PASSWORD)

    logger.log(
        logging.INFO if allowed else logging.WARNING,
        "login %s for username=%r",
        "succeeded" if allowed else "REJECTED",
        username,
    )

    return allowed


EXAMPLES = [
    {"display_text": "Reply in Japanese", "text": "Introduce yourself in Japanese."},
    {
        "display_text": "Reason step by step",
        "text": "Is 1,001 a prime number? Work it out step by step.",
    },
    {
        "display_text": "Write some code",
        "text": "Write a Python function that reverses a singly linked list.",
    },
    {
        "display_text": "Keep it short",
        "text": "Explain what a hash map is in two sentences.",
    },
]

EMPTY_STATE = (
    "**Pick a model and send a message**\n\n"
    "A sleeping endpoint can take a few minutes to answer the first message."
)

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.teal,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace"],
    radius_size=gr.themes.sizes.radius_md,
).set(
    body_background_fill="#f2f4f6",
    body_background_fill_dark="#0e1216",
    block_border_width="0px",
    block_shadow="none",
    button_primary_background_fill="#0b6e79",
    button_primary_background_fill_hover="#095a63",
    button_primary_background_fill_dark="#14909d",
    button_primary_background_fill_hover_dark="#18a5b3",
    button_primary_text_color="#ffffff",
)

CSS = """
/* Gradio caps .main at 640px, which squeezes the chat. */
.gradio-container, .gradio-container > .main {
    max-width: 880px !important; width: 100%; margin: 0 auto; }

#header { align-items: center; gap: 12px; }
#header h1 { font-size: 1.35rem; font-weight: 600; margin: 0; }
#header p { margin: 2px 0 0; color: var(--body-text-color-subdued); }
#header button { background: var(--block-background-fill);
    border: 1px solid var(--border-color-primary); border-radius: 8px; }
#header button:hover { border-color: var(--button-primary-background-fill); }

#model-picker .wrap { gap: 4px; padding: 4px; border-radius: 10px;
    background: var(--block-background-fill);
    border: 1px solid var(--border-color-primary); width: fit-content; }
#model-picker label { border: 0 !important; box-shadow: none !important;
    background: transparent !important; padding: 6px 14px !important;
    border-radius: 7px !important; cursor: pointer; }
#model-picker label.selected { background: var(--button-primary-background-fill) !important;
    color: var(--button-primary-text-color) !important; }
#model-picker input[type="radio"] { position: absolute; opacity: 0;
    width: 1px; height: 1px; pointer-events: none; }
#model-picker label:has(input:focus-visible) {
    outline: 2px solid var(--button-primary-background-fill); outline-offset: 2px; }
#model-id { margin-top: -6px; }
#model-id p, #model-id code { font-size: 0.8rem; color: var(--body-text-color-subdued); }
#model-id code { background: none; border: 0; padding: 0; }

#chat {
    background: var(--background-fill-primary) !important;
    border: 1px solid var(--border-color-primary) !important;
    border-radius: 12px !important; }
#chat .bubble-wrap { background: transparent; }

#chat .placeholder-content { justify-content: center; gap: 20px; padding: 24px; }
#chat .placeholder { flex: 0 0 auto !important; height: auto !important;
    text-align: center; }
#chat .placeholder p { margin: 4px 0; color: var(--body-text-color-subdued); }
#chat .placeholder strong { color: var(--body-text-color); font-size: 1.05rem; }
#chat .examples { width: 100%; max-width: 560px; margin: 0 auto; gap: 10px; }
#chat .example { background: var(--background-fill-secondary) !important;
    border: 1px solid var(--border-color-primary) !important;
    border-radius: 10px !important; text-align: left; }
#chat .example:hover { border-color: var(--button-primary-background-fill) !important; }

#composer { align-items: stretch; gap: 8px; }
#composer textarea { font-size: 0.975rem; padding: 12px 14px !important;
    background: var(--block-background-fill) !important;
    border: 1px solid var(--border-color-primary) !important;
    border-radius: 10px !important; }
#composer textarea:focus { border-color: var(--button-primary-background-fill) !important;
    box-shadow: 0 0 0 3px color-mix(in srgb,
        var(--button-primary-background-fill) 20%, transparent) !important; }
#composer button { min-width: 88px; border-radius: 10px; }
"""


def create_app():
    with gr.Blocks(title="Endpoint chat") as demo:
        with gr.Row(elem_id="header", equal_height=True):
            gr.Markdown(
                "# Endpoint chat\nChat with your Hugging Face Inference Endpoints.",
            )
            new_chat_button = gr.Button("New chat", size="sm", scale=0, min_width=110)

        model_selector = gr.Radio(
            choices=[(MODEL_LABELS[model], model) for model in AVAILABLE_MODELS],
            value=DEFAULT_MODEL,
            label="Model",
            show_label=False,
            container=False,
            elem_id="model-picker",
        )
        model_id = gr.Markdown(model_caption(DEFAULT_MODEL), elem_id="model-id")

        chatbot = gr.Chatbot(
            show_label=False,
            height="calc(100dvh - 330px)",
            min_height=320,
            placeholder=EMPTY_STATE,
            examples=EXAMPLES,
            reasoning_tags=[(THINK_OPEN, THINK_CLOSE)],
            buttons=["copy", "copy_all"],
            feedback_options=None,
            elem_id="chat",
        )

        with gr.Row(elem_id="composer"):
            msg = gr.Textbox(
                label="Message",
                show_label=False,
                placeholder=INPUT_PLACEHOLDER,
                lines=1,
                max_lines=8,
                autofocus=True,
                container=False,
                scale=1,
            )
            submit_button = gr.Button("Send", variant="primary", scale=0)
            stop_button = gr.Button(
                "Stop", variant="stop", scale=0, interactive=False, visible=False
            )

        controls = [msg, submit_button, stop_button]
        chat_outputs = [chatbot, *controls]
        respond_inputs = [msg, chatbot, model_selector]

        send_events = [
            trigger(fn=respond, inputs=respond_inputs, outputs=chat_outputs)
            for trigger in (submit_button.click, msg.submit)
        ]
        send_events.append(
            chatbot.example_select(fn=example_text, outputs=msg).then(
                fn=respond, inputs=respond_inputs, outputs=chat_outputs
            )
        )

        # Cancelling skips the generator's cleanup, so Stop does it itself.
        stop_button.click(
            fn=cancel_stream,
            inputs=[chatbot],
            outputs=chat_outputs,
            cancels=send_events,
        )
        new_chat_button.click(fn=new_chat, outputs=chat_outputs, cancels=send_events)
        model_selector.change(fn=model_caption, inputs=model_selector, outputs=model_id)

    return demo


# Spaces' reloader looks up `demo` in __main__.
demo = create_app()


if __name__ == "__main__":
    if not LOGIN_PASSWORD:
        print(
            "WARNING: LOGIN_PASSWORD is not set - the app is starting without "
            "authentication and anyone who can reach it can spend endpoint quota."
        )

    demo.launch(auth=verify_login if LOGIN_PASSWORD else None, theme=THEME, css=CSS)