"""Tencent-branded Hy3 chat demo, served from a :class:`gradio.Server`.

Goal: ``only the UI and the transport change`` — the chat behaviour is
identical to the previous ``gr.Blocks`` app (``app.py`` + ``chat.py``).
We reuse the exact same incremental **delta-op** streaming protocol the
old app shipped through its hidden ``#hy-chat-delta`` ``gr.HTML`` bridge:
the same op types (``user``, ``assistant_begin``, ``reasoning_delta``,
``content_delta``, ``assistant_end``, ``tool_call``, ``remove_bubble``,
``reset``), the same ``seq`` de-dup, the same ``epoch`` staleness guard,
the same ``preprocess_latex`` on content deltas, the same tool-call
markdown from ``display.format_tool_calls_for_display``, and — critically
— the same ``is_cancelled`` epoch mechanism so clicking "+" mid-stream
abandons the upstream API call server-side instead of billing tokens for
an answer nobody will see.

The only real change vs. the Blocks app is the transport: instead of a
``gr.HTML`` component whose value we mutate, we stream the same JSON
delta payload through ``@app.api`` SSE frames, and the frontend consumes
them with the same renderer (ported from ``static/chat.js``).

Why not reuse ``chat.send_message`` directly? It yields the
7-tuple-of-``gr.update`` shape the Blocks event handlers expect, plus
``gr.Warning`` toasts — all Gradio-coupled. The *pure* turn logic it
wraps (``_run_streaming_turn``, ``_delta``, the op assembly) is already
Gradio-free, so we call that directly and serialise its op-lists to
JSON ourselves.
"""

from __future__ import annotations

import json
import logging
import threading
from pathlib import Path
from typing import Any, Iterator

from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles

from gradio import Server

import chat as _chat  # for _delta / _next_bubble_id (pure helpers reused below)
from chat import _run_streaming_turn
from core.chat import (
    ChatState,
    finalize_response,
    flush_tool_results,
    init_state,
    record_tool_result,
    reset_state_in_place,
)
from display import format_tool_call_prompt
from i18n import t

logger = logging.getLogger(__name__)

_STATIC_DIR = Path(__file__).parent / "static"
_INDEX_HTML = _STATIC_DIR / "index.html"

CHAT_CONCURRENCY = 64

# ---------------------------------------------------------------------------
# Per-session state (identical in spirit to the Blocks app's gr.State)
# ---------------------------------------------------------------------------
# Each browser tab holds a random session id (see static/index.html). We keep
# an in-process ChatState per id so multi-turn + tool-calling work across the
# stateless HTTP boundary. The lock guards the *non-streaming* bookkeeping
# (record/flush tool results, new_chat reset); streaming itself does NOT hold
# the lock (see /chat) so a new_chat can interrupt an in-flight stream.

_SESSIONS: dict[str, dict] = {}
_SESSIONS_LOCK = threading.Lock()


def _get_session(session_id: str) -> dict:
    with _SESSIONS_LOCK:
        sess = _SESSIONS.get(session_id)
        if sess is None:
            sess = {"state": init_state(), "lock": threading.Lock()}
            _SESSIONS[session_id] = sess
        return sess


def _resolve_state(session_id: Any) -> ChatState:
    return _get_session(session_id or "")["state"]


# Reuse chat.py's HTML-safe JSON delta encoder + bubble-id counter so the
# wire format is byte-identical to the Blocks app's #hy-chat-delta payload.
_delta = _chat._delta
_next_bubble_id = _chat._next_bubble_id


# ---------------------------------------------------------------------------
# Server + routes
# ---------------------------------------------------------------------------
app = Server(title="Hy3 · Tencent")


@app.get("/", response_class=HTMLResponse)
async def homepage() -> HTMLResponse:
    """Serve the custom Tencent-branded frontend."""
    return HTMLResponse(_INDEX_HTML.read_text(encoding="utf-8"))


# Static asset passthrough (logo, vendored libs the frontend pulls from the
# same origin instead of a CDN, etc.).
app.mount("/static", StaticFiles(directory=str(_STATIC_DIR)), name="hy-static")


