// Browser front end. Every request goes straight from the visitor to
// huggingface.co — this page has no backend and holds no token, so it
// reads exactly what the visitor could read themselves.

import {
  KNOWN_METHODS, EXCLUSION_FIELD, covers, moduleTree, modulePaths,
  resolveMethod, sameModule, unmatchedExclusions,
} from './pack_check.js';

const HF = 'https://huggingface.co';
const $ = (id) => document.getElementById(id);

// Tensor tails that mark a converted weight, used to compare a pack
// against its source by module rather than by tensor name.
const EXPECTED_UNQUANTIZED = new Set([
  'embed_tokens', 'lm_head', 'embed_out', 'A_log', 'dt_bias', 'conv1d',
]);

function expectedUnquantized(module) {
  const tail = module.split('.').pop();
  return EXPECTED_UNQUANTIZED.has(tail)
    || tail.endsWith('norm') || module.includes('.norm')
    || module.includes('embed')   // embeddings are not LinearBase
    || tail.includes('conv');     // convolutions are not either
}

const esc = (s) => String(s).replace(/[&<>]/g,
  (c) => ({ '&': '&amp;', '<': '&lt;', '>': '&gt;' }[c]));

async function json(url) {
  const r = await fetch(url);
  if (!r.ok) throw new Error(`${r.status} ${r.statusText}`);
  return r.json();
}

/** Tensor name -> shape, from safetensors headers only.
 *
 *  Two ranged requests per shard: eight bytes for the header length,
 *  then the header itself. A file of any size costs kilobytes. */
async function safetensorNames(repo) {
  const info = await json(`${HF}/api/models/${repo}`);
  const files = (info.siblings ?? [])
    .map((s) => s.rfilename)
    .filter((f) => f.endsWith('.safetensors'));

  const out = new Map();
  await Promise.all(files.map(async (file) => {
    const url = `${HF}/${repo}/resolve/main/${file}`;
    const head = await fetch(url, { headers: { Range: 'bytes=0-7' } });
    if (!head.ok) throw new Error(`${file}: ${head.status}`);
    const len = Number(new DataView(await head.arrayBuffer())
      .getBigUint64(0, true));
    const body = await fetch(url, {
      headers: { Range: `bytes=8-${8 + len - 1}` },
    });
    if (!body.ok) throw new Error(`${file}: ${body.status}`);
    for (const [name, entry] of Object.entries(JSON.parse(await body.text()))) {
      if (name !== '__metadata__') out.set(name, entry.shape ?? []);
    }
  }));
  return { names: out, info };
}

function declared(quant, method, module, entries) {
  if (covers(quant, method, module)) return true;
  return entries.some((e) => {
    const t = String(e);
    return !t.startsWith('re:') && !t.startsWith('-:') && !t.startsWith('+:')
      && sameModule(t, module);
  });
}

function card(kind, title, html) {
  return `<div class="card ${kind}"><h3>${esc(title)}</h3>${html}</div>`;
}

const list = (items, cap = 14) =>
  '<ul>' + items.slice(0, cap).map((i) => `<li>${esc(i)}</li>`).join('')
  + (items.length > cap ? `<li>… and ${items.length - cap} more</li>` : '')
  + '</ul>';

