'use strict';

const APP = Object.freeze({
  modelUrl: './ms_se_efficientnet_b0.onnx?v=ac619ed4',
  ortVersion: '1.27.0',
  imageSize: 224,
  mean: [0.485, 0.456, 0.406],
  std: [0.229, 0.224, 0.225],
  labels: [
    'battery', 'biological', 'cardboard', 'clothes', 'glass',
    'metal', 'paper', 'plastic', 'shoes', 'trash',
  ],
  metadata: {
    input: { width: 224, height: 224, channels: 3 },
    complexity: { parameters: 4234159 },
  },
});

const state = {
  session: null,
  labels: APP.labels,
  metadata: APP.metadata,
  bitmap: null,
  objectUrl: null,
  file: null,
  modelReady: false,
  predicting: false,
};

const els = {
  runtimePill: document.querySelector('#runtimePill'),
  runtimeText: document.querySelector('#runtimeText'),
  statusDot: document.querySelector('#statusDot'),
  statusCard: document.querySelector('.status-card'),
  progressBar: document.querySelector('#progressBar'),
  loadPercent: document.querySelector('#loadPercent'),
  statusDetail: document.querySelector('#statusDetail'),
  fileInput: document.querySelector('#fileInput'),
  dropZone: document.querySelector('#dropZone'),
  clearButton: document.querySelector('#clearButton'),
  previewWrap: document.querySelector('#previewWrap'),
  imagePreview: document.querySelector('#imagePreview'),
  fileName: document.querySelector('#fileName'),
  fileDetails: document.querySelector('#fileDetails'),
  predictButton: document.querySelector('#predictButton'),
  predictButtonText: document.querySelector('#predictButtonText'),
  buttonSpinner: document.querySelector('#buttonSpinner'),
  inputMessage: document.querySelector('#inputMessage'),
  emptyResult: document.querySelector('#emptyResult'),
  results: document.querySelector('#results'),
  predictedClass: document.querySelector('#predictedClass'),
  confidenceRing: document.querySelector('#confidenceRing'),
  confidenceValue: document.querySelector('#confidenceValue'),
  topThree: document.querySelector('#topThree'),
  probabilityList: document.querySelector('#probabilityList'),
  preprocessTime: document.querySelector('#preprocessTime'),
  inferenceTime: document.querySelector('#inferenceTime'),
  totalTime: document.querySelector('#totalTime'),
  providerBadge: document.querySelector('#providerBadge'),
  classChips: document.querySelector('#classChips'),
  modelFacts: document.querySelector('#modelFacts'),
  canvas: document.querySelector('#preprocessCanvas'),
};

function setProgress(percent, detail) {
  const safePercent = Math.max(0, Math.min(100, Math.round(percent)));
  els.progressBar.style.width = `${safePercent}%`;
  els.loadPercent.textContent = `${safePercent}%`;
  if (detail) els.statusDetail.textContent = detail;
}

function setRuntimeState(kind, text) {
  els.runtimePill.classList.remove('ready', 'error');
  if (kind) els.runtimePill.classList.add(kind);
  els.runtimeText.textContent = text;
}

function setError(message, error) {
  console.error(message, error || '');
  state.modelReady = false;
  els.statusCard.classList.add('error');
  setRuntimeState('error', 'Runtime unavailable');
  setProgress(100, message);
  els.inputMessage.textContent = 'The model could not be initialized. Refresh the page or try another modern browser.';
  els.inputMessage.classList.add('error');
  updatePredictButton();
}

function titleCase(value) {
  return String(value)
    .replace(/[_-]+/g, ' ')
    .replace(/\b\w/g, character => character.toUpperCase());
}

function formatBytes(bytes) {
  if (!Number.isFinite(bytes) || bytes <= 0) return '';
  const units = ['B', 'KB', 'MB', 'GB'];
  const index = Math.min(Math.floor(Math.log(bytes) / Math.log(1024)), units.length - 1);
  return `${(bytes / (1024 ** index)).toFixed(index >= 2 ? 1 : 0)} ${units[index]}`;
}

function formatMilliseconds(value) {
  if (!Number.isFinite(value)) return '—';
  if (value < 10) return `${value.toFixed(2)} ms`;
  if (value < 100) return `${value.toFixed(1)} ms`;
  return `${Math.round(value)} ms`;
}

function updatePredictButton() {
  const enabled = state.modelReady && Boolean(state.bitmap) && !state.predicting;
  els.predictButton.disabled = !enabled;
  els.clearButton.disabled = !state.bitmap && !state.file;
}

