"""Another TULLUS HF Space Original — CodeParrot Mini Lab.

A small Gradio coding-completion interface around the published
TULLUS/codeparrot-small-multi checkpoint. This is inference only.
"""

from __future__ import annotations

import logging
import threading
from typing import Any

import gradio as gr
import spaces
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    StoppingCriteria,
    StoppingCriteriaList,
)


# ---------------------------------------------------------------------------
# Model/runtime setup. ZeroGPU's documented pattern places the model on CUDA
# at module import time; only calls to generate() are inside @spaces.GPU.
# ---------------------------------------------------------------------------

MODEL_ID = "TULLUS/codeparrot-small-multi-F32"
CHECKPOINT_DTYPE = torch.float32  # Preserve the repository config.json dtype: F32.
MAX_NEW_TOKENS_UI = 384
MAX_CONTEXT_PROMPT_CHARS = 40_000  # Secondary guard before tokenization.

logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger("tullus_codeparrot_mini_lab")

TOKENIZER = None
MODEL = None
MODEL_DEVICE = "unavailable"
MODEL_LOAD_ERROR: str | None = None
CONTEXT_TOKENS = 1024  # Safe display fallback; replaced from loaded config.
PARAMETER_COUNT = 0
RUNTIME_DTYPE = "unavailable"
HAS_CHAT_TEMPLATE = False

# A stop event per Gradio session lets the Stop button request cooperative
# cancellation during model.generate(), without one visitor stopping another.
_STOP_EVENTS: dict[str, threading.Event] = {}
_STOP_EVENTS_LOCK = threading.Lock()

LANGUAGES = [
    "Auto",
    "Python",
    "JavaScript",
    "PHP",
    "Java",
    "C#",
    "C++",
    "Go",
    "Ruby",
    "TypeScript",
]
COMMENT_PREFIX = {
    "Python": "#",
    "Ruby": "#",
    "JavaScript": "//",
    "PHP": "//",
    "Java": "//",
    "C#": "//",
    "C++": "//",
    "Go": "//",
    "TypeScript": "//",
}


class StopOnEvent(StoppingCriteria):
    """Stop generation once the current session's Stop button is pressed."""

    def __init__(self, event: threading.Event):
        self.event = event

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs: Any) -> bool:
        return self.event.is_set()


def _get_context_capacity(config: Any) -> int:
    """Read the positional limit from the loaded architecture config."""
    for attribute in ("max_position_embeddings", "n_positions", "n_ctx"):
        value = getattr(config, attribute, None)
        if isinstance(value, int) and 0 < value < 1_000_000:
            return value
    return 1024


def _load_model_once() -> None:
    """Initialize shared model/tokenizer once per Space process."""
    global TOKENIZER, MODEL, MODEL_DEVICE, MODEL_LOAD_ERROR
    global CONTEXT_TOKENS, PARAMETER_COUNT, RUNTIME_DTYPE, HAS_CHAT_TEMPLATE

    try:
        LOGGER.info("Loading tokenizer from %s", MODEL_ID)
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)

        LOGGER.info("Loading published Safetensors checkpoint from %s", MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype=CHECKPOINT_DTYPE,
            use_safetensors=True,
        )

        # Hugging Face ZeroGPU emulates CUDA at process initialization. The CPU
        # branch keeps import/startup diagnosable in non-ZeroGPU environments;
        # intended hosted deployment is a Space configured for ZeroGPU.
        if torch.cuda.is_available():
            model = model.to("cuda")
            MODEL_DEVICE = "cuda"
        else:
            model = model.to("cpu")
            MODEL_DEVICE = "cpu"

        model.eval()
        
        # Diagnostic only: compare tokenizer and model special-token settings.
        generation_config = getattr(model, "generation_config", None)

        LOGGER.info(
            "Special-token audit: "
            "tokenizer(eos=%r, eos_id=%r, bos_id=%r, pad_id=%r); "
            "model.config(eos_id=%r, bos_id=%r, pad_id=%r); "
            "model.generation_config(eos_id=%r, pad_id=%r); "
            "vocab(tokenizer=%r, tokenizer_len=%r, model=%r)",
            tokenizer.eos_token,
            tokenizer.eos_token_id,
            tokenizer.bos_token_id,
            tokenizer.pad_token_id,
            getattr(model.config, "eos_token_id", None),
            getattr(model.config, "bos_token_id", None),
            getattr(model.config, "pad_token_id", None),
            getattr(generation_config, "eos_token_id", None),
            getattr(generation_config, "pad_token_id", None),
            tokenizer.vocab_size,
            len(tokenizer),
            getattr(model.config, "vocab_size", None),
        )

        TOKENIZER = tokenizer
        MODEL = model
        TOKENIZER = tokenizer
        MODEL = model
        CONTEXT_TOKENS = _get_context_capacity(model.config)
        PARAMETER_COUNT = sum(parameter.numel() for parameter in model.parameters())
        RUNTIME_DTYPE = str(next(model.parameters()).dtype).replace("torch.", "")
        HAS_CHAT_TEMPLATE = bool(getattr(tokenizer, "chat_template", None))
        LOGGER.info(
            "Model ready: parameters=%s context=%s dtype=%s device=%s chat_template=%s",
            f"{PARAMETER_COUNT:,}",
            CONTEXT_TOKENS,
            RUNTIME_DTYPE,
            MODEL_DEVICE,
            HAS_CHAT_TEMPLATE,
        )
    except Exception as exc:  # Keep the interface alive and show a useful status.
        MODEL_LOAD_ERROR = f"{type(exc).__name__}: {exc}"
        LOGGER.exception("Model initialization failed")
        TOKENIZER = None
        MODEL = None
        MODEL_DEVICE = "unavailable"
        RUNTIME_DTYPE = "unavailable"