async function check(repo) {
  const { names, info } = await safetensorNames(repo);
  if (names.size === 0) {
    return card('problem', 'Nothing to read',
      '<p>This repo has no <code>.safetensors</code> files.</p>');
  }

  const config = await json(`${HF}/${repo}/resolve/main/config.json`);
  const quant = config.quantization_config ?? {};
  const method = resolveMethod(quant);
  const known = KNOWN_METHODS.has(method);
  const base = [].concat(info.cardData?.base_model ?? [])[0];
  const tied = Boolean(config.tie_word_embeddings
    ?? config.text_config?.tie_word_embeddings);

  const meta = `<div class="meta">
    <span><code>${esc(repo)}</code></span>
    <span>quantization: <code>${esc(method)}</code></span>
    <span>tensors: <code>${names.size.toLocaleString()}</code></span>
    ${base ? `<span>source: <code>${esc(base)}</code></span>` : ''}
  </div>`;

  if (Object.keys(quant).length === 0) {
    return meta + card('clear', 'Not a quantized pack',
      '<p>No <code>quantization_config</code>, so nothing here could have '
      + 'been quantized away.</p>');
  }
  if (!known) {
    return meta + card('check', `Format “${method}” is not implemented here`,
      '<p>The checks reproduce the exclusion rules of compressed-tensors, '
      + 'awq, gptq, bitsandbytes, modelopt and auto-round from vLLM\'s '
      + 'source. Yours is not among them, so they are skipped rather than '
      + 'guessed at.</p>');
  }

  const out = [];
  const tree = moduleTree(names.keys());

  // 1 — entries that name nothing
  const stray = unmatchedExclusions(quant, method, tree, tied);
  if (stray.length) {
    out.push(card('problem',
      `${stray.length} exclusion ${stray.length === 1 ? 'entry names' : 'entries name'} no module in this model`,
      '<p>They protect nothing, so whatever they were meant to keep was '
      + 'quantized:</p>' + list(stray)
      + '<p>A common cause is a vision tower: Llama and Pixtral call it '
      + '<code>vision_tower</code>, Qwen calls it <code>visual</code>.</p>'));
  }

  // 2 — modules that did not survive from the source
  if (base && base !== repo) {
    try {
      const src = await safetensorNames(base);
      const have = modulePaths(names.keys());
      const dropped = [...modulePaths(src.names.keys())]
        .filter((m) => !have.has(m)).sort();
      if (dropped.length) {
        const groups = {};
        for (const m of dropped) {
          const head = m.split('.')[0];
          groups[head] = (groups[head] ?? 0) + 1;
        }
        out.push(card('problem',
          `${dropped.length} modules of the source are absent here`,
          '<p><code>from_pretrained</code> does not materialise tensors the '
          + 'AutoModel class has no slot for, so they never reach the '
          + 'quantizer and never reach the artifact. For a '
          + 'multi-token-prediction head the symptom is 0% draft acceptance, '
          + 'with nothing in the logs.</p>'
          + list(Object.entries(groups)
            .sort((a, b) => b[1] - a[1])
            .map(([k, v]) => `${k} — ${v} modules`), 10)));
      }
    } catch (e) {
      out.push(card('check', 'Could not read the source model',
        `<p><code>${esc(base)}</code> is named as the base but could not be `
        + `read (${esc(e.message)}), so that check was skipped.</p>`));
    }
  }

  // 3 — present, unquantized, undeclared
  const entries = Object.keys(quant).length
    ? [].concat(quant[EXCLUSION_FIELD[method]] ?? []).flatMap(
        (v) => (typeof v === 'object' && v !== null ? Object.keys(v) : v))
    : [];
  const raw = quant[EXCLUSION_FIELD[method]];
  const entryList = Array.isArray(raw) ? raw
    : (raw && typeof raw === 'object' ? Object.keys(raw) : entries);
  const undeclared = [...new Set([...names.entries()]
    .filter(([n, shape]) => n.endsWith('.weight') && shape.length === 2)
    .map(([n]) => n.slice(0, -'.weight'.length))
    .filter((m) => !expectedUnquantized(m)
      && !declared(quant, method, m, entryList)))].sort();

  if (undeclared.length) {
    out.push(card('check',
      `${undeclared.length} modules are at source precision but not declared`,
      '<p>A runtime builds these quantized, looks for a packed weight, finds '
      + 'a plain one, and skips it — leaving the module randomly initialised '
      + 'while the model still answers correctly.</p>' + list(undeclared, 12)
      + '<p>This is a <strong>prediction</strong>, not a measurement. Confirm '
      + 'it by serving the pack and grepping the startup log for '
      + '<code>not found in params_dict</code>.</p>'));
  }

  if (!out.length) {
    out.push(card('clear', 'Nothing found',
      '<p>Every exclusion entry names a real module, no source module is '
      + 'missing, and nothing is left at source precision without being '
      + 'declared.</p>'));
  }
  return meta + out.join('');
}

async function submit(event) {
  event?.preventDefault();
  let repo = $('repo').value.trim().replace(/\/+$/, '');
  if (!repo) return;
  if (repo.startsWith('http')) repo = repo.split('huggingface.co/').pop();

  const out = $('out');
  const button = $('go');
  button.disabled = true;
  out.innerHTML = '<p class="spin">Reading headers…</p>';
  try {
    out.innerHTML = await check(repo);
  } catch (e) {
    out.innerHTML = card('problem', `Could not check ${repo}`,
      `<p><code>${esc(e.message)}</code></p><p>If the repo is gated or `
      + `private, this page cannot read it — it has no token.</p>`);
  } finally {
    button.disabled = false;
  }
}

$('form').addEventListener('submit', submit);
// Run once on load against the pre-filled example. A checker whose
// first screen is empty asks the visitor to trust it before it has
// shown them anything.
submit();
for (const b of document.querySelectorAll('[data-ex]')) {
  b.addEventListener('click', () => {
    $('repo').value = b.dataset.ex;
    submit();
  });
}