// vp2vi browser workspace — paste → per-line guarded polish.
import { SpmTokenizer } from './engine/spm_tokenizer.js';
import { initDevice } from './engine/device.js';
import { loadWeights } from './engine/weights.js';
import { SRC_CAP } from './engine/constants.js';
import { guardedDecodeBatch, readTensorF32 } from './guarded-engine.js';
import { planDeviceBatches } from './device-fit.js';
import { prepareGuardedTopK } from './guarded-topk.js';

// Speculative-window executor (receipt ape_a2_webgpu_p23_v2: bit-identical
// to per-step on both builds, 12x fewer GPU<->CPU syncs).
const SPEC_WINDOW = 16;

const SAMPLES = Object.freeze({
  story: [
    'Hắn đối với chuyện này cũng không có biện pháp nào.',
    'Nghe được lời này, mọi người ở đây đều là hơi sững sờ.',
  ].join('\n'),
  dialogue: [
    '“Ngươi xác định muốn cùng ta động thủ?” hắn nhíu mày hỏi.',
    '“Nếu đã như vậy, vậy thì không cần nhiều lời.” nàng nhàn nhạt nói.',
  ].join('\n'),
  action: [
    'Kiếm quang xẹt qua, đem trước mặt cự thạch sinh sinh bổ thành hai nửa.',
    'Đối mặt một chưởng này, hắn không lùi mà tiến tới.',
    'Chỉ thấy thân hình hắn lóe lên, trong nháy mắt liền xuất hiện tại đối phương trước người.',
  ].join('\n'),
});

const $ = (id) => document.getElementById(id);
const status = (message, state = '', short = message) => {
  $('status').textContent = message;
  $('status').dataset.short = short;
  document.body.classList.remove('ready', 'running', 'failed');
  if (state) document.body.classList.add(state);
};

let copyText = '';

function assetBases() {
  // The source-tree dev server exposes artifacts at /artifacts; the built
  // Space places model and weights beside index.html.
  if (new URL(import.meta.url).pathname.includes('/src/webgpu-ar/')) {
    return {
      model: '/artifacts/ape_a2_webgpu/model/vp2vi',
      weights: '/artifacts/ape_a2_webgpu/weights_f16',
    };
  }
  return { model: './model/vp2vi', weights: './weights' };
}

function badge(text, info = false) {
  return `<span class="badge${info ? ' info' : ''}">${text}</span>`;
}

function meaningfulLines(value) {
  return value.split('\n').filter((line) => line.trim()).length;
}

function updateInputState() {
  const value = $('input').value;
  const lines = meaningfulLines(value);
  $('input-count').textContent = `${lines} dòng${value.length ? ` · ${value.length.toLocaleString('vi-VN')} ký tự` : ''}`;
  $('clear').disabled = !value || $('run').disabled && document.body.classList.contains('running');
}

function setSampleButtonsDisabled(disabled) {
  for (const button of document.querySelectorAll('[data-sample]')) {
    button.disabled = disabled;
  }
}

function loadSample(name) {
  const value = SAMPLES[name];
  if (!value) return;
  const input = $('input');
  input.value = value;
  updateInputState();
  input.focus();
  input.setSelectionRange(value.length, value.length);
}

function setProgress(done, total) {
  const active = total > 0 && done < total;
  $('progress').hidden = !active;
  $('progress').setAttribute('aria-hidden', String(!active));
  $('progress-bar').style.width = `${total ? Math.round((100 * done) / total) : 0}%`;
}

function adapterLabel(ctx) {
  const info = ctx.adapterInfo;
  const value = info.description || info.device || info.vendor || 'WebGPU';
  return value.replace(/\s+/g, ' ').trim();
}

async function writeClipboard(text) {
  try {
    await navigator.clipboard.writeText(text);
  } catch {
    const helper = document.createElement('textarea');
    helper.value = text;
    helper.setAttribute('readonly', '');
    helper.style.position = 'fixed';
    helper.style.opacity = '0';
    document.body.appendChild(helper);
    helper.select();
    document.execCommand('copy');
    helper.remove();
  }
}

