"""
Chatterbox — HF Space chat app that writes code as downloadable files
and previews HTML live, similar to Claude's "artifacts".

Why the old version broke:
- `mistralai/Mistral-7B-Instruct-v0.1` was dropped from HF's free
  serverless inference tier.
- `client.text_generation(...)` with a hand-built `[INST]...[/INST]`
  prompt is the old (pre chat-template) way of calling models.
  `InferenceClient.chat_completion()` is the current, stable API and
  works across providers HF routes to (`provider="auto"`).

Requirements (requirements.txt):
    gradio>=4.44
    huggingface_hub>=0.26
"""

import os
import re
import tempfile
import zipfile
from pathlib import Path

import gradio as gr
from huggingface_hub import InferenceClient

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

# Reads a token from the Space's secrets if you add one (needed for higher
# rate limits / gated models). Works without it too, just more rate-limited.
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")

# Models known to work through HF's free/routed serverless chat-completion
# as of 2026. If one of these goes away, just remove it from the list —
# HF periodically retires models from the free tier.
MODEL_CHOICES = [
    "meta-llama/Llama-3.1-8B-Instruct",
    "Qwen/Qwen2.5-Coder-32B-Instruct",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "HuggingFaceH4/zephyr-7b-beta",
]

DEFAULT_SYSTEM_PROMPT = (
    "You are an expert software engineer, similar to Claude's coding "
    "assistant. When the user asks you to build or write something, answer "
    "with a short explanation and then one or more fenced code blocks.\n\n"
    "Every code block MUST start with a language tag, and its very first "
    "line must be a comment giving a filename, like this:\n\n"
    "```html\n<!-- filename: index.html -->\n<!doctype html>...\n```\n"
    "```python\n# filename: app.py\nprint('hi')\n```\n\n"
    "Always write complete, self-contained, runnable files (inline CSS/JS "
    "for HTML pages). Never truncate code for brevity."
)

CODE_BLOCK_RE = re.compile(r"```([a-zA-Z0-9_+-]*)\n(.*?)```", re.DOTALL)
FILENAME_RE = re.compile(r"^\s*(?:#|//|<!--)\s*filename:\s*([^\s>]+)", re.IGNORECASE)

EXT_BY_LANG = {
    "html": "html", "css": "css", "javascript": "js", "js": "js",
    "python": "py", "py": "py", "json": "json", "bash": "sh", "sh": "sh",
    "yaml": "yml", "yml": "yml", "typescript": "ts", "ts": "ts", "": "txt",
}

OUT_DIR = Path(tempfile.gettempdir()) / "chatterbox_outputs"
OUT_DIR.mkdir(exist_ok=True)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def extract_files(text: str, tag: str):
    """Pull fenced code blocks out of a model reply, write each to disk.

    Returns (list_of_written_file_paths, html_source_or_None).
    """
    files = []
    html_source = None
    for i, m in enumerate(CODE_BLOCK_RE.finditer(text)):
        lang = (m.group(1) or "").strip().lower()
        body = m.group(2)
        first_line, _, rest = body.partition("\n")
        fname_match = FILENAME_RE.match(first_line)
        if fname_match:
            filename = fname_match.group(1)
            body = rest
        else:
            ext = EXT_BY_LANG.get(lang, lang or "txt")
            filename = f"snippet_{i + 1}.{ext}"

        safe_name = re.sub(r"[^\w.\-]", "_", filename)
        path = OUT_DIR / f"{tag}_{safe_name}"
        path.write_text(body, encoding="utf-8")
        files.append(str(path))

        if lang in ("html", "") and safe_name.endswith("html") and html_source is None:
            html_source = body
        elif lang == "html" and html_source is None:
            html_source = body

    return files, html_source