# ---------------------------------------------------------------------------
# Validation — mirror chat._validate_user_message but return the error string
# instead of emitting a gr.Warning toast.
# ---------------------------------------------------------------------------
def _validate_message(message: str, functions_json_str: str) -> str | None:
    """Return None if valid, else a localized error string."""
    if not message or not message.strip():
        return t("warn.empty_msg")
    if functions_json_str and functions_json_str.strip():
        try:
            json.loads(functions_json_str)
        except json.JSONDecodeError as e:
            return t("warn.invalid_fn_json", err=str(e))
    return None


# ---------------------------------------------------------------------------
# A streaming turn = a sequence of JSON delta payloads (the SAME wire format
# the old app pushed through #hy-chat-delta). Each yielded string is a single
# {"seq","ops"[,"epoch"]} document the frontend applies with applyDelta.
# ---------------------------------------------------------------------------

def _frame(ops: list[dict], epoch: int) -> str:
    """Encode a frame's ops as an epoch-tagged delta payload."""
    return _delta(ops, epoch=epoch)


def _stream_turn(
    state: ChatState,
    bubble_id: str,
    epoch: int,
    system_prompt: str,
    think_level: str,
    temperature: float | None,
    max_tokens: int,
    top_p: float,
    preserved_thinking: bool | None,
    functions_json_str: str,
) -> Iterator[str]:
    """Run one assistant turn, yielding delta payloads.

    This is the pure core of the Blocks app's ``_run_streaming_turn`` +
    ``send_message`` streaming section, translated to JSON frames. It:

      * collapses the temperature/preserved_thinking "use default" tristate,
      * runs ``preprocess_latex`` on every content delta (so bare
        ``\\begin{equation}`` and negation macros render correctly),
      * suppresses tool-call snapshots mid-stream and materialises them as a
        ``tool_call`` op at end-of-turn (keeping wire payloads tiny),
      * yields heartbeats (empty-ops frames) so the connection stays alive,
      * bails the moment ``state["epoch"]`` diverges (new_chat clicked) so
        the upstream API stream is abandoned server-side.
    """
    # The Blocks UI collapses (slider, "use default" checkbox) into a tristate
    # BEFORE calling build_api_kwargs. The Server frontend sends the tristate
    # directly (null = use default), so we only need to pass it through.
    temperature_use_default = temperature is None
    preserved_thinking_use_default = preserved_thinking is None

    show_tool = False
    tool_md = ""
    for ops, this_show_tool, this_tool_md in _run_streaming_turn(
        state, bubble_id, epoch,
        system_prompt, think_level,
        # _run_streaming_turn takes the raw slider value + use-default flag;
        # reconstruct them from the tristate the frontend sent.
        float(temperature) if temperature is not None else 0.9,
        temperature_use_default,
        max_tokens, top_p,
        bool(preserved_thinking) if preserved_thinking is not None else False,
        preserved_thinking_use_default,
        functions_json_str,
    ):
        show_tool = show_tool or this_show_tool
        if this_tool_md:
            tool_md = this_tool_md
        yield _frame(ops, epoch), this_show_tool, this_tool_md


