/**
 * The page: one repo id in, one honest breakdown out.
 *
 * Every rule lives in precision_map.js, which the parity gate holds to
 * the Python implementation. Nothing here decides anything — it fetches,
 * calls, and renders, so a rule can never be quietly different in the UI
 * from the one that was tested.
 */
import { FORMATS, RULES_MEASURED_ON, inspect } from './precision_map.js';
import { progress } from './progress.js';
import { ensureNextStepStyles, nextSteps, repoFromUrl } from './tools.js';

const $ = (id) => document.getElementById(id);

// A link from another tool carries the model over; arriving at an empty
// form is what makes five tools feel like five tools instead of one.
ensureNextStepStyles();
const carried = repoFromUrl();
if (carried) { $('repo').value = carried; }

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

const GIB = 1024 ** 3;
const gib = (b) => (b / GIB >= 0.01 ? `${(b / GIB).toFixed(2)} GiB`
  : `${(b / (1024 ** 2)).toFixed(0)} MiB`);
const pct = (part, whole) => `${((part / whole) * 100).toFixed(1)}%`;

function renderFormats() {
  $('formats').innerHTML = Object.entries(FORMATS).map(([method, f]) =>
    `<tr>
       <td><code>${esc(method)}</code></td>
       <td>${esc(f.tool)}</td>
       <td>${f.producer ? `<code>${esc(f.producer.join('.').replace('config.', ''))}</code>`
    : '<span class="muted">records none</span>'}</td>
       <td class="small">${esc(f.verifiedOn)}</td>
     </tr>`).join('');
}

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

/** The stacked bar. Three categories, in the order they are explained. */
function bar(b, total) {
  const seg = (k, cls) => (b[k] > 0
    ? `<span class="seg ${cls}" style="width:${(b[k] / total) * 100}%"
         title="${esc(k)} — ${gib(b[k])}"></span>` : '');
  return `<div class="bar">${seg('quantized', 'q')}${seg('metadata', 'm')}${seg('full', 'f')}${seg('buffer', 'b')}</div>`;
}

function render(r) {
  const b = r.buckets;
  const t = r.totalBytes;

  const version = r.producerVersion
    ? `<code>${esc(r.producerVersion)}</code>`
    : '<span class="muted">no version recorded in the pack</span>';

  const meta = `<div class="meta">
      <span>repo <code>${esc(r.repo)}</code></span>
      <span>declared <code>${esc(r.method)}</code></span>
      <span>${r.shards} shard${r.shards === 1 ? '' : 's'}, ${r.tensorCount} tensors</span>
    </div>`;

  if (!r.knownFormat && r.method !== 'none') {
    return meta + card('check', `No convention on file for ${esc(r.method)}`,
      `<p>This reader has never been verified against a
       <code>${esc(r.method)}</code> pack, so the split below rests on dtypes
       alone and the metadata line may be undercounted. Treat it as an
       estimate, not a reading.</p>` + bar(b, t) + table(r));
  }

  if (!r.recognised) {
    return meta + card('problem', 'This layout is not understood',
      `<p>The pack declares <code>${esc(r.method)}</code>, which quantizes
       something — but no tensor here is in a packed dtype or under a packed
       name. Either the convention moved since ${esc(RULES_MEASURED_ON)}, or
       this pack stores its weights somewhere this cannot see.</p>
       <p>No breakdown is shown, on purpose. A number here would look like a
       measurement and be a guess.</p>`);
  }

  if (r.method === 'none') {
    return meta + card('clear', 'Not quantized',
      `<p>Every tensor is at full precision — ${esc(gib(b.full))} of it. No
       quantization method is declared and none is detectable from the
       dtypes.</p>` + tableFamilies(r, b.full));
  }

  const headline = `${pct(b.full, t)} of this pack is still full precision`;

  return meta + card('check', headline,
    `<p>Written by <strong>${esc(r.tool)}</strong>, ${version}.</p>`
    + bar(b, t)
    + table(r)
    + `<h4>What stayed whole</h4>`
    + tableFamilies(r, b.full)
    + (r.headQuantized
      ? `<p class="warn"><strong>The output head is quantized.</strong> That is
         a real choice with a real cost: measured on this repo's own weights,
         quantizing <code>lm_head</code> flips the model's chosen token on
         about a fifth of positions. It buys size; it is worth knowing you
         bought it.</p>`
      : `<p class="note">The output head was left at full precision — the
         usual choice, and the reason a "4-bit" pack is never 4-bit
         throughout.</p>`)
    + (r.unknownDtypes.length
      ? `<p class="warn">Unknown dtypes counted at 2 bytes:
         <code>${r.unknownDtypes.map(esc).join('</code>, <code>')}</code>.</p>`
      : ''));
}

function table(r) {
  const b = r.buckets;
  const t = r.totalBytes;
  const row = (label, key, why) => (b[key] > 0
    ? `<tr><td>${label}</td><td class="num">${gib(b[key])}</td>
         <td class="num">${pct(b[key], t)}</td><td class="small">${why}</td></tr>`
    : '');
  return `<div class="scroll"><table class="rows">
    ${row('Quantized payload', 'quantized', 'the weights that were compressed')}
    ${row('Quantization data', 'metadata', 'scales, zero points, group indices')}
    ${row('Full precision', 'full', 'never quantized, at source precision')}
    ${row('Index buffers', 'buffer', 'positions, masks, stored shapes')}
  </table></div>`;
}

function tableFamilies(r, fullTotal) {
  const rows = Object.entries(r.fullByFamily);
  if (!rows.length) return '';
  return `<div class="scroll"><table class="rows">${rows.map(([fam, size]) =>
    `<tr><td>${esc(fam)}</td><td class="num">${gib(size)}</td>
       <td class="num">${pct(size, fullTotal)}</td></tr>`).join('')}
  </table></div>`;
}

async function run(repo) {
  const go = $('go');
  go.disabled = true;
  // Indeterminate until the shard count is known — reading config.json
  // and the shard index says nothing about how much work follows.
  const pg = progress(out).start(`reading ${repo}`);
  try {
    const r = await inspect(repo, (done, total) => {
      if (done === 0) {
        pg.total(total).start(total === 1
          ? 'reading the safetensors header'
          : `reading ${total} safetensors headers`);
      } else {
        pg.step();
      }
    });
    pg.done();
    out.innerHTML = render(r) + nextSteps('precision', repo);
  } catch (e) {
    out.innerHTML = card('problem', 'Could not read the repo',
      `<p>${esc(e.message || e)}</p>
       <p>Gated and private repos cannot be read from the browser.</p>`);
  } finally {
    go.disabled = false;
  }
}

$('form').addEventListener('submit', (e) => {
  e.preventDefault();
  const v = $('repo').value.trim();
  if (v) run(v);
});
document.querySelectorAll('[data-ex]').forEach((btn) =>
  btn.addEventListener('click', () => {
    $('repo').value = btn.dataset.ex;
    run(btn.dataset.ex);
  }));

$('measured').textContent = RULES_MEASURED_ON;
renderFormats();