async function warmEngine(ctx, weights, pieceTable, bias, tok) {
  const draft = 'Hắn gật đầu.';
  const srcIds = Array.from(tok.encode(draft, { addSpecialTokens: false }));
  const candidates = Array.from({ length: 64 }, (_, li) => ({
    id: `warm-${li}`,
    draft,
    srcIds,
    li,
  }));
  const { batches } = planDeviceBatches(candidates, ctx, weights.dtype);
  await guardedDecodeBatch(ctx, weights, pieceTable, bias, batches[0], { specWindow: SPEC_WINDOW });
}

async function boot() {
  $('input').addEventListener('input', updateInputState);
  for (const button of document.querySelectorAll('[data-sample]')) {
    button.addEventListener('click', () => loadSample(button.dataset.sample));
  }
  $('clear').addEventListener('click', () => {
    $('input').value = '';
    $('input').focus();
    updateInputState();
  });
  $('copy').addEventListener('click', async () => {
    if (!copyText) return;
    await writeClipboard(copyText);
    const label = $('copy').textContent;
    $('copy').textContent = 'Đã chép';
    setTimeout(() => { $('copy').textContent = label; }, 1400);
  });

  try {
    const bases = assetBases();
    status('Đang chuẩn bị tokenizer…');
    const [tokJson, tokConfig, pieceTable] = await Promise.all([
      fetch(`${bases.model}/tokenizer.json`).then((r) => r.json()),
      fetch(`${bases.model}/tokenizer_config.json`).then((r) => r.json()),
      fetch(`${bases.model}/piece_table.json`).then((r) => r.json()),
    ]);
    const tok = new SpmTokenizer(tokJson, tokConfig);
    const ctx = await initDevice();
    const topKReady = prepareGuardedTopK(ctx.device);
    status('Đang tải model…');
    const weights = await loadWeights(ctx.device, bases.weights, {
      targetDtype: ctx.hasF16 ? 'f16' : 'f32',
      // Value-preserving transposes let the decoder reuse weight tiles across
      // rows. They cost a small one-time upload, not output quality.
      ffnWT: ctx.hasF16,
      projWT: ctx.hasF16,
      onProgress: (loaded, total) => {
        const percent = Math.round((100 * loaded) / total);
        status(`Đang tải model · ${percent}%`, '', `Tải model · ${percent}%`);
      },
    });
    await topKReady;
    const bias = await readTensorF32(ctx.device, weights, 'final_logits_bias');
    status('Đang hoàn tất…');
    await warmEngine(ctx, weights, pieceTable, bias, tok);
    const readyText = `Sẵn sàng · ${adapterLabel(ctx)}`;
    status(readyText, 'ready');
    $('run').disabled = false;
    updateInputState();

    const run = async () => {
      const lines = $('input').value.split('\n').map((line) => line.trim());
      const jobs = [];
      lines.forEach((text, li) => { if (text) jobs.push({ li, text }); });
      if (!jobs.length) {
        status('Chưa có dòng nào để polish', 'ready');
        $('input').focus();
        return;
      }

      $('run').disabled = true;
      $('clear').disabled = true;
      $('copy').disabled = true;
      setSampleButtonsDisabled(true);
      copyText = '';
      status(`Đang polish · 0/${jobs.length} dòng`, 'running', `Đang chạy · 0/${jobs.length}`);
      setProgress(0, jobs.length);
      await new Promise((resolve) => requestAnimationFrame(resolve));

      const t0 = performance.now();
      const results = new Map();
      const runnable = [];
      for (const job of jobs) {
        const ids = Array.from(tok.encode(job.text, { addSpecialTokens: false }));
        if (!ids.length || ids.length > SRC_CAP) {
          results.set(job.li, { skipped: true, draft: job.text });
        } else {
          runnable.push({ id: String(job.li), draft: job.text, srcIds: ids, li: job.li });
        }
      }

      const { batches } = planDeviceBatches(runnable, ctx, weights.dtype);
      let doneRows = jobs.length - runnable.length;
      let outTokens = 0;
      let outChars = 0;
      let engineMs = 0;
      if (results.size) render(jobs, results);

      try {
        for (const chunk of batches) {
          status(
            `Đang polish · ${doneRows}/${jobs.length} dòng`,
            'running',
            `Đang chạy · ${doneRows}/${jobs.length}`,
          );
          const res = await guardedDecodeBatch(ctx, weights, pieceTable, bias, chunk, { specWindow: SPEC_WINDOW });
          engineMs += res.timing.encoderMs + res.timing.decodeMs;
          res.rows.forEach((row, i) => {
            const text = tok.decode(row.outputs, { skip_special_tokens: true }).trim();
            outTokens += row.outputs.length;
            outChars += text.length;
            results.set(chunk[i].li, {
              draft: chunk[i].draft,
              text,
              stats: row.stats,
              finished: row.finished,
            });
          });
          doneRows += chunk.length;
          setProgress(doneRows, jobs.length);
          render(jobs, results);
        }

        const seconds = (performance.now() - t0) / 1000;
        const throughputSeconds = Math.max(.001, engineMs / 1000);
        status(
          `Xong · ${jobs.length} dòng / ${seconds.toFixed(1)} giây · ${Math.round(outTokens / throughputSeconds)} token/giây · ${Math.round(outChars / throughputSeconds)} ký tự/giây`,
          'ready',
          `Xong · ${jobs.length} dòng · ${seconds.toFixed(1)}s`,
        );
        copyText = jobs.map((job) => {
          const result = results.get(job.li);
          return result?.skipped ? result.draft : result?.text ?? '';
        }).join('\n');
        $('copy').disabled = !copyText;
      } catch (err) {
        status(`Không thể hoàn tất · ${err?.message ?? err}`, 'failed');
        throw err;
      } finally {
        setProgress(1, 1);
        $('run').disabled = false;
        setSampleButtonsDisabled(false);
        updateInputState();
      }
    };

    $('run').addEventListener('click', run);
    $('input').addEventListener('keydown', (event) => {
      if ((event.ctrlKey || event.metaKey) && event.key === 'Enter' && !$('run').disabled) {
        event.preventDefault();
        run();
      }
    });
  } catch (err) {
    status(`Không thể khởi động · ${err?.message ?? err}`, 'failed');
    throw err;
  }
}

