// Sirius Studio — Assistant panel.
//
// Two providers, both on the user's own hardware, no key and no spend:
//   webllm  — the model runs inside this tab on the GPU through WebGPU
//             (vendor/web-llm, Apache-2.0). Weights download once from the
//             MLC mirror and stay in the browser cache.
//   local   — any OpenAI-compatible endpoint on the user's machine (Ollama,
//             LM Studio, llama.cpp). The page talks to localhost directly; the
//             runtime has to allow this origin (OLLAMA_ORIGINS=https://silentmode.st).
//
// The model returns data (HTML + CSS or text as JSON); Studio inserts it
// through the editor API. Model output is never executed as code.

const LS_KEY = "siriusAi";
// Base ids; the quantisation suffix is chosen at load time: q4f16 when the
// GPU exposes the shader-f16 feature (half the memory), q4f32 otherwise.
const WEBLLM_MODELS = [
  { id: "Qwen2.5-Coder-1.5B-Instruct", label: "Qwen2.5 Coder 1.5B · best HTML/CSS (≈1 GB download)" },
  { id: "Qwen2.5-1.5B-Instruct", label: "Qwen2.5 1.5B · balanced copy + layout (≈1 GB)" },
  { id: "Qwen2.5-Coder-0.5B-Instruct", label: "Qwen2.5 Coder 0.5B · light, still builds sections (≈400 MB, 1 GB GPU)" },
  { id: "SmolLM2-360M-Instruct", label: "SmolLM2 360M · weakest machines; rewrites text, rarely builds sections (≈250 MB)" },
  { id: "Qwen2.5-Coder-3B-Instruct", label: "Qwen2.5 Coder 3B · strongest here (≈2 GB, needs 4 GB GPU)" },
  { id: "Llama-3.2-3B-Instruct", label: "Llama 3.2 3B · better prose (≈2 GB, needs 4 GB GPU)" },
];
let gpuF16 = null;
async function gpuSupportsF16() {
  if (gpuF16 != null) return gpuF16;
  try { const a = await navigator.gpu.requestAdapter(); gpuF16 = !!a?.features?.has("shader-f16"); } catch { gpuF16 = false; }
  return gpuF16;
}
const fullModelId = async (base) => `${base}-q4${(await gpuSupportsF16()) ? "f16" : "f32"}_1-MLC`;