_load_model_once()


def _startup_status() -> str:
    if MODEL is None or TOKENIZER is None:
        return (
            "🟠 **MODEL OFFLINE** — the interface loaded, but model initialization failed. "
            "Check the Space logs for the specific loading error."
        )
    if MODEL_DEVICE == "cuda":
        return "🟢 **SYSTEMS ONLINE** · model loaded · ZeroGPU inference ready"
    return (
        "🟡 **MODEL LOADED ON CPU** · the app did not detect CUDA at startup. "
        "For hosted GPU inference, set this Space's hardware to ZeroGPU."
    )


def _as_text(content: Any) -> str:
    """Convert Gradio message content to ordinary text without role templating."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        chunks = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                chunks.append(str(item.get("text", "")))
        return "\n".join(chunks)
    return str(content or "")


def _recent_completed_pairs(history: list[dict[str, Any]]) -> list[tuple[str, str]]:
    """Extract only completed user/assistant pairs from recent UI history."""
    pairs: list[tuple[str, str]] = []
    pending_user: str | None = None
    for message in history or []:
        role = message.get("role")
        body = _as_text(message.get("content", "")).strip("\n")
        if role == "user":
            pending_user = body
        elif role == "assistant" and pending_user is not None:
            # Don't inject prior application error messages as model context.
            if body and not body.startswith("⚠️"):
                pairs.append((pending_user, body))
            pending_user = None
    return pairs[-2:]


def _language_prefix(language: str) -> str:
    if language == "Auto" or language not in COMMENT_PREFIX:
        return ""
    marker = COMMENT_PREFIX[language]
    return f"{marker} Target language: {language}\n"


def _make_prompt(
    current_message: str,
    history: list[dict[str, Any]],
    language: str,
    use_recent_context: bool,
) -> tuple[str, bool]:
    """Build a transparent completion prompt; never apply a chat template.

    Optional history is represented as commented request headings and raw prior
    completions. Oldest context is removed first if required by the real limit.
    The current prompt is never silently truncated.
    """
    language_prefix = _language_prefix(language)
    base_prompt = language_prefix + current_message
    base_ids = TOKENIZER(base_prompt, add_special_tokens=True, truncation=False)["input_ids"]
    if len(base_ids) >= CONTEXT_TOKENS:
        raise ValueError(
            f"Your current prompt uses {len(base_ids)} tokens, but this model has a "
            f"{CONTEXT_TOKENS}-token context. Shorten the current prompt to leave room "
            "for generated tokens, then try again. The app did not truncate your code."
        )

    pairs = _recent_completed_pairs(history) if use_recent_context else []
    marker = COMMENT_PREFIX.get(language, "#")
    trimmed = False

    def assemble(selected_pairs: list[tuple[str, str]]) -> str:
        if not selected_pairs:
            return base_prompt
        context_sections: list[str] = [f"{marker} Earlier coding context (reference only):"]
        for request_text, completion_text in selected_pairs:
            context_sections.append(f"{marker} Earlier request:")
            context_sections.extend(
                f"{marker} {line}" if line else marker
                for line in request_text.splitlines()
            )
            context_sections.append(f"{marker} Earlier completion follows:")
            context_sections.append(completion_text)
            context_sections.append(f"\n{marker} End earlier completion.")
        context_sections.append(f"{marker} Current request follows:")
        return "\n".join(context_sections) + "\n" + base_prompt

    prompt = assemble(pairs)
    while pairs:
        count = len(TOKENIZER(prompt, add_special_tokens=True, truncation=False)["input_ids"])
        if count < CONTEXT_TOKENS:
            break
        pairs.pop(0)
        trimmed = True
        prompt = assemble(pairs)

    prompt_count = len(TOKENIZER(prompt, add_special_tokens=True, truncation=False)["input_ids"])
    if prompt_count >= CONTEXT_TOKENS:
        # This should only happen if the current prompt itself fills the window;
        # retain an explicit guard if tokenizer behavior changes.
        raise ValueError(
            f"The assembled prompt uses {prompt_count} tokens; context capacity is "
            f"{CONTEXT_TOKENS}. Shorten the current prompt."
        )
    return prompt, trimmed


def _error_result(history: list[dict[str, Any]], message: str, status: str):
    updated = list(history or [])
    updated.append({"role": "user", "content": message})
    updated.append({"role": "assistant", "content": f"⚠️ {status}"})
    return updated, "", f"🟠 **INPUT / INFERENCE ISSUE** · {status}"


@spaces.GPU(duration=60)
def respond(
    message: str,
    history: list[dict[str, Any]] | None,
    language: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    sample: bool,
    use_recent_context: bool,
    request: gr.Request,
    progress=gr.Progress(),
):
    """Generate a completion inside a ZeroGPU-decorated function."""
    chat_history = list(history or [])
    message = message if isinstance(message, str) else str(message or "")

    if not message.strip():
        return chat_history, "", "🟡 **READY** · Enter a code prefix or a coding request to begin."
    if len(message) > MAX_CONTEXT_PROMPT_CHARS:
        return _error_result(
            chat_history,
            message,
            f"The prompt is longer than the {MAX_CONTEXT_PROMPT_CHARS:,}-character safety limit. "
            "Please shorten it before retrying.",
        )
    if MODEL is None or TOKENIZER is None:
        return _error_result(
            chat_history,
            message,
            "The model did not initialize. Check the Space logs, then restart the Space.",
        )

    session_id = getattr(request, "session_hash", None) or "local-session"
    stop_event = threading.Event()
    with _STOP_EVENTS_LOCK:
        _STOP_EVENTS[session_id] = stop_event

    try:
        progress(0.05, desc="Preparing the coding prompt")
        prompt, prior_context_trimmed = _make_prompt(
            message,
            chat_history,
            language,
            bool(use_recent_context),
        )
        tokenized = TOKENIZER(prompt, return_tensors="pt", add_special_tokens=True, truncation=False)
        input_length = int(tokenized["input_ids"].shape[-1])
        if input_length >= CONTEXT_TOKENS:
            raise ValueError(
                f"Prompt length is {input_length} tokens; context capacity is {CONTEXT_TOKENS}. "
                "Shorten the current prompt and retry."
            )

        requested_new_tokens = int(max_new_tokens)
        requested_new_tokens = max(1, min(requested_new_tokens, MAX_NEW_TOKENS_UI))
        available_new_tokens = CONTEXT_TOKENS - input_length
        actual_new_tokens = min(requested_new_tokens, available_new_tokens)
        if actual_new_tokens < 1:
            raise ValueError(
                f"No room remains to generate tokens: prompt={input_length}, "
                f"context={CONTEXT_TOKENS}. Shorten the prompt."
            )

        # Inputs are transferred to the device chosen at module initialization.
        model_inputs = {key: value.to(MODEL_DEVICE) for key, value in tokenized.items()}
        generation_kwargs: dict[str, Any] = {
            "max_new_tokens": actual_new_tokens,
            "repetition_penalty": float(repetition_penalty),
            "do_sample": bool(sample),
            "stopping_criteria": StoppingCriteriaList([StopOnEvent(stop_event)]),
            "use_cache": True,
        }
        if sample:
            generation_kwargs["temperature"] = float(temperature)
            generation_kwargs["top_p"] = float(top_p)

        eos_id = TOKENIZER.eos_token_id
        if eos_id is None:
            eos_id = getattr(MODEL.config, "eos_token_id", None)
        pad_id = TOKENIZER.pad_token_id if TOKENIZER.pad_token_id is not None else eos_id
        if eos_id is not None:
            generation_kwargs["eos_token_id"] = eos_id
        if pad_id is not None:
            generation_kwargs["pad_token_id"] = pad_id

        progress(0.25, desc="Requesting ZeroGPU and generating")
        with torch.inference_mode():
            generated = MODEL.generate(**model_inputs, **generation_kwargs)

        # Slice by the tokenized input length, never by fragile string removal.
        new_token_ids = generated[0, input_length:]
        generated_token_count = int(new_token_ids.numel())
        
        completion = TOKENIZER.decode(
            new_token_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )
        
        was_stopped = stop_event.is_set()
        
        # Determine whether generation ended on EOS or reached its token ceiling.
        eos_ids = (
            {int(x) for x in eos_id}
            if isinstance(eos_id, (list, tuple, set))
            else ({int(eos_id)} if eos_id is not None else set())
        )
        
        # ended_with_eos = (
            # generated_token_count > 0
            # and int(new_token_ids[-1]) in eos_ids
        # )
        
        # hit_token_limit = (
            # generated_token_count >= actual_new_tokens
            # and not ended_with_eos
            # and not was_stopped
        # )


        raw_new_ids = [int(token_id) for token_id in new_token_ids.tolist()]

        ended_with_eos = (
            bool(raw_new_ids)
            and raw_new_ids[-1] in eos_ids
        )

        hit_token_limit = (
            generated_token_count >= actual_new_tokens
            and not ended_with_eos
            and not was_stopped
        )

        eos_positions = [
            index
            for index, token_id in enumerate(raw_new_ids)
            if token_id in eos_ids
        ]

        generation_config = getattr(MODEL, "generation_config", None)

        LOGGER.info(
            "Generation-stop audit: "
            "effective_eos_id=%r, "
            "tokenizer_eos_id=%r, "
            "model_config_eos_id=%r, "
            "generation_config_eos_id=%r, "
            "generated=%d/%d, "
            "tail_token_ids=%s, "
            "eos_positions=%s, "
            "ended_with_eos=%s, "
            "stop_event=%s, "
            "hit_token_limit=%s",
            eos_id,
            TOKENIZER.eos_token_id,
            getattr(MODEL.config, "eos_token_id", None),
            getattr(generation_config, "eos_token_id", None),
            generated_token_count,
            actual_new_tokens,
            raw_new_ids[-8:],
            eos_positions,
            ended_with_eos,
            was_stopped,
            hit_token_limit,
        )


        if not completion:
            completion = (
                "(Generation stopped before visible text was produced.)"
                if was_stopped
                else "(The model ended without producing visible text.)"
            )

        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": completion})

        notes = []
        if prior_context_trimmed:
            notes.append("older context was dropped to fit the model window")
        if actual_new_tokens < requested_new_tokens:
            notes.append(f"output capped at {actual_new_tokens} tokens by the context limit")
        if was_stopped:
            notes.append("stop signal received")
        note_text = " · " + " · ".join(notes) if notes else ""
        #state_word = "STOPPED" if was_stopped else "COMPLETE"
        if was_stopped:
            state_word = "STOPPED"
        elif hit_token_limit:
            state_word = "TOKEN LIMIT REACHED"
        else:
            state_word = "COMPLETE"
        
        notes.append(f"generated {generated_token_count} tokens")
        
        if hit_token_limit:
            notes.append("generation reached its token ceiling; output may be truncated")
        elif ended_with_eos:
            notes.append("model emitted EOS")
        else:
            notes.append("ended before the token ceiling")
        progress(1.0, desc="Completion ready")
        return (
            chat_history,
            "",
            f"🟢 **{state_word}** · input {input_length}/{CONTEXT_TOKENS} tokens · "
            f"new-token limit {actual_new_tokens} · runtime {RUNTIME_DTYPE}{note_text}",
        )
    except ValueError as exc:
        LOGGER.info("Prompt rejected: %s", exc)
        return _error_result(chat_history, message, str(exc))
    except Exception:
        LOGGER.exception("Generation failed")
        return _error_result(
            chat_history,
            message,
            "Generation failed unexpectedly. Please try again; if it repeats, check the Space logs.",
        )
    finally:
        with _STOP_EVENTS_LOCK:
            if _STOP_EVENTS.get(session_id) is stop_event:
                _STOP_EVENTS.pop(session_id, None)


def stop_generation(request: gr.Request) -> str:
    """Set this browser session's event; StoppingCriteria checks it per token."""
    session_id = getattr(request, "session_hash", None) or "local-session"
    with _STOP_EVENTS_LOCK:
        event = _STOP_EVENTS.get(session_id)
    if event is None:
        return "🟡 **NO ACTIVE GENERATION** · nothing to stop right now."
    event.set()
    return "🟠 **STOP REQUESTED** · the current generation will stop at its next token check."