function render(jobs, results) {
  const parts = [];
  for (const job of jobs) {
    const result = results.get(job.li);
    if (!result) continue;
    if (result.skipped) {
      parts.push(`<article class="result-row flagged"><div class="draft">${esc(result.draft)}</div><div class="out">${esc(result.draft)}${badge('Dòng quá dài')}</div></article>`);
      continue;
    }
    const badges = [];
    if (result.stats.nameBroken) badges.push(badge('Kiểm tra tên riêng'));
    if (result.stats.digitFallback) badges.push(badge('Kiểm tra chữ số'));
    if (!result.finished) badges.push(badge('Chạm giới hạn độ dài'));
    if (result.stats.lengthStops > 0) badges.push(badge('Đã rút gọn', true));
    const flagged = result.stats.nameBroken || result.stats.digitFallback || !result.finished;
    parts.push(`<article class="result-row${flagged ? ' flagged' : ''}"><div class="draft">${esc(result.draft)}</div><div class="out">${esc(result.text)}${badges.join('')}</div></article>`);
  }
  $('results').innerHTML = parts.length
    ? parts.join('')
    : '<div class="empty-state"><span class="empty-glyph" aria-hidden="true">↗</span><p>Kết quả sẽ xuất hiện ở đây.</p></div>';
}

const esc = (value) => value.replace(/[&<>"]/g, (char) => ({
  '&': '&amp;', '<': '&lt;', '>': '&gt;', '"': '&quot;',
}[char]));

boot();