def zip_files(file_paths, tag: str) -> str:
    """Bundle several generated files into a single .zip for one-click download.

    Files are stored under their original name (the shared `tag_` prefix used
    on disk to avoid collisions between turns is stripped inside the zip).
    """
    zip_path = OUT_DIR / f"{tag}_bundle.zip"
    prefix = f"{tag}_"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for fp in file_paths:
            arcname = Path(fp).name
            if arcname.startswith(prefix):
                arcname = arcname[len(prefix):]
            zf.write(fp, arcname=arcname)
    return str(zip_path)


def build_client(model: str) -> InferenceClient:
    kwargs = {"provider": "auto"}
    if HF_TOKEN:
        kwargs["token"] = HF_TOKEN
    return InferenceClient(model=model, **kwargs)


def make_preview_html(html_source: str | None) -> str:
    if not html_source:
        return "<div style='opacity:.6;padding:1rem;'>Нет HTML в последнем ответе — превью появится, когда модель вернёт html-блок.</div>"
    escaped = html_source.replace('"', "&quot;")
    return (
        f'<iframe style="width:100%;height:520px;border:1px solid #4443;'
        f'border-radius:10px;background:#fff;" srcdoc="{escaped}"></iframe>'
    )


# ---------------------------------------------------------------------------
# Core chat function (generator, streams tokens)
# ---------------------------------------------------------------------------

def respond(message, chat_history, model, system_prompt, temperature, max_tokens, top_p):
    chat_history = chat_history or []

    if not message or not message.strip():
        yield chat_history, "", gr.update(), gr.update()
        return

    system_prompt = system_prompt.strip() or DEFAULT_SYSTEM_PROMPT

    messages = [{"role": "system", "content": system_prompt}]
    messages.extend(chat_history)
    messages.append({"role": "user", "content": message})

    chat_history = chat_history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": ""},
    ]

    try:
        client = build_client(model)
        stream = client.chat_completion(
            messages=messages,
            max_tokens=int(max_tokens),
            temperature=float(max(temperature, 1e-2)),
            top_p=float(top_p),
            stream=True,
        )

        header = f"_Модель: `{model}`_\n\n"
        partial = header
        chat_history[-1]["content"] = partial
        for chunk in stream:
            if not chunk.choices:
                # some providers send trailing chunks with no `choices`
                # (e.g. usage-only chunks) — nothing to append, skip them.
                continue
            delta = chunk.choices[0].delta.content or ""
            if not delta:
                continue
            partial += delta
            chat_history[-1]["content"] = partial
            yield chat_history, "", gr.update(), gr.update()

        tag = f"turn{len(chat_history)}"
        files, html_source = extract_files(partial, tag)

        if len(files) > 1:
            # Multiple files: bundle them so it's one click, not one download
            # per file.
            download_value = zip_files(files, tag)
        elif files:
            download_value = files[0]
        else:
            download_value = None

        yield (
            chat_history,
            "",
            gr.update(value=download_value, visible=bool(download_value)),
            gr.update(value=make_preview_html(html_source)),
        )

    except Exception as e:
        err = str(e)
        if "api_key" in err.lower() or "log in" in err.lower():
            hint = (
                "\n\n👉 Похоже, не задан токен. Добавьте секрет `HF_TOKEN` в "
                "настройках Space (Settings → Variables and secrets), создав "
                "токен на huggingface.co/settings/tokens с правом *Make calls "
                "to Inference Providers*, затем перезапустите Space."
            )
        else:
            hint = (
                "\n\nЕсли модель недоступна на бесплатном инференсе — "
                "попробуйте другую из выпадающего списка сверху."
            )
        chat_history[-1]["content"] = f"⚠️ Ошибка запроса к модели `{model}`: {e}{hint}"
        yield chat_history, "", gr.update(), gr.update()


def clear_all():
    return [], None, make_preview_html(None)


DEFAULT_SETTINGS = {
    "model": MODEL_CHOICES[0],
    "system_prompt": DEFAULT_SYSTEM_PROMPT,
    "temperature": 0.7,
    "max_tokens": 2048,
    "top_p": 0.9,
}


