// Runs SmolLM2-360M-Instruct (q4f16, ~273 MB) on the WASM backend, off the main thread.
import * as tf from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3";

const { pipeline, TextStreamer, InterruptableStoppingCriteria } = tf;

const MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct";
let generator = null;

self.onmessage = async (e) => {
  const msg = e.data;
  try {
    if (msg.type === "load") {
      const progress_callback = (data) => {
        // Cache hits send non-cloneable payloads sometimes; keep it plain.
        const { status, file, progress, loaded, total } = data;
        self.postMessage({ type: "progress", data: { status, file, progress, loaded, total } });
      };
      // Prefer WebGPU when the adapter supports fp16 (q4f16 needs it);
      // fall back to WASM otherwise or if WebGPU init fails.
      let device = "wasm";
      try {
        const adapter = await self.navigator?.gpu?.requestAdapter?.();
        if (adapter?.features?.has("shader-f16")) device = "webgpu";
      } catch {}
      try {
        generator = await pipeline("text-generation", MODEL_ID, {
          dtype: "q4f16",
          device,
          progress_callback,
        });
      } catch (err) {
        if (device === "wasm") throw err;
        device = "wasm";
        generator = await pipeline("text-generation", MODEL_ID, {
          dtype: "q4f16",
          device,
          progress_callback,
        });
      }
      self.postMessage({ type: "ready", device });
    } else if (msg.type === "condense") {
      const line = await condensePair(msg.pair, msg.before, msg.after);
      self.postMessage({ type: "result", id: msg.id, line });
    } else if (msg.type === "expand") {
      const lines = await expandLine(msg.line, msg.before, msg.after);
      self.postMessage({ type: "result", id: msg.id, lines });
    }
  } catch (err) {
    self.postMessage({ type: "error", id: msg.id, message: err.message ?? String(err) });
  }
};

// Apply the chat template ourselves, then append `prefill` so the model's first
// generated token is already inside the JSON value — it cannot add a preamble.
// `stopPattern` interrupts generation as soon as the JSON value is closed.
async function generateConstrained(messages, prefill, stopPattern, maxNewTokens) {
  const tokenizer = generator.tokenizer;
  const prompt =
    tokenizer.apply_chat_template(messages, { tokenize: false, add_generation_prompt: true }) +
    prefill;

  let acc = "";
  const opts = {
    max_new_tokens: maxNewTokens,
    do_sample: true,
    temperature: 0.7,
    top_p: 0.9,
    repetition_penalty: 1.15,
    no_repeat_ngram_size: 3,
    return_full_text: false,
  };
  if (typeof InterruptableStoppingCriteria === "function" && typeof TextStreamer === "function") {
    const stopper = new InterruptableStoppingCriteria();
    opts.stopping_criteria = stopper;
    opts.streamer = new TextStreamer(tokenizer, {
      skip_prompt: true,
      skip_special_tokens: true,
      callback_function: (text) => {
        acc += text;
        if (stopPattern.test(acc)) stopper.interrupt();
      },
    });
  }
  const out = await generator(prompt, opts);
  return out[0].generated_text;
}

function unescapeJson(s) {
  try {
    return JSON.parse('"' + s + '"');
  } catch {
    return s;
  }
}