async function fetchBinaryWithProgress(url, onProgress) {
  const response = await fetch(url, { cache: 'force-cache' });
  if (!response.ok) throw new Error(`Model download failed: HTTP ${response.status}`);

  const total = Number(response.headers.get('content-length')) || 0;
  if (!response.body || !total) {
    const buffer = await response.arrayBuffer();
    onProgress(1, buffer.byteLength);
    return new Uint8Array(buffer);
  }

  const reader = response.body.getReader();
  const chunks = [];
  let received = 0;

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    received += value.length;
    onProgress(received / total, total);
  }

  const bytes = new Uint8Array(received);
  let offset = 0;
  for (const chunk of chunks) {
    bytes.set(chunk, offset);
    offset += chunk.length;
  }
  return bytes;
}

function renderClassChips() {
  els.classChips.replaceChildren();
  for (const label of state.labels) {
    const chip = document.createElement('span');
    chip.className = 'class-chip';
    chip.textContent = label;
    els.classChips.appendChild(chip);
  }
}

function renderModelFacts() {
  if (!state.metadata) return;
  const facts = [
    ['Input', `${state.metadata.input.width} × ${state.metadata.input.height} RGB`],
    ['Classes', String(state.labels.length)],
    ['Parameters', `${(state.metadata.complexity.parameters / 1e6).toFixed(3)}M`],
    ['Runtime', 'WebAssembly'],
  ];
  els.modelFacts.replaceChildren();
  for (const [term, value] of facts) {
    const wrapper = document.createElement('div');
    const dt = document.createElement('dt');
    const dd = document.createElement('dd');
    dt.textContent = term;
    dd.textContent = value;
    wrapper.append(dt, dd);
    els.modelFacts.appendChild(wrapper);
  }
}

async function initializeModel() {
  try {
    if (!window.ort) throw new Error('ONNX Runtime Web did not load from the pinned CDN.');

    setRuntimeState('', 'Preparing model');
    setProgress(3, 'Preparing class labels and model information…');
    renderClassChips();
    renderModelFacts();

    ort.env.logLevel = 'error';
    ort.env.wasm.numThreads = 1;
    ort.env.wasm.proxy = false;
    ort.env.wasm.wasmPaths = `https://cdn.jsdelivr.net/npm/onnxruntime-web@${APP.ortVersion}/dist/`;

    setRuntimeState('', 'Downloading model');
    const modelBytes = await fetchBinaryWithProgress(APP.modelUrl, (fraction, totalBytes) => {
      const percent = 5 + (fraction * 66);
      setProgress(percent, `Downloading ${formatBytes(totalBytes) || 'ONNX model'}…`);
    });

    setProgress(75, 'Creating optimized WebAssembly inference session…');
    setRuntimeState('', 'Creating session');
    state.session = await ort.InferenceSession.create(modelBytes, {
      executionProviders: ['wasm'],
      graphOptimizationLevel: 'all',
    });

    const inputNames = Array.isArray(state.session.inputNames) ? state.session.inputNames : [];
    const outputNames = Array.isArray(state.session.outputNames) ? state.session.outputNames : [];
    if (!inputNames.includes('input') || !outputNames.includes('logits')) {
      throw new Error(`Unexpected ONNX interface: inputs ${inputNames.join(', ') || 'none'}, outputs ${outputNames.join(', ') || 'none'}`);
    }

    // Run one deterministic smoke test before enabling the interface. This catches
    // invalid model files, unsupported operators, and output-shape mismatches.
    setProgress(92, 'Verifying model execution and output shape…');
    const testData = new Float32Array(3 * APP.imageSize * APP.imageSize);
    const testTensor = new ort.Tensor('float32', testData, [1, 3, APP.imageSize, APP.imageSize]);
    const testOutputs = await state.session.run({ input: testTensor });
    const testLogits = testOutputs.logits && testOutputs.logits.data;
    if (!testLogits || testLogits.length !== state.labels.length) {
      const actualLength = testLogits ? testLogits.length : 0;
      throw new Error(`ONNX smoke test failed: expected ${state.labels.length} logits, received ${actualLength}.`);
    }
    if (!Array.from(testLogits).every(Number.isFinite)) {
      throw new Error('ONNX smoke test failed: output contains non-finite values.');
    }

    state.modelReady = true;
    els.statusCard.classList.remove('error');
    setProgress(100, 'Ready. The ONNX model has loaded successfully.');
    setRuntimeState('ready', 'Model ready');
    els.providerBadge.textContent = 'WASM';
    updatePredictButton();
  } catch (error) {
    setError(error.message || 'Model initialization failed.', error);
  }
}