def clear_conversation():
    return [], "", _startup_status()


# ---------------------------------------------------------------------------
# Space Station UI
# ---------------------------------------------------------------------------

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.cyan,
    secondary_hue=gr.themes.colors.orange,
    neutral_hue=gr.themes.colors.slate,
    font=["Inter", "ui-sans-serif", "sans-serif"],
    font_mono=["ui-monospace", "SFMono-Regular", "Menlo", "Consolas", "monospace"],
).set(
    body_background_fill="#070b12",
    body_text_color="#e4edf6",
    block_background_fill="#101923",
    block_border_color="#243545",
    block_label_text_color="#a9c2d2",
    input_background_fill="#0b111a",
    input_border_color="#294152",
    button_primary_background_fill="#0c7d8d",
    button_primary_background_fill_hover="#1195a6",
    button_primary_text_color="#f4fdff",
    button_secondary_background_fill="#192633",
    button_secondary_text_color="#d4e3ed",
)

CSS = r"""
:root {
  --space-bg: #070b12;
  --space-panel: rgba(12, 20, 30, 0.94);
  --space-line: #234151;
  --space-cyan: #69e4f1;
  --space-muted: #8ca5b6;
  --space-amber: #f0b86a;
}
html, body, .gradio-container {
  background-color: var(--space-bg) !important;
  color: #e4edf6 !important;
}
body {
  background-image:
    radial-gradient(1px 1px at 12% 16%, rgba(141, 224, 239, .45) 50%, transparent 100%),
    radial-gradient(1px 1px at 83% 22%, rgba(240, 184, 106, .45) 50%, transparent 100%),
    radial-gradient(1px 1px at 72% 79%, rgba(141, 224, 239, .28) 50%, transparent 100%),
    radial-gradient(1px 1px at 27% 68%, rgba(141, 224, 239, .25) 50%, transparent 100%),
    radial-gradient(ellipse at 50% -12%, rgba(16, 77, 105, .26), transparent 54%);
}
.gradio-container { max-width: 1440px !important; }
#station { max-width: 1180px; margin: 0 auto; }
#station-header {
  position: relative; overflow: hidden; border: 1px solid #2d5262;
  border-radius: 18px; padding: 25px 28px 22px;
  background: linear-gradient(115deg, rgba(14, 32, 45, .98), rgba(10, 17, 27, .97) 63%, rgba(31, 27, 24, .93));
  box-shadow: 0 14px 50px rgba(0,0,0,.25), inset 0 1px rgba(255,255,255,.025);
  margin-bottom: 12px;
}
#station-header:after {
  content: ""; position: absolute; right: -38px; top: -76px; width: 270px; height: 270px;
  border: 1px solid rgba(105, 228, 241, .2); border-radius: 50%;
  box-shadow: 0 0 0 18px rgba(105,228,241,.025), 0 0 0 44px rgba(105,228,241,.018);
  pointer-events: none;
}
#station-kicker { color: var(--space-amber); letter-spacing: .2em; font: 600 11px ui-monospace, monospace; text-transform: uppercase; }
#station-title { margin: 7px 0 4px; font-size: clamp(25px, 4vw, 39px); line-height: 1.08; font-weight: 700; letter-spacing: .035em; color: #f1f8fc; }
#station-title span { color: var(--space-cyan); }
#station-subtitle { color: #a9c3d0; font-size: 14px; margin-top: 8px; }
#telemetry { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 18px; position: relative; z-index: 1; }
.telemetry-chip { border: 1px solid #284c5a; background: rgba(5,13,21,.58); color: #a7c6d2; border-radius: 5px; padding: 6px 9px; font: 11px ui-monospace, monospace; }
.telemetry-chip b { color: var(--space-cyan); font-weight: 500; }
#system-status { border: 1px solid #284858; background: rgba(9, 18, 28, .9); border-radius: 9px; padding: 10px 14px; font-size: 13px; margin: 10px 0 14px; }
#conversation { border: 1px solid #284252; border-radius: 14px; overflow: hidden; background: rgba(9, 15, 23, .96); }
#conversation .wrap { border: 0 !important; }
#conversation .message { font-size: 14px; line-height: 1.55; }
#conversation pre, #conversation code { font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace !important; }
#conversation pre { border: 1px solid #243d4b !important; border-radius: 8px !important; background: #080e16 !important; }
#conversation .message.user { border-color: #25505d !important; }
#prompt-box textarea { font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace !important; line-height: 1.55 !important; }
#prompt-box textarea:focus { border-color: var(--space-cyan) !important; box-shadow: 0 0 0 1px rgba(105,228,241,.16) !important; }
#send-button { min-height: 47px; font-weight: 650; letter-spacing: .035em; }
#stop-button { min-height: 47px; border-color: #725232 !important; color: #f0bf7e !important; }
#clear-button { min-height: 47px; }
#lab-panel { border: 1px solid #263d4b; border-radius: 12px; background: rgba(10,17,26,.85); padding: 8px 12px; }
#model-card { border-top: 1px solid #263b49; margin-top: 14px; padding-top: 14px; color: var(--space-muted); font-size: 12px; line-height: 1.7; }
#model-card strong { color: #d1e7ef; }
#model-card a { color: var(--space-cyan); }
#station .gr-accordion { border-color: #284252 !important; background: rgba(12,20,30,.8) !important; }
#station footer { opacity: .65; }
@media (max-width: 700px) {
  #station-header { padding: 20px 17px; }
  #station-title { font-size: 27px; }
  #station-header:after { right: -150px; opacity: .5; }
  #conversation { min-height: 390px; }
}
"""