function tidy(s) {
  return s
    .replace(/\s+/g, " ")
    // JSON debris the parse fallback can drag along, e.g. a trailing '} or "]}.
    .replace(/["'’\s]*[\]}]+["'’\s]*$/, "")
    .trim();
}

// If the model looped and never closed the quote, the parse fallback swallows
// the whole ramble. Cut an overlong line back at the seam nearest the budget.
function clampLine(line, budget) {
  if (line.length <= budget) return line;
  const seams = [...line.matchAll(/[,;.—–]\s+|\s+(?:and|but|or|while|when|so)\s+/g)]
    .filter((m) => m.index > 0 && m.index <= budget);
  if (seams.length) return line.slice(0, seams.at(-1).index).trim();
  const cut = line.lastIndexOf(" ", budget);
  return line.slice(0, cut > 0 ? cut : budget).trim();
}

function contextNote(before, after) {
  let note = "";
  if (before?.length) note += `\nFor context, the lines just before them:\n${before.join("\n")}\n`;
  if (after?.length) note += `\nFor context, the lines just after them:\n${after.join("\n")}\n`;
  return note;
}

// A literal placeholder like {"line": "your new line here"} gets parroted
// verbatim by a 360M model, so the format is taught with a real few-shot
// exchange instead — the condense and expand examples are inverses.
const EXAMPLE_TWO = ["the rain fell all night on the tin roof", "and we listened without speaking"];
const EXAMPLE_ONE = "all night we listened to the rain speak on the tin roof";

async function condensePair([a, b], before, after) {
  const messages = [
    {
      role: "system",
      content:
        "You are a poet who condenses poems. You reply with exactly one JSON object and nothing else.",
    },
    {
      role: "user",
      content: `Two consecutive lines of a poem:

Line A: ${EXAMPLE_TWO[0]}
Line B: ${EXAMPLE_TWO[1]}

Combine Line A and Line B into ONE short poetic line that keeps their imagery, meaning, and voice. Reply with JSON only.`,
    },
    {
      role: "assistant",
      content: `{"line": "${EXAMPLE_ONE}"}`,
    },
    {
      role: "user",
      content: `Two consecutive lines of a poem:

Line A: ${a}
Line B: ${b}
${contextNote(before, after)}
Combine Line A and Line B into ONE short poetic line that keeps their imagery, meaning, and voice. Reply with JSON only.`,
    },
  ];

  // Generation continues from `{"line": "` — capture up to the first unescaped quote.
  const raw = await generateConstrained(messages, '{"line": "', /(?:^|[^\\])"/, 60);
  const m = raw.match(/^((?:[^"\\]|\\.)*)"/s);
  const line = clampLine(
    tidy(unescapeJson(m ? m[1] : raw.split("\n")[0])),
    Math.max(50, a.length + b.length),
  );
  // A condensation that vanished is worse than no condensation: keep the originals.
  return line || tidy(`${a} ${b}`);
}

// Split one string into two lines at a natural seam (for fallback parsing).
function splitInTwo(s) {
  const mid = s.length / 2;
  const seams = [...s.matchAll(/[,;—–]\s+|\s+(?:and|but|or|while|when|so)\s+/g)];
  if (seams.length) {
    const best = seams.reduce((p, c) => (Math.abs(c.index - mid) < Math.abs(p.index - mid) ? c : p));
    return [s.slice(0, best.index + (best[0].startsWith(" ") ? 0 : 1)).trim(), s.slice(best.index + best[0].length).trim()];
  }
  const words = s.split(" ");
  const cut = Math.ceil(words.length / 2);
  return [words.slice(0, cut).join(" "), words.slice(cut).join(" ")];
}

async function expandLine(line, before, after) {
  const messages = [
    {
      role: "system",
      content:
        "You are a poet who elaborates poems. You reply with exactly one JSON object and nothing else.",
    },
    {
      role: "user",
      content: `One line of a poem:

Line: ${EXAMPLE_ONE}

Rewrite this single line as TWO short poetic lines that together carry the same meaning, imagery, and voice. Reply with JSON only.`,
    },
    {
      role: "assistant",
      content: `{"lines": ["${EXAMPLE_TWO[0]}", "${EXAMPLE_TWO[1]}"]}`,
    },
    {
      role: "user",
      content: `One line of a poem:

Line: ${line}
${contextNote(before, after)}
Rewrite this single line as TWO short poetic lines that together carry the same meaning, imagery, and voice. Reply with JSON only.`,
    },
  ];

  // Generation continues from `{"lines": ["` — stop once the array closes.
  const raw = await generateConstrained(messages, '{"lines": ["', /"\s*\]/, 90);
  const two = raw.match(/^((?:[^"\\]|\\.)*)"\s*,\s*"((?:[^"\\]|\\.)*)/s);
  const budget = Math.max(60, line.length * 2);
  if (two) {
    const first = clampLine(tidy(unescapeJson(two[1])), budget);
    const second = clampLine(tidy(unescapeJson(two[2])), budget);
    if (first && second) return [first, second];
  }
  const one = raw.match(/^((?:[^"\\]|\\.)*)"/s);
  const text = clampLine(tidy(unescapeJson(one ? one[1] : raw.split("\n")[0])), budget * 2);
  if (!text) return [tidy(line)];
  return splitInTwo(text);
}