def load_settings(saved):
    s = {**DEFAULT_SETTINGS, **(saved or {})}
    return (
        gr.update(value=s["model"]),
        gr.update(value=s["system_prompt"]),
        gr.update(value=s["temperature"]),
        gr.update(value=s["max_tokens"]),
        gr.update(value=s["top_p"]),
    )


def save_settings(model, system_prompt, temperature, max_tokens, top_p):
    return {
        "model": model,
        "system_prompt": system_prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "top_p": top_p,
    }


# Thin, rounded, theme-colored scrollbars instead of the bulky default OS
# ones, plus a layout that fits the viewport instead of forcing page scroll.
CUSTOM_CSS = """
:root {
  --sb-thumb: #8a8fb8;
  --sb-track: transparent;
}
@media (prefers-color-scheme: dark) {
  :root { --sb-thumb: #55597a; }
}
* {
  scrollbar-width: thin;
  scrollbar-color: var(--sb-thumb) var(--sb-track);
}
*::-webkit-scrollbar { width: 8px; height: 8px; }
*::-webkit-scrollbar-track { background: transparent; }
*::-webkit-scrollbar-thumb {
  background-color: var(--sb-thumb);
  border-radius: 8px;
  border: 2px solid transparent;
  background-clip: content-box;
}
*::-webkit-scrollbar-thumb:hover { background-color: #9fa3c7; }
"""

# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="Chatterbox", theme=gr.themes.Soft(), css=CUSTOM_CSS, fill_height=True) as demo:
    settings_state = gr.BrowserState(DEFAULT_SETTINGS)

    gr.Markdown("## 🤖 Chatterbox — чат + генерация файлов кода")

    if not HF_TOKEN:
        gr.Markdown(
            "⚠️ **Секрет `HF_TOKEN` не найден.** HF Inference Providers требуют "
            "токен даже для бесплатных запросов. Добавьте его в "
            "Settings → Variables and secrets вашего Space (токен создаётся на "
            "huggingface.co/settings/tokens с правом *Make calls to Inference "
            "Providers*), затем перезапустите Space."
        )

    with gr.Row(equal_height=False):
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                type="messages",
                show_label=False,
                show_copy_button=True,
                scale=1,
            )
            msg = gr.Textbox(
                placeholder="Опишите, что построить, например: «сделай лендинг с формой подписки»",
                show_label=False,
                lines=2,
            )
            with gr.Row():
                send_btn = gr.Button("Отправить", variant="primary")
                clear_btn = gr.Button("Очистить")

        with gr.Column(scale=2):
            with gr.Tab("📄 Файлы"):
                files_out = gr.File(
                    label="Сгенерированные файлы (несколько — упаковываются в .zip)",
                    file_count="single",
                )
            with gr.Tab("👁 Превью (HTML)"):
                preview = gr.HTML(value=make_preview_html(None))

            with gr.Accordion("Настройки модели", open=False):
                model_dd = gr.Dropdown(
                    MODEL_CHOICES, value=MODEL_CHOICES[0], label="Модель",
                    allow_custom_value=True,
                )
                system_box = gr.Textbox(
                    label="System prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6,
                )
                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
                max_tokens = gr.Slider(128, 8192, value=2048, step=64, label="Max new tokens")
                top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p")

    settings_widgets = [model_dd, system_box, temperature, max_tokens, top_p]

    demo.load(load_settings, inputs=[settings_state], outputs=settings_widgets)
    for widget in settings_widgets:
        widget.change(save_settings, inputs=settings_widgets, outputs=[settings_state])

    inputs = [msg, chatbot, model_dd, system_box, temperature, max_tokens, top_p]
    outputs = [chatbot, msg, files_out, preview]

    msg.submit(respond, inputs, outputs)
    send_btn.click(respond, inputs, outputs)
    clear_btn.click(clear_all, None, [chatbot, files_out, preview])

if __name__ == "__main__":
    demo.launch(show_api=False)