"""Northwind Ops POC chat — Hugging Face Space (app.py)."""

from __future__ import annotations

import os
from functools import lru_cache
from typing import Any

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    import spaces
except ImportError:  # CPU Space / local without ZeroGPU

    class _SpacesShim:
        @staticmethod
        def GPU(fn=None, **_kwargs):
            if fn is None:
                return lambda f: f
            return fn

    spaces = _SpacesShim()  # type: ignore

MODEL_ID = os.environ.get("NORTHWIND_MODEL_ID", "UnaverageTech411/northwind-ops")
SYSTEM = os.environ.get(
    "NORTHWIND_SYSTEM",
    (
        "You are Northwind Ops Assistant for Northwind Traders — the showcase company model "
        "from the Arriella custom-model factory. For Northwind IT, HR, and finance procedures, "
        "answer from documented company knowledge and name the official tool (ServiceNow, Concur, Workday). "
        "Give crisp step paths (for example ServiceNow → IT → VPN Token Reset with Northwind SSO). "
        "For ordinary world knowledge, math, spelling, and general chat, answer normally and concisely. "
        "Only say a fact is not documented when it is a Northwind-internal detail missing from training "
        "(salary bands, PTO balances, unpublished policies). Never say not documented for VPN, Concur, "
        "or Workday procedures that are in your training. Never refuse common general-knowledge questions."
    ),
)

# (button label, prompt) — training-data / gate probes
EXAMPLE_PROMPTS: list[tuple[str, str]] = [
    ("VPN reset", "How do I reset my VPN token at Northwind in ServiceNow?"),
    ("VPN by email?", "Can I reset my Northwind VPN token by emailing IT?"),
    ("Expense / Concur", "How do I submit an expense report in Concur?"),
    ("PTO / Workday", "Where do I check PTO in Workday?"),
    ("Salary band", "What is my exact salary band code?"),
    ("SaaS access", "Where do I request SaaS access?"),
    ("Who is CTO?", "Who is the Northwind CTO?"),
    ("Capital of France", "What is the capital of France?"),
]

CSS = """
.gradio-container {
  max-width: 920px !important;
  margin: auto;
  font-family: "IBM Plex Sans", "Segoe UI", sans-serif !important;
}
footer { display: none !important; }
#poc-title h1 {
  font-family: "Syne", "Arial Narrow", sans-serif !important;
  letter-spacing: -0.02em;
  margin-bottom: 0.25rem !important;
}
#poc-lede {
  color: #9aabbd !important;
  font-size: 0.98rem !important;
  line-height: 1.5 !important;
  margin-top: 0 !important;
}
#chatbot { border: 2px solid rgba(244,247,251,0.22) !important; }
#probe-row button, #probe-row-2 button {
  font-size: 0.82rem !important;
}
"""


@lru_cache(maxsize=1)
def _load() -> tuple[Any, Any]:
    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype,
        device_map="auto" if torch.cuda.is_available() else None,
        trust_remote_code=True,
    )
    if not torch.cuda.is_available():
        model = model.to("cpu")
    model.eval()
    return tok, model


def _history_to_messages(history: list[dict[str, str]]) -> list[dict[str, str]]:
    messages = [{"role": "system", "content": SYSTEM}]
    for turn in history or []:
        role = turn.get("role")
        content = (turn.get("content") or "").strip()
        if role in {"user", "assistant"} and content:
            messages.append({"role": role, "content": content})
    return messages


@spaces.GPU(duration=90)
def _generate(messages: list[dict[str, str]], max_new_tokens: int = 220) -> str:
    tok, model = _load()
    prompt = tok.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tok(prompt, return_tensors="pt")
    device = next(model.parameters()).device
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tok.eos_token_id,
        )
    gen = out[0, inputs["input_ids"].shape[-1] :]
    return tok.decode(gen, skip_special_tokens=True).strip()


def respond(message: str, history: list[dict[str, str]]):
    text = (message or "").strip()
    if not text:
        return history, ""
    history = list(history or [])
    history.append({"role": "user", "content": text})
    try:
        reply = _generate(_history_to_messages(history))
    except Exception as exc:  # noqa: BLE001
        reply = f"(Demo error: {exc})"
    history.append({"role": "assistant", "content": reply or "(empty reply)"})
    return history, ""


def use_example(prompt: str, history: list[dict[str, str]]):
    return respond(prompt, history)


def clear_chat():
    return [], ""


try:
    _load()
except Exception:
    pass

theme = gr.themes.Soft(
    primary_hue=gr.themes.colors.amber,
    secondary_hue=gr.themes.colors.emerald,
    neutral_hue=gr.themes.colors.slate,
).set(
    body_background_fill="#06090d",
    body_background_fill_dark="#06090d",
    block_background_fill="#111922",
    block_background_fill_dark="#111922",
    body_text_color="#f4f7fb",
    body_text_color_dark="#f4f7fb",
    border_color_primary="rgba(244,247,251,0.22)",
    border_color_primary_dark="rgba(244,247,251,0.22)",
)

with gr.Blocks(title="Northwind Ops Assistant", theme=theme, css=CSS) as demo:
    with gr.Column(elem_id="poc-title"):
        gr.Markdown("# Northwind Ops — try the POC")
        gr.Markdown(
            "1B company model: ServiceNow / Concur / Workday procedures, "
            "Pile-hardened, Heretic ×2, openness-gated. "
            "Tap a probe from the training set, or type your own. Clear anytime.",
            elem_id="poc-lede",
        )

    chatbot = gr.Chatbot(
        label="Northwind Ops",
        height=440,
        type="messages",
        elem_id="chatbot",
        value=[
            {
                "role": "assistant",
                "content": (
                    "Ask a company procedure and a general fact — that contrast is the demo. "
                    "Try **VPN reset** or **Capital of France** below."
                ),
            }
        ],
    )

    with gr.Row():
        msg = gr.Textbox(
            label="Message",
            show_label=False,
            placeholder="How do I reset my VPN token at Northwind in ServiceNow?",
            scale=5,
            autofocus=True,
            container=False,
        )
        send = gr.Button("Send", variant="primary", scale=1)
        clear_btn = gr.Button("Clear chat", variant="secondary", scale=1)

    gr.Markdown("**Training-data probes**")
    with gr.Row(elem_id="probe-row"):
        btns_a = [gr.Button(label, size="sm") for label, _ in EXAMPLE_PROMPTS[:4]]
    with gr.Row(elem_id="probe-row-2"):
        btns_b = [gr.Button(label, size="sm") for label, _ in EXAMPLE_PROMPTS[4:]]
    example_btns = btns_a + btns_b

    gr.Markdown(
        "Weights: [`UnaverageTech411/northwind-ops`](https://huggingface.co/UnaverageTech411/northwind-ops) · "
        "Local: `ollama run northwind-ops` · "
        "Arriella custom-model factory POC"
    )

    send.click(respond, inputs=[msg, chatbot], outputs=[chatbot, msg])
    msg.submit(respond, inputs=[msg, chatbot], outputs=[chatbot, msg])
    clear_btn.click(clear_chat, outputs=[chatbot, msg])
    for btn, (_label, prompt) in zip(example_btns, EXAMPLE_PROMPTS):
        btn.click(
            lambda hist, p=prompt: use_example(p, hist),
            inputs=[chatbot],
            outputs=[chatbot, msg],
        )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()