async function decodeImage(file) {
  if ('createImageBitmap' in window) {
    try {
      return await createImageBitmap(file, { imageOrientation: 'from-image' });
    } catch (_) {
      return createImageBitmap(file);
    }
  }

  return new Promise((resolve, reject) => {
    const url = URL.createObjectURL(file);
    const image = new Image();
    image.onload = () => {
      URL.revokeObjectURL(url);
      resolve(image);
    };
    image.onerror = () => {
      URL.revokeObjectURL(url);
      reject(new Error('The selected file could not be decoded as an image.'));
    };
    image.src = url;
  });
}

function resetResults() {
  els.results.classList.add('hidden');
  els.emptyResult.classList.remove('hidden');
  els.topThree.replaceChildren();
  els.probabilityList.replaceChildren();
}

function clearImage() {
  if (state.bitmap && typeof state.bitmap.close === 'function') state.bitmap.close();
  if (state.objectUrl) URL.revokeObjectURL(state.objectUrl);
  state.bitmap = null;
  state.objectUrl = null;
  state.file = null;
  els.fileInput.value = '';
  els.imagePreview.removeAttribute('src');
  els.previewWrap.classList.add('hidden');
  els.dropZone.classList.remove('hidden');
  els.fileName.textContent = '';
  els.fileDetails.textContent = '';
  els.inputMessage.textContent = '';
  els.inputMessage.classList.remove('error');
  resetResults();
  updatePredictButton();
}

async function handleFile(file) {
  if (!file) return;
  if (!file.type.startsWith('image/')) {
    els.inputMessage.textContent = 'Select a valid image file.';
    els.inputMessage.classList.add('error');
    return;
  }

  try {
    els.inputMessage.textContent = 'Decoding image…';
    els.inputMessage.classList.remove('error');

    if (state.bitmap && typeof state.bitmap.close === 'function') state.bitmap.close();
    if (state.objectUrl) URL.revokeObjectURL(state.objectUrl);

    state.bitmap = await decodeImage(file);
    state.file = file;
    state.objectUrl = URL.createObjectURL(file);

    els.imagePreview.src = state.objectUrl;
    els.fileName.textContent = file.name;
    els.fileDetails.textContent = `${state.bitmap.width} × ${state.bitmap.height} · ${formatBytes(file.size)}`;
    els.dropZone.classList.add('hidden');
    els.previewWrap.classList.remove('hidden');
    els.inputMessage.textContent = state.modelReady ? 'Image ready for classification.' : 'Image ready. Waiting for the model to finish loading.';
    resetResults();
    updatePredictButton();
  } catch (error) {
    els.inputMessage.textContent = error.message || 'Could not read this image.';
    els.inputMessage.classList.add('error');
    clearImage();
  }
}

function preprocessImage(bitmap) {
  const canvas = els.canvas;
  const context = canvas.getContext('2d', { willReadFrequently: true });
  canvas.width = APP.imageSize;
  canvas.height = APP.imageSize;
  context.clearRect(0, 0, APP.imageSize, APP.imageSize);
  context.imageSmoothingEnabled = true;
  context.imageSmoothingQuality = 'high';
  context.drawImage(bitmap, 0, 0, APP.imageSize, APP.imageSize);

  const rgba = context.getImageData(0, 0, APP.imageSize, APP.imageSize).data;
  const planeSize = APP.imageSize * APP.imageSize;
  const nchw = new Float32Array(3 * planeSize);

  for (let pixel = 0; pixel < planeSize; pixel += 1) {
    const source = pixel * 4;
    nchw[pixel] = ((rgba[source] / 255) - APP.mean[0]) / APP.std[0];
    nchw[planeSize + pixel] = ((rgba[source + 1] / 255) - APP.mean[1]) / APP.std[1];
    nchw[(2 * planeSize) + pixel] = ((rgba[source + 2] / 255) - APP.mean[2]) / APP.std[2];
  }

  return new ort.Tensor('float32', nchw, [1, 3, APP.imageSize, APP.imageSize]);
}

function softmax(logits) {
  const maximum = Math.max(...logits);
  const exponentials = logits.map(value => Math.exp(value - maximum));
  const denominator = exponentials.reduce((sum, value) => sum + value, 0);
  return exponentials.map(value => value / denominator);
}