parameter_display = f"{PARAMETER_COUNT / 1_000_000:.1f}M" if PARAMETER_COUNT else "~111M (model-page metadata)"
model_info_html = f"""
<div id="model-card">
  <strong>MODEL BAY / REFERENCE CARD</strong><br>
  Source checkpoint: <a href="https://huggingface.co/{MODEL_ID}" target="_blank" rel="noreferrer">{MODEL_ID}</a><br>
  Architecture: GPT-2 causal language model · Parameter count: {parameter_display}<br>
  Published artifact: Safetensors, listed as F32 on the model page · Runtime dtype: {RUNTIME_DTYPE}<br>
  Context capacity: {CONTEXT_TOKENS:,} tokens from the loaded configuration · Device at startup: {MODEL_DEVICE}<br>
  No chat template is used. Detected tokenizer chat-template metadata: {"present (not applied)" if HAS_CHAT_TEMPLATE else "not present"}.<br>
  Code completion is the primary behavior; instruction-following and correctness are not guaranteed.
  Model metadata identifies Apache-2.0 licensing and the CodeParrot code / HumanEval datasets.
</div>
"""

with gr.Blocks(
    title="TULLUS Space Station · CodeParrot Mini Lab",
    analytics_enabled=False,
) as demo:
    with gr.Column(elem_id="station"):
        gr.HTML(
            """
            <header id="station-header">
              <div id="station-kicker">TULLUS / ORBITAL RESEARCH DIVISION · LAB 09</div>
              <div id="station-title">TULLUS SPACE <span>STATION</span></div>
              <div style="font:600 12px ui-monospace,monospace;letter-spacing:.17em;color:#c3dce5">CODEPARROT MINI LAB</div>
              <div id="station-subtitle">Code generation from orbit. Laboratory systems online. 🛸</div>
              <div id="telemetry">
                <div class="telemetry-chip"><b>SYS</b> / CODE COMPLETION</div>
                <div class="telemetry-chip"><b>CORE</b> / GPT-2 · 111M CLASS</div>
                <div class="telemetry-chip"><b>LINK</b> / HUGGING FACE</div>
                <div class="telemetry-chip"><b>AMBIENCE</b> / SPACE JAZZ ▪︎ LOW</div>
              </div>
            </header>
            """
        )
        status = gr.Markdown(_startup_status(), elem_id="system-status")

        gr.Markdown("### ◉ Mission console  /  coding channel")
        chatbot = gr.Chatbot(
            label="Conversation log",
            height=450,
            elem_id="conversation",
            buttons=["copy", "copy_all"],
            render_markdown=True,
            placeholder=(
                "<div style='text-align:center;padding:30px 10px;color:#8ca5b6'>"
                "<div style='font:12px monospace;letter-spacing:.15em;color:#69e4f1'>CHANNEL OPEN</div>"
                "<h3 style='color:#e4edf6'>Awaiting your first coding task</h3>"
                "Drop in a function prefix, existing code, or a plain-language request."
                "</div>"
            ),
        )

        with gr.Row(equal_height=True):
            language = gr.Dropdown(
                choices=LANGUAGES,
                value="Auto",
                label="Language hint",
                info="Auto passes your prompt unchanged. A language selection adds a visible plain-text comment prefix (not a special token).",
                scale=1,
            )
        prompt_box = gr.Textbox(
            label="Prompt / code prefix",
            placeholder="Paste code, start a function, or describe the snippet you want to explore…",
            lines=5,
            max_lines=12,
            elem_id="prompt-box",
            autofocus=True,
        )
        with gr.Row():
            send_button = gr.Button("↗  TRANSMIT PROMPT", variant="primary", elem_id="send-button", scale=3)
            stop_button = gr.Button("■  STOP GENERATION", variant="secondary", elem_id="stop-button", scale=1)
            clear_button = gr.Button("↺  NEW TASK", variant="secondary", elem_id="clear-button", scale=1)

        with gr.Accordion("Generation controls / lab instruments", open=False, elem_id="lab-panel"):
            with gr.Row():
                max_new_tokens = gr.Slider(
                    minimum=16, maximum=MAX_NEW_TOKENS_UI, value=128, step=16,
                    label="Maximum new tokens", info="The context window may reduce this for long prompts.",
                )
                sample = gr.Checkbox(value=True, label="Use sampling", info="Turn off for greedy decoding.")
            with gr.Row():
                temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.2, step=0.05, label="Temperature", info="Lower is less varied.")
                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p", info="Sampling nucleus size.")
                repetition_penalty = gr.Slider(minimum=1.0, maximum=1.5, value=1.1, step=0.05, label="Repetition penalty", info="Higher discourages repeated text.")
            use_recent_context = gr.Checkbox(
                value=False,
                label="Include recent conversation context",
                info="Off by default. When enabled, up to two completed exchanges may be included as commented reference context; older context is dropped first to fit the model window.",
            )

        gr.Examples(
            examples=[
                ["def fibonacci(n):", "Python"],
                ["function reverseString(str) {", "JavaScript"],
                ["public static int factorial(int n) {", "Java"],
                ["int calculate_sum(", "C++"],
                ["function formatUser(user: User) {", "TypeScript"],
                [
                    'def is_palindrome(s):\n    """Return True if s reads the same forwards and backwards, ignoring capitalization and punctuation."""\n    ',
                    "Python",
                ],
            ],
            inputs=[prompt_box, language],
            label="Docking bay / starter prompts",
        )

        gr.HTML(model_info_html)
        gr.Markdown(
            "*Experimental compact code-generation model, not a frontier reasoning chatbot. "
            "Review, test, and validate generated code before use.*"
        )

    generation_inputs = [
        prompt_box,
        chatbot,
        language,
        max_new_tokens,
        temperature,
        top_p,
        repetition_penalty,
        sample,
        use_recent_context,
    ]
    generation_outputs = [chatbot, prompt_box, status]
    send_event = prompt_box.submit(
        fn=respond,
        inputs=generation_inputs,
        outputs=generation_outputs,
        concurrency_id="model-generation",
        concurrency_limit=1,
    )
    send_button.click(
        fn=respond,
        inputs=generation_inputs,
        outputs=generation_outputs,
        concurrency_id="model-generation",
        concurrency_limit=1,
    )
    stop_button.click(fn=stop_generation, inputs=None, outputs=status, concurrency_limit=4)
    clear_button.click(fn=clear_conversation, inputs=None, outputs=[chatbot, prompt_box, status])

# Keep a small amount of concurrency available so the Stop event can be handled
# while the model-generation event is running. GPU generation itself is limited
# to one concurrent request above.
demo.queue(default_concurrency_limit=4)

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS)