"""Gradio entrypoint for the TuringDNA assistant Space.

Two surfaces:

1.  Gradio ChatInterface at "/" — for humans visiting
    huggingface.co/spaces/winter4000/turingdna-assistant directly to
    sanity-check the model.

2.  Auto-generated Gradio API at "/api/predict" (or /run/predict on
    older Gradio) — what the Flask app's /api/chat proxy calls via
    gradio_client.Client(...).predict(...).

The Flask side uses `api_name="/chat"`, so the ChatInterface event is
registered with that api_name.

ZeroGPU note: the actual inference is in llm.generate() which carries
the @spaces.GPU(duration=60) decorator. From here we just hand control
to ChatInterface — Gradio handles the request queueing and ZeroGPU
acquires the A100 when generate() is called.
"""

import os
import sys

import gradio as gr
import fastapi
import starlette
import jinja2

# Print the installed versions of the web stack so the next failure
# log tells us exactly what we're running against. The Gradio 4.44 +
# Starlette + jinja2 interaction has been the root cause of every
# build issue so far; this single line saves a lot of debugging if
# the constellation changes underneath us.
print(
    f"[turingdna.assistant] gradio={gr.__version__} fastapi={fastapi.__version__} "
    f"starlette={starlette.__version__} jinja2={jinja2.__version__} python={sys.version.split()[0]}",
    flush=True,
)

from llm import chat, MODEL_ID_IN_USE, SYSTEM_PROMPT


def chat_fn(message, history=None) -> str:
    """Gradio ChatInterface callback.

    Deliberately IGNORES Gradio's session `history` — the Flask client
    embeds conversation context into the message string itself. Single
    source of truth = the message string.

    Bulletproof type coercion: Gradio 4.44 has inconsistent type
    handling depending on call path (UI vs API), version, and theme.
    We've seen message arrive as: a str, a {'role','content'} dict, a
    list-of-dicts, and an empty string. Coerce everything to a plain
    str BEFORE anything else touches it.

    Also: this entire function is wrapped in a top-level exception
    handler that logs the traceback to stderr so the next mystery
    'list' object has no attribute 'get' tells us which LINE
    blew up (the previous error format was useless for debugging
    because it stripped the traceback).
    """
    import sys
    import traceback

    # Logging payload (truncated for log readability).
    print(
        f"[chat_fn] in  message_type={type(message).__name__} "
        f"history_type={type(history).__name__} "
        f"history_len={len(history) if hasattr(history, '__len__') else 'n/a'}",
        flush=True,
    )
    try:
        # ── Coerce message to a plain string, no matter what shape ──
        if message is None:
            text = ""
        elif isinstance(message, str):
            text = message
        elif isinstance(message, dict):
            # OpenAI-style {role, content} — take content.
            text = str(message.get("content") or "")
        elif isinstance(message, (list, tuple)):
            # A list of messages — take the last one's content.
            if not message:
                text = ""
            else:
                last = message[-1]
                if isinstance(last, dict):
                    text = str(last.get("content") or "")
                else:
                    text = str(last)
        else:
            text = str(message)

        text = text.strip()
        if not text:
            return "Ask me a question about your sequence, mutations, or library."

        # history=[] always — Flask owns conversation context.
        return chat(text, history=[])

    except Exception as exc:  # noqa: BLE001
        # FULL traceback to stderr so the next failure log is debuggable.
        # Returning a clean message to the user but the operator sees
        # exactly which line raised what.
        tb = traceback.format_exc()
        print(
            f"[chat_fn] EXCEPTION {type(exc).__name__}: {exc}\n{tb}",
            file=sys.stderr,
            flush=True,
        )
        return f"Generation error: {type(exc).__name__}: {exc}"


# Modest demo prompts shown above the composer — give the visiting
# researcher (or me, when sanity-checking) a one-click way to see what
# the assistant is good at without having to compose a question from
# scratch.
DEMO_PROMPTS = [
    "Explain the ESM-2 ΔLL score in one paragraph for someone with a biochem background.",
    "I have a 280-aa enzyme and the engine is suggesting M1L, V8I, and K42R. Which of these is most likely to affect activity, and why?",
    "What's the difference between site-saturation mutagenesis and the combinatorial library approach this engine uses?",
    "Why does codon optimization for E. coli sometimes hurt expression in yeast?",
]


# Note: we used to wrap this ChatInterface in a multi-tab gr.Blocks (Chat
# tab + System-prompt-viewer tab). That pattern is shaky in Gradio 4.x —
# .render() on a ChatInterface inside another Blocks context didn't
# initialise the queue cleanly and produced 500s on every submit. Dropped
# the wrapper; the ChatInterface launches directly. System prompt is
# documented in README.md instead.
demo = gr.ChatInterface(
    fn=chat_fn,
    type="messages",
    title="TuringDNA Assistant",
    description=(
        f"In-app biology assistant powered by **{MODEL_ID_IN_USE}** on ZeroGPU. "
        "Ask about protein sequences, ΔLL scores, codon optimization, or cloning strategy. "
        "First call after the Space wakes up takes 10–30 s while the GPU is acquired."
    ),
    examples=DEMO_PROMPTS,
    cache_examples=False,
    # NOTE: api_name kwarg is a Gradio 5.x feature; we're pinned to 4.44.
    # In 4.x the ChatInterface submit endpoint auto-registers as "/chat"
    # via the function name fallback. gradio_client.predict(api_name="/chat")
    # on the Flask side matches.
)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        server_name="0.0.0.0",
        # On HF Spaces the platform sets GRADIO_SERVER_PORT; locally we
        # default to 7860.
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
        show_api=True,
    )