export function initAssistant(editor, { esc, name }) {
  const $ = (id) => document.getElementById(id);
  const panel = $("ai"); if (!panel) return;
  const settings = Object.assign({ provider: "webllm", model: WEBLLM_MODELS[0].id, url: "http://localhost:11434/v1", lmodel: "" }, load());
  let engine = null, engineModel = null, worker = null, webllm = null, busy = false, abort = null, lastApplied = false;

  // ---------- ui ----------
  $("ai-model").innerHTML = WEBLLM_MODELS.map((m) => `<option value="${m.id}">${esc(m.label)}</option>`).join("");
  $("ai-model").value = settings.model;
  if (!$("ai-model").value) { $("ai-model").value = WEBLLM_MODELS[0].id; settings.model = WEBLLM_MODELS[0].id; }
  $("ai-provider").value = settings.provider;
  $("ai-url").value = settings.url;
  $("ai-lmodel").value = settings.lmodel;
  const showProvider = () => { const w = $("ai-provider").value === "webllm"; $("ai-webllm").hidden = !w; $("ai-local").hidden = w; };
  showProvider();
  const st = (text, cls = "") => { const el = $("ai-status"); el.textContent = text; el.className = "ai-status " + cls; };
  const out = (text) => { const el = $("ai-out"); el.textContent = text; el.scrollTop = 1e6; };
  const persist = () => { settings.provider = $("ai-provider").value; settings.model = $("ai-model").value; settings.url = $("ai-url").value.trim().replace(/\/+$/, ""); settings.lmodel = $("ai-lmodel").value.trim(); save(settings); };

  $("btn-ai").addEventListener("click", () => { panel.hidden = !panel.hidden; $("btn-ai").classList.toggle("on", !panel.hidden); setTimeout(() => editor.refresh(), 60); if (!panel.hidden) paintTarget(); });
  $("ai-close").addEventListener("click", () => { panel.hidden = true; $("btn-ai").classList.remove("on"); setTimeout(() => editor.refresh(), 60); });
  $("ai-provider").addEventListener("change", () => { showProvider(); persist(); st(readyText(), ""); });
  $("ai-model").addEventListener("change", persist);
  $("ai-url").addEventListener("change", persist);
  $("ai-lmodel").addEventListener("change", persist);
  $("ai-load").addEventListener("click", () => ensureWebLLM().catch((e) => st(e.message || String(e), "err")));
  $("ai-connect").addEventListener("click", () => probeLocal().catch((e) => st(e.message || String(e), "err")));
  $("ai-undo").addEventListener("click", () => { if (lastApplied) { editor.UndoManager.undo(); lastApplied = false; $("ai-undo").hidden = true; } });
  $("ai-stop").addEventListener("click", () => { try { abort?.abort(); engine?.interruptGenerate?.(); } catch {} });
  panel.querySelectorAll("[data-verb]").forEach((b) => b.addEventListener("click", () => run(b.dataset.verb).catch((e) => { st(friendly(e), "err"); })));
  function friendly(e) {
    const m = e?.message || String(e);
    if (/disposed|out of memory|OutOfMemory|device lost|DeviceLost/i.test(m)) {
      try { worker?.terminate(); } catch {} engine = null; engineModel = null;
      return "Your GPU ran out of memory for this model. Pick a smaller one (Qwen2.5 Coder 0.5B or SmolLM2) and press Load model again.";
    }
    return m;
  }
  editor.on("component:selected component:deselected component:toggled", paintTarget);
  if (!("gpu" in navigator)) { const o = $("ai-provider").querySelector('option[value="webllm"]'); if (o) o.textContent += " — no WebGPU here"; if (settings.provider === "webllm") st("This browser has no WebGPU. Use a local endpoint, or a browser with WebGPU.", "warn"); else st(readyText()); }
  else st(readyText());

  function readyText() { return $("ai-provider").value === "webllm" ? (engine && engineModel === $("ai-model").value ? "Model loaded — ask away" : "Model not loaded yet — press Load model (once)") : "Press Connect to check the local endpoint"; }
  function selected() { return editor.getSelected(); }
  function isText(c) { return !!c && (c.get("type") === "text" || c.is?.("text")); }
  function paintTarget() {
    const c = selected();
    const t = $("ai-target");
    if (!c) { t.textContent = "Nothing selected — “Make a section” adds at the end of the page."; return; }
    const tag = c.get("tagName") || "element", cls = (c.getClasses?.() || []).slice(0, 2).join("."), text = (c.getEl?.()?.innerText || "").trim().slice(0, 60);
    t.textContent = `Selected: <${tag}${cls ? "." + cls : ""}>${text ? " “" + text + (text.length >= 60 ? "…" : "") + "”" : ""}`;
  }

  // ---------- providers ----------
  async function ensureWebLLM() {
    if (!("gpu" in navigator)) throw new Error("WebGPU is not available in this browser.");
    const base = $("ai-model").value;
    if (engine && engineModel === base) return engine;
    const modelId = await fullModelId(base);
    if (!webllm) { st("Loading the engine…", "busy"); webllm = await import("../vendor/web-llm/index.js?v=0.2.85"); }
    $("ai-load").disabled = true;
    try {
      const progress = (r) => st((r.text || "Loading…").replace(/\[.*?\]\s*/g, "").slice(0, 140), "busy");
      // Switching models: a fresh worker every time. Reloading inside the same
      // engine left the runtime with disposed objects ("Object has already
      // been disposed" on the next generation).
      if (engine) { try { await engine.unload?.(); } catch {} try { worker?.terminate(); } catch {} engine = null; engineModel = null; }
      worker = new Worker(new URL("./ai-worker.js?v=20260920ai", import.meta.url), { type: "module" });
      engine = await webllm.CreateWebWorkerMLCEngine(worker, modelId, { initProgressCallback: progress });
      engineModel = base;
      st(`Model loaded (${(await gpuSupportsF16()) ? "f16" : "f32 — this GPU has no f16 shaders, so the slightly larger build"}) — runs on this device, nothing leaves it`, "ok");
      return engine;
    } finally { $("ai-load").disabled = false; }
  }
  async function probeLocal() {
    persist();
    st("Checking " + settings.url + "…", "busy");
    const r = await fetch(settings.url + "/models", { signal: AbortSignal.timeout(6000) }).catch((e) => { throw new Error(`Cannot reach ${settings.url}: ${e.message}. Is the runtime running and is this origin allowed (e.g. OLLAMA_ORIGINS=${location.origin})?`); });
    if (!r.ok) throw new Error(`Endpoint answered ${r.status}`);
    const j = await r.json().catch(() => ({}));
    const ids = (j.data || []).map((m) => m.id).filter(Boolean);
    $("ai-lmodels").innerHTML = ids.map((id) => `<option value="${esc(id)}">`).join("");
    if (!settings.lmodel && ids.length) { $("ai-lmodel").value = ids[0]; persist(); }
    st(`Connected — ${ids.length} model${ids.length === 1 ? "" : "s"} available${settings.lmodel ? ", using " + settings.lmodel : ""}`, "ok");
  }
  async function complete(messages, onToken) {
    const provider = $("ai-provider").value;
    let text = "";
    if (provider === "webllm") {
      const eng = await ensureWebLLM();
      st("Thinking on your GPU…", "busy");
      const stream = await eng.chat.completions.create({ messages, stream: true, temperature: 0.5, max_tokens: 1400 });
      for await (const chunk of stream) { const d = chunk.choices?.[0]?.delta?.content || ""; if (d) { text += d; onToken(text); } }
      return text;
    }
    persist();
    if (!settings.lmodel) throw new Error("Pick a model name for the local endpoint first (press Connect to list them).");
    abort = new AbortController();
    st(`Asking ${settings.lmodel} at ${settings.url}…`, "busy");
    const r = await fetch(settings.url + "/chat/completions", { method: "POST", headers: { "content-type": "application/json" }, signal: abort.signal, body: JSON.stringify({ model: settings.lmodel, messages, stream: true, temperature: 0.5 }) });
    if (!r.ok) throw new Error(`Endpoint answered ${r.status}: ${(await r.text().catch(() => "")).slice(0, 200)}`);
    const reader = r.body.getReader(); const dec = new TextDecoder(); let buf = "";
    for (;;) {
      const { value, done } = await reader.read(); if (done) break;
      buf += dec.decode(value, { stream: true });
      let nl; while ((nl = buf.indexOf("\n")) >= 0) {
        const line = buf.slice(0, nl).trim(); buf = buf.slice(nl + 1);
        if (!line.startsWith("data:")) continue;
        const data = line.slice(5).trim(); if (data === "[DONE]") continue;
        try { const d = JSON.parse(data).choices?.[0]?.delta?.content || ""; if (d) { text += d; onToken(text); } } catch {}
      }
    }
    return text;
  }

  // ---------- prompts + apply ----------
  const SYSTEM = `You are a web designer working inside a visual site builder. Reply with a single JSON object and nothing else: no prose, no markdown fences, no explanations. Never emit <script>, <iframe>, <link>, <style> tags, event handlers or external URLs other than images already given. Write clean semantic HTML and plain CSS. Every CSS class you create starts with the given prefix and every rule targets those classes. Keep copy short and honest; keep the language of the page.`;
  function palette() {
    try {
      const d = document.querySelector("#gjs iframe").contentDocument; const cs = getComputedStyle(d.body);
      return `page background ${cs.backgroundColor}, text ${cs.color}, font ${cs.fontFamily.split(",")[0]}`;
    } catch { return "unknown"; }
  }
  const strip = (raw) => String(raw).replace(/<think>[\s\S]*?<\/think>/g, "").replace(/```(?:json|html|css)?/gi, "").trim();
  function parseJson(raw) {
    let s = strip(raw);
    const a = s.indexOf("{"), b = s.lastIndexOf("}");
    if (a >= 0 && b > a) s = s.slice(a, b + 1);
    try { const j = JSON.parse(s); return j && typeof j === "object" ? j : null; } catch { return null; }
  }
  // Small models often answer with bare HTML or CSS instead of the JSON we
  // asked for. Accept that too: tags → html (+ any <style> as css); a block
  // of rules without tags → css.
  function parseLoose(raw) {
    const s = strip(raw);
    if (/<[a-z][^>]*>/i.test(s)) {
      const css = (s.match(/<style[^>]*>[\s\S]*?<\/style>/gi) || []).map((m) => m.replace(/<\/?style[^>]*>/gi, "")).join("\n");
      const html = s.replace(/<style[\s\S]*?<\/style>/gi, "").replace(/^[\s\S]*?(?=<)/, "").trim();
      return html ? { html, css } : null;
    }
    if (/\{[\s\S]*\}/.test(s) && /[:;]/.test(s)) return { css: s.slice(s.indexOf(".") >= 0 ? 0 : 0) };
    return null;
  }
  const sanitize = (html) => String(html || "").replace(/<\s*(script|iframe|object|embed|link|meta|style)[^>]*>[\s\S]*?<\s*\/\s*\1\s*>/gi, "").replace(/<\s*(script|iframe|object|embed|link|meta)[^>]*\/?>/gi, "").replace(/\son\w+\s*=\s*("[^"]*"|'[^']*'|[^\s>]+)/gi, "").replace(/javascript:/gi, "");
  const sanitizeCss = (css) => String(css || "").replace(/@import[^;]*;/gi, "").replace(/expression\s*\(/gi, "").replace(/url\(\s*['"]?\s*(?!data:image|https?:)[^)]*\)/gi, "none");

  async function run(verb) {
    if (busy) return;
    const ask = $("ai-prompt").value.trim();
    const sel = selected();
    if (verb !== "section" && !sel) { st("Select something on the page first.", "warn"); return; }
    if (verb === "text" && !isText(sel)) { st("Select a text block (heading, paragraph, button label) to rewrite.", "warn"); return; }
    if (!ask && verb !== "text") { st("Describe what you want first.", "warn"); return; }
    busy = true; $("ai-stop").hidden = false; $("ai-undo").hidden = true; out("");
    const prefix = "ai" + Math.random().toString(36).slice(2, 6);
    try {
      let messages, raw, data;
      if (verb === "section") {
        messages = [{ role: "system", content: SYSTEM }, { role: "user", content: `Create one <section> for the site "${name}". Request: ${ask}. Match this palette unless told otherwise: ${palette()}. Class prefix: "${prefix}-". Answer with a JSON object with two string keys, html (the complete section markup) and css (the rules for its classes).` }];
        raw = await complete(messages, out); data = parseJson(raw); if (!/<[a-z][^>]*>/i.test(data?.html || "")) data = parseLoose(raw);
        if (!/<[a-z][^>]*>/i.test(data?.html || "")) throw new Error("The model did not return a usable section. Try again, rephrase, or pick a larger model (Qwen2.5 Coder).");
        const html = sanitize(data.html), css = sanitizeCss(data.css);
        const parent = sel?.parent?.(); const at = sel ? sel.index() + 1 : undefined;
        const added = parent ? parent.components().add(html, { at }) : editor.addComponents(html);
        if (css) editor.Css.addRules(css);
        const first = Array.isArray(added) ? added[0] : added; if (first) { editor.select(first); first.getEl?.()?.scrollIntoView?.({ behavior: "smooth", block: "center" }); }
        lastApplied = true; st("Section added — edit it like any other block", "ok");
      } else if (verb === "text") {
        const current = (sel.getEl()?.innerText || sel.toHTML().replace(/<[^>]+>/g, " ")).trim();
        messages = [{ role: "system", content: SYSTEM }, { role: "user", content: `Rewrite this text${ask ? " — instruction: " + ask : " to be clearer and shorter"}. Keep its language and meaning, no quotes around it. Text: """${current.slice(0, 1500)}""" Answer with a JSON object with one string key, text, holding the rewritten text.` }];
        raw = await complete(messages, out); data = parseJson(raw);
        const text = (data?.text ?? "").trim() || String(raw).replace(/<think>[\s\S]*?<\/think>/g, "").trim();
        if (!text) throw new Error("Empty answer.");
        sel.components(esc(text));
        lastApplied = true; st("Text replaced", "ok");
      } else if (verb === "style") {
        sel.addClass(prefix);
        const snippet = sel.toHTML().slice(0, 1500);
        messages = [{ role: "system", content: SYSTEM }, { role: "user", content: `Restyle this element. Request: ${ask}. Page palette: ${palette()}. The element already has the class "${prefix}"; write CSS rules for ".${prefix}" and its descendants only. Element: ${snippet} Answer with a JSON object with one string key, css, holding the rules.` }];
        raw = await complete(messages, out); data = parseJson(raw); if (!/\{[\s\S]*:[\s\S]*\}/.test(data?.css || "")) data = parseLoose(raw);
        if (!/\{[\s\S]*:[\s\S]*\}/.test(data?.css || "")) { sel.removeClass(prefix); throw new Error("The model did not return usable CSS. Try again, rephrase, or pick a larger model."); }
        editor.Css.addRules(sanitizeCss(data.css));
        lastApplied = true; st("Styles applied to the selection", "ok");
      }
      $("ai-undo").hidden = false;
    } finally { busy = false; $("ai-stop").hidden = true; abort = null; }
  }

  function load() { try { return JSON.parse(localStorage.getItem(LS_KEY) || "{}"); } catch { return {}; } }
  function save(s) { try { localStorage.setItem(LS_KEY, JSON.stringify(s)); } catch {} }
}