function renderResults(probabilities, timings) {
  const ranked = probabilities
    .map((probability, index) => ({ label: state.labels[index], probability }))
    .sort((a, b) => b.probability - a.probability);

  const winner = ranked[0];
  const confidencePercent = winner.probability * 100;

  els.predictedClass.textContent = titleCase(winner.label);
  els.confidenceValue.textContent = `${confidencePercent.toFixed(1)}%`;
  els.confidenceRing.style.setProperty('--confidence', `${winner.probability * 360}deg`);
  els.confidenceRing.setAttribute('aria-label', `${confidencePercent.toFixed(1)} percent confidence`);

  els.topThree.replaceChildren();
  for (const item of ranked.slice(0, 3)) {
    const row = document.createElement('li');
    const name = document.createElement('span');
    const value = document.createElement('span');
    name.className = 'class-name';
    value.className = 'class-prob';
    name.textContent = item.label;
    value.textContent = `${(item.probability * 100).toFixed(2)}%`;
    row.append(name, value);
    els.topThree.appendChild(row);
  }

  els.probabilityList.replaceChildren();
  for (const item of ranked) {
    const row = document.createElement('div');
    row.className = 'probability-row';
    const name = document.createElement('span');
    name.className = 'name';
    name.textContent = item.label;
    const track = document.createElement('div');
    track.className = 'bar-track';
    const fill = document.createElement('div');
    fill.className = 'bar-fill';
    fill.style.width = `${Math.max(item.probability * 100, 0.3)}%`;
    track.appendChild(fill);
    const value = document.createElement('span');
    value.className = 'value';
    value.textContent = `${(item.probability * 100).toFixed(2)}%`;
    row.append(name, track, value);
    els.probabilityList.appendChild(row);
  }

  els.preprocessTime.textContent = formatMilliseconds(timings.preprocess);
  els.inferenceTime.textContent = formatMilliseconds(timings.inference);
  els.totalTime.textContent = formatMilliseconds(timings.total);
  els.emptyResult.classList.add('hidden');
  els.results.classList.remove('hidden');
}

async function predict() {
  if (!state.modelReady || !state.bitmap || state.predicting) return;

  state.predicting = true;
  updatePredictButton();
  els.predictButtonText.textContent = 'Classifying…';
  els.buttonSpinner.classList.remove('hidden');
  els.inputMessage.textContent = 'Running browser-side inference…';
  els.inputMessage.classList.remove('error');

  try {
    const totalStart = performance.now();
    const preprocessStart = performance.now();
    const inputTensor = preprocessImage(state.bitmap);
    const preprocessEnd = performance.now();

    const inferenceStart = performance.now();
    const outputs = await state.session.run({ input: inputTensor });
    const inferenceEnd = performance.now();

    const outputTensor = outputs.logits;
    if (!outputTensor || !outputTensor.data) {
      throw new Error('The ONNX runtime did not return the expected logits output.');
    }
    const logits = Array.from(outputTensor.data);
    if (logits.length !== state.labels.length) {
      throw new Error(`Expected ${state.labels.length} logits but received ${logits.length}.`);
    }

    const probabilities = softmax(logits);
    const totalEnd = performance.now();
    renderResults(probabilities, {
      preprocess: preprocessEnd - preprocessStart,
      inference: inferenceEnd - inferenceStart,
      total: totalEnd - totalStart,
    });
    els.inputMessage.textContent = 'Prediction completed locally in this browser.';
  } catch (error) {
    console.error(error);
    els.inputMessage.textContent = error.message || 'Inference failed.';
    els.inputMessage.classList.add('error');
  } finally {
    state.predicting = false;
    els.predictButtonText.textContent = 'Classify image';
    els.buttonSpinner.classList.add('hidden');
    updatePredictButton();
  }
}

function preventDefaults(event) {
  event.preventDefault();
  event.stopPropagation();
}

['dragenter', 'dragover'].forEach(eventName => {
  els.dropZone.addEventListener(eventName, event => {
    preventDefaults(event);
    els.dropZone.classList.add('dragging');
  });
});

['dragleave', 'drop'].forEach(eventName => {
  els.dropZone.addEventListener(eventName, event => {
    preventDefaults(event);
    els.dropZone.classList.remove('dragging');
  });
});

els.dropZone.addEventListener('drop', event => handleFile(event.dataTransfer.files[0]));
els.dropZone.addEventListener('keydown', event => {
  if (event.key === 'Enter' || event.key === ' ') {
    event.preventDefault();
    els.fileInput.click();
  }
});
els.fileInput.addEventListener('change', event => handleFile(event.target.files[0]));
els.clearButton.addEventListener('click', clearImage);
els.predictButton.addEventListener('click', predict);
window.addEventListener('beforeunload', () => {
  if (state.bitmap && typeof state.bitmap.close === 'function') state.bitmap.close();
  if (state.objectUrl) URL.revokeObjectURL(state.objectUrl);
});

initializeModel();