import { loadRuntime, encodePrompt, renderStory, MODEL_ID, REVISION } from './runtime.js';
import { sample, seededRandom, validateSettings } from './sampling.js';

let runtime, loading = false, running = false, stopped = false;
const send = (type, data = {}) => postMessage({ type, ...data });
const yieldToEvents = () => new Promise(resolve => setTimeout(resolve, 0));

self.onmessage = async ({ data }) => {
  try {
    if (data.type === 'stop') { stopped = true; return; }
    if (data.type === 'load') {
      if (running || loading) return;
      loading = true;
      runtime?.decoder.past.dispose();
      await runtime?.model.dispose();
      runtime = null;
      try {
        runtime = await loadRuntime(data.device, info => send('loading', { info }));
        send('ready', { device: runtime.device, fallback: runtime.fallback, vocabulary: runtime.vocabulary, provenance: runtime.provenance });
      } finally { loading = false; }
      return;
    }
    if (!runtime) throw new Error('Load the model first.');
    if (data.type === 'inspect') {
      if (!running) send('inspection', { ...encodePrompt(runtime.tokenizer, data.prompt), prompt: data.prompt });
      return;
    }
    if (data.type !== 'generate' || running) return;
    validateSettings(data.settings);
    const prompt = encodePrompt(runtime.tokenizer, data.prompt);
    if (prompt.ids.some(id => [0, 1, 2].includes(id))) throw new Error('Leave start, end and pad tokens out of your story opening.');
    const input = [1, ...prompt.ids];
    if (input.length >= 2048) throw new Error('Your opening fills the 2,048-token context. Shorten it to leave room for a continuation.');
    running = true;
    stopped = false;
    const { decoder, tokenizer, device, provenance } = runtime;
    decoder.reset();
    const generated = [], random = seededRandom(data.settings.seed);
    const started = performance.now();
    let logits, firstTokenMs = null, reason = 'token_limit';
    const emit = (type) => {
      const elapsedMs = performance.now() - started;
      const story = renderStory(tokenizer, [...prompt.ids, ...generated]);
      send(type, {
        ...story, ids: [...generated], promptIds: prompt.ids, unknown: prompt.unknown,
        settings: data.settings, prompt: data.prompt, device, reason,
        elapsedMs, firstTokenMs, tokens: generated.length,
        tokensPerSecond: generated.length ? generated.length / (elapsedMs / 1000) : 0,
        contextUsed: input.length + generated.length,
        model: { id: MODEL_ID, revision: REVISION, ...provenance },
      });
    };
    try {
      for (let i = 0; i < input.length; i++) {
        if (stopped) break;
        logits = await decoder.forward(input[i]);
        if (i % 8 === 0) await yieldToEvents();
      }
      const limit = Math.min(data.settings.maxTokens, 2048 - input.length);
      let lastUpdate = 0;
      for (let i = 0; i < limit && !stopped; i++) {
        const token = sample(logits, data.settings, random);
        generated.push(token);
        firstTokenMs ??= performance.now() - started;
        if (token === 2) { reason = 'eos'; break; }
        if (performance.now() - lastUpdate > 40) {
          emit('update');
          lastUpdate = performance.now();
          await yieldToEvents();
        }
        // No extra forward call is needed after the last sampled token.
        if (i + 1 < limit && !stopped) logits = await decoder.forward(token);
      }
      if (stopped) reason = 'stopped';
      else if (reason !== 'eos' && input.length + generated.length === 2048) reason = 'context_limit';
      emit('done');
    } finally { running = false; decoder.reset(); }
  } catch (error) {
    send('error', { message: error.message ?? String(error) });
  }
};