# ---------------------------------------------------------------------------
# /chat — the streaming chat endpoint. Mirrors chat.send_message's flow.
# ---------------------------------------------------------------------------
@app.api(
    name="chat",
    concurrency_limit=CHAT_CONCURRENCY,
    concurrency_id="chat",
    stream_every=0.06,
)
def chat(
    message: str,
    session_id: str = "",
    system_prompt: str = "",
    think_level: str = "high",
    temperature: float | None = None,
    max_tokens: int = 0,
    top_p: float = 0,
    preserved_thinking: bool | None = None,
    functions_json_str: str = "",
) -> Iterator[tuple]:
    """Stream a chat turn as delta-op frames.

    Yields ``(delta_json, tool_area_state)`` where ``delta_json`` is the
    JSON payload the frontend applies with ``applyDelta`` (same shape the
    Blocks app pushed through #hy-chat-delta) and ``tool_area_state`` is
    ``None`` normally, or a dict ``{"prompt": <md>, "pending": <n>}`` when
    the model requested tool calls at end-of-turn.
    """
    state = _resolve_state(session_id)

    # ── Validation (mirrors chat._validate_user_message) ──
    err = _validate_message(message, functions_json_str)
    if err is not None:
        # Never started streaming: send a single idle marker frame + the
        # error. The frontend surfaces the error as a toast and keeps the
        # user's message in the input box.
        yield _frame([], state["epoch"]), {"error": err}
        return

    message = message.strip()
    state["messages"].append({"role": "user", "content": message})
    bubble_id = _next_bubble_id()
    epoch = state["epoch"]

    def _cancelled() -> bool:
        return state.get("epoch") != epoch

    # First frame: append the user bubble + open an assistant bubble.
    yield _frame([
        {"type": "user", "text": message},
        {"type": "assistant_begin", "id": bubble_id},
    ], epoch), None

    tool_area = None
    try:
        for frame, show_tool, tool_md in _stream_turn(
            state, bubble_id, epoch, system_prompt, think_level,
            temperature, max_tokens, top_p,
            preserved_thinking, functions_json_str,
        ):
            if _cancelled():
                return
            yield frame, None
            if show_tool:
                tool_area = {"prompt": tool_md}

    except Exception as e:
        if _cancelled():
            return
        logger.exception("chat failed: %s", e)
        # Drop the in-progress assistant bubble, restore the user's message
        # to the input, surface a localized failure — exactly the Blocks
        # app's error path (chat.py:364-381).
        if state["messages"] and state["messages"][-1].get("role") == "user":
            state["messages"].pop()
        yield _frame(
            [{"type": "remove_bubble", "id": bubble_id}], epoch,
        ), {"error": t("warn.request_failed"), "restore_input": message}
        return

    if _cancelled():
        return

    # If the model requested tool calls, mirror the pending queue the Blocks
    # app's finalize_response would have built so submit_tool_result works.
    _sync_pending_tool_calls(state)
    if tool_area and state.get("pending_tool_calls"):
        pending = state["pending_tool_calls"]
        tool_area = {
            "prompt": format_tool_call_prompt(pending[0], 1, len(pending)),
            "pending": len(pending),
        }
        yield _frame([], epoch), tool_area


def _sync_pending_tool_calls(state: ChatState) -> None:
    """Re-derive ``pending_tool_calls`` from the last assistant message.

    ``_run_streaming_turn`` already called ``finalize_response`` (which
    mutates ``state``), so we just read off what it set — identical to the
    state the Blocks app would have after a tool-calling turn.
    """
    # finalize_response (core/chat.py) already populated pending_tool_calls
    # when the model called tools; nothing to re-derive here, but keep the
    # hook for symmetry with the earlier snapshot-based implementation.
    if not state.get("pending_tool_calls"):
        state["pending_assistant_msg"] = None


# ---------------------------------------------------------------------------
# /new_chat — reset. Critically this does NOT block on a running stream:
# the stream's _cancelled() polls state["epoch"], so bumping it here makes
# the in-flight /chat abandon the upstream API on its next yield.
# ---------------------------------------------------------------------------
@app.api(name="new_chat", concurrency_limit=CHAT_CONCURRENCY)
def new_chat(session_id: str = "") -> dict:
    sess = _get_session(session_id or "")
    # Only the bookkeeping is locked; the streaming /chat never holds this
    # lock, so a reset lands instantly even mid-stream.
    with sess["lock"]:
        new_epoch = reset_state_in_place(sess["state"])
    return {"epoch": str(new_epoch), "ok": True}


# ---------------------------------------------------------------------------
# /submit_tool_result — mirrors chat.submit_tool_result.
# ---------------------------------------------------------------------------
@app.api(
    name="submit_tool_result",
    concurrency_limit=CHAT_CONCURRENCY,
    concurrency_id="chat",
    stream_every=0.06,
)
def submit_tool_result(
    result_text: str,
    session_id: str = "",
    system_prompt: str = "",
    think_level: str = "high",
    temperature: float | None = None,
    max_tokens: int = 0,
    top_p: float = 0,
    preserved_thinking: bool | None = None,
    functions_json_str: str = "",
) -> Iterator[tuple]:
    sess = _get_session(session_id or "")
    state = sess["state"]

    with sess["lock"]:
        pending = state.get("pending_tool_calls", [])
        if not pending:
            yield _frame([], state["epoch"]), None
            return

        record_tool_result(state, pending[0], result_text)
        state["pending_tool_calls"] = pending[1:]

        if state["pending_tool_calls"]:
            # More results still pending — keep the tool area open for the
            # next call. Emit a tool_area update, no streaming yet.
            remaining = state["pending_tool_calls"]
            submitted = state.get("submitted_tool_results", [])
            total = len(remaining) + len(submitted)
            current_idx = len(submitted) + 1
            yield _frame([], state["epoch"]), {
                "prompt": format_tool_call_prompt(remaining[0], current_idx, total),
                "pending": len(remaining),
            }
            return

        flush_tool_results(state)

    # All results in — stream the follow-up turn (same protocol as /chat).
    bubble_id = _next_bubble_id()
    epoch = state["epoch"]

    def _cancelled() -> bool:
        return state.get("epoch") != epoch

    yield _frame([{"type": "assistant_begin", "id": bubble_id}], epoch), None

    tool_area = None
    try:
        for frame, show_tool, tool_md in _stream_turn(
            state, bubble_id, epoch, system_prompt, think_level,
            temperature, max_tokens, top_p,
            preserved_thinking, functions_json_str,
        ):
            if _cancelled():
                return
            yield frame, None
            if show_tool:
                tool_area = {"prompt": tool_md}

    except Exception as e:
        if _cancelled():
            return
        logger.exception("submit_tool_result failed: %s", e)
        yield _frame(
            [{"type": "remove_bubble", "id": bubble_id}], epoch,
        ), {"error": t("warn.request_failed")}
        return

    if _cancelled():
        return

    _sync_pending_tool_calls(state)
    if tool_area and state.get("pending_tool_calls"):
        pending = state["pending_tool_calls"]
        yield _frame([], epoch), {
            "prompt": format_tool_call_prompt(pending[0], 1, len(pending)),
            "pending": len(pending),
        }


# ---------------------------------------------------------------------------
# /config — example prompts + i18n. Loaded once on frontend boot.
# ---------------------------------------------------------------------------
@app.api(name="config", concurrency_limit=32)
def config() -> dict:
    """Return the showcase example prompts + i18n strings for the frontend."""
    from app import EXAMPLES  # noqa: WPS433 — lazy so server import stays light
    from i18n import LANG, TRANSLATIONS

    i18n = {k: t(k) for k in TRANSLATIONS.get(LANG, TRANSLATIONS["en"])}
    return {
        "examples": list(EXAMPLES),
        "i18n": i18n,
        "lang": LANG,
        # The tool-call display label + per-op helpers the renderer needs.
        "tool_call_label": t("tool.call_label"),
    }


# ---------------------------------------------------------------------------
# /validate_functions — i18n-aware validation + pretty-print.
# ---------------------------------------------------------------------------
@app.api(name="validate_functions", concurrency_limit=8)
def validate_functions(functions_json_str: str) -> dict:
    """Validate + pretty-print a tools-JSON string. i18n-aware."""
    from tools import parse_functions_json

    if not functions_json_str or not functions_json_str.strip():
        return {"ok": False, "message": t("warn.fn.enter_json"), "formatted": ""}
    try:
        data = json.loads(functions_json_str)
    except json.JSONDecodeError as e:
        return {"ok": False, "message": t("warn.fn.invalid_format", err=str(e)), "formatted": ""}
    if isinstance(data, dict):
        data = [data]
    if not isinstance(data, list):
        return {"ok": False, "message": t("warn.fn.must_be_array"), "formatted": ""}
    names: list[str] = []
    for i, item in enumerate(data):
        if not isinstance(item, dict):
            return {"ok": False, "message": t("warn.fn.item_not_object", i=i + 1), "formatted": ""}
        fn = _normalize(item)
        if not fn:
            return {"ok": False, "message": t("warn.fn.item_invalid", i=i + 1), "formatted": ""}
        if fn["name"] in names:
            return {"ok": False, "message": t("warn.fn.duplicate_name", name=fn["name"]), "formatted": ""}
        names.append(fn["name"])
    formatted = json.dumps(data, indent=2, ensure_ascii=False)
    return {
        "ok": True,
        "message": t("info.fn.validation_passed", n=len(names), names=", ".join(names)),
        "formatted": formatted,
    }


def _normalize(item: dict) -> dict | None:
    """Accept either {type, function} or a bare {name, parameters} entry."""
    if item.get("type") == "function" and "function" in item:
        fn = item["function"]
        if isinstance(fn, dict) and fn.get("name"):
            return fn
        return None
    if item.get("name"):
        return item
    return None


if __name__ == "__main__":
    # HuggingFace Spaces invokes ``python server.py``.
    app.launch(show_error=True)