import { MODEL_OPTIONS, classifyImage, getModelOption } from './classifier.js';

const imageInput = document.querySelector('#image-input');
const modelSelect = document.querySelector('#model-select');
const modelNote = document.querySelector('#model-note');
const clipControls = document.querySelector('#clip-controls');
const candidateLabels = document.querySelector('#candidate-labels');
const topK = document.querySelector('#top-k');
const topKValue = document.querySelector('#top-k-value');
const runButton = document.querySelector('#run-button');
const statusEl = document.querySelector('#status');
const progressWrap = document.querySelector('#progress-wrap');
const progressBar = document.querySelector('#progress-bar');
const previewWrap = document.querySelector('#preview-wrap');
const previewImage = document.querySelector('#preview-image');
const summary = document.querySelector('#summary');
const tbody = document.querySelector('#predictions-body');
const runtimeBadge = document.querySelector('#runtime-badge');
const sampleGrid = document.querySelector('#sample-grid');
const shuffleButton = document.querySelector('#shuffle-examples');
const runMeta = document.querySelector('#run-meta');
const metaModel = document.querySelector('#meta-model');
const metaImage = document.querySelector('#meta-image');
const metaRuntime = document.querySelector('#meta-runtime');
const historyBody = document.querySelector('#history-body');
const clearHistory = document.querySelector('#clear-history');

const state = { dataUrl: null, sourceLabel: null, running: false, history: [] };

const EXAMPLES = [
  { src: './examples/chelsea.png', label: 'Cat' },
  { src: './examples/coffee.png', label: 'Coffee' },
  { src: './examples/rocket.jpg', label: 'Rocket' },
  { src: './examples/astronaut.png', label: 'Astronaut' },
  { src: './examples/cheetah.jpg', label: 'Cheetah' },
  { src: './examples/lion.jpg', label: 'Lion' },
  { src: './examples/tower.jpg', label: 'Tower' },
  { src: './examples/bus.png', label: 'Bus illustration' },
  { src: './examples/dog.jpg', label: 'Dog' },
  { src: './examples/horse.jpg', label: 'Horse' },
  { src: './examples/elephant.jpg', label: 'Elephant' },
  { src: './examples/rabbit.jpg', label: 'Rabbit' },
  { src: './examples/bird.jpg', label: 'Parrot' },
  { src: './examples/cow.jpg', label: 'Cow' },
  { src: './examples/airplane.jpg', label: 'Airplane' },
  { src: './examples/car.jpg', label: 'Car' },
  { src: './examples/bicycle.jpg', label: 'Bicycle' },
  { src: './examples/train.jpg', label: 'Train' },
  { src: './examples/motorcycle.jpg', label: 'Motorcycle' },
  { src: './examples/boat.jpg', label: 'Boat' },
  { src: './examples/truck.jpg', label: 'Truck' },
  { src: './examples/pizza.jpg', label: 'Pizza' },
  { src: './examples/hamburger.jpg', label: 'Hamburger' },
  { src: './examples/apple.jpg', label: 'Apple' },
  { src: './examples/banana.jpg', label: 'Banana' },
  { src: './examples/orange.jpg', label: 'Orange' },
  { src: './examples/sandwich.jpg', label: 'Sandwich' },
  { src: './examples/bread.jpg', label: 'Bread' },
  { src: './examples/strawberry.jpg', label: 'Strawberry' },
  { src: './examples/ice_cream.jpg', label: 'Ice cream' },
  { src: './examples/tiger.jpg', label: 'Tiger' },
  { src: './examples/zebra.jpg', label: 'Zebra' },
  { src: './examples/giraffe.jpg', label: 'Giraffe' },
  { src: './examples/panda.jpg', label: 'Giant panda' },
  { src: './examples/koala.jpg', label: 'Koala' },
  { src: './examples/kangaroo.jpg', label: 'Kangaroo' },
  { src: './examples/camel.jpg', label: 'Camel' },
  { src: './examples/goat.jpg', label: 'Goat' },
  { src: './examples/sheep.jpg', label: 'Sheep' },
  { src: './examples/pig.jpg', label: 'Pig' },
  { src: './examples/duck.jpg', label: 'Duck' },
  { src: './examples/chicken.jpg', label: 'Chicken' },
  { src: './examples/flamingo.jpg', label: 'Flamingo' },
  { src: './examples/penguin.jpg', label: 'Penguin' },
  { src: './examples/owl.jpg', label: 'Owl' },
  { src: './examples/eagle.jpg', label: 'Eagle' },
  { src: './examples/butterfly.jpg', label: 'Butterfly' },
  { src: './examples/bee.jpg', label: 'Bee' },
  { src: './examples/turtle.jpg', label: 'Turtle' },
  { src: './examples/frog.jpg', label: 'Frog' },
  { src: './examples/city_bus.jpg', label: 'City bus' },
  { src: './examples/taxi.jpg', label: 'Taxi' },
  { src: './examples/tram.jpg', label: 'Tram' },
  { src: './examples/subway.jpg', label: 'Subway train' },
  { src: './examples/helicopter.jpg', label: 'Helicopter' },
  { src: './examples/sailboat.jpg', label: 'Sailboat' },
  { src: './examples/ship.jpg', label: 'Ship' },
  { src: './examples/canoe.jpg', label: 'Canoe' },
  { src: './examples/tractor.jpg', label: 'Tractor' },
  { src: './examples/fire_truck.jpg', label: 'Fire truck' },
  { src: './examples/ambulance.jpg', label: 'Ambulance' },
  { src: './examples/police_car.jpg', label: 'Police car' },
  { src: './examples/van.jpg', label: 'Van' },
  { src: './examples/scooter.jpg', label: 'Scooter' },
  { src: './examples/hot_air_balloon.jpg', label: 'Hot-air balloon' },
  { src: './examples/excavator.jpg', label: 'Excavator' },
  { src: './examples/bulldozer.jpg', label: 'Bulldozer' },
  { src: './examples/forklift.jpg', label: 'Forklift' },
  { src: './examples/jet.jpg', label: 'Jet aircraft' },
  { src: './examples/cable_car.jpg', label: 'Cable car' },
  { src: './examples/pear.jpg', label: 'Pear' },
  { src: './examples/peach.jpg', label: 'Peach' },
  { src: './examples/pineapple.jpg', label: 'Pineapple' },
  { src: './examples/watermelon.jpg', label: 'Watermelon' },
  { src: './examples/grapes.jpg', label: 'Grapes' },
  { src: './examples/mango.jpg', label: 'Mango' },
  { src: './examples/lemon.jpg', label: 'Lemon' },
  { src: './examples/avocado.jpg', label: 'Avocado' },
  { src: './examples/carrot.jpg', label: 'Carrot' },
  { src: './examples/broccoli.jpg', label: 'Broccoli' },
  { src: './examples/tomato.jpg', label: 'Tomato' },
  { src: './examples/corn.jpg', label: 'Corn' },
  { src: './examples/donut.jpg', label: 'Doughnut' },
  { src: './examples/cookie.jpg', label: 'Cookie' },
  { src: './examples/cupcake.jpg', label: 'Cupcake' },
  { src: './examples/sushi.jpg', label: 'Sushi' },
  { src: './examples/pasta.jpg', label: 'Pasta' },
  { src: './examples/salad.jpg', label: 'Salad' },
  { src: './examples/popcorn.jpg', label: 'Popcorn' },
  { src: './examples/egg.jpg', label: 'Egg' },
  { src: './examples/laptop.jpg', label: 'Laptop' },
  { src: './examples/camera.jpg', label: 'Camera' },
  { src: './examples/guitar.jpg', label: 'Guitar' },
  { src: './examples/piano.jpg', label: 'Piano' },
  { src: './examples/clock.jpg', label: 'Clock' },
  { src: './examples/chair.jpg', label: 'Chair' },
  { src: './examples/lamp.jpg', label: 'Lamp' },
  { src: './examples/backpack.jpg', label: 'Backpack' },
  { src: './examples/umbrella.jpg', label: 'Umbrella' },
  { src: './examples/book.jpg', label: 'Book' },
];

function chooseRandom(items, count) {
  return [...items]
    .map((item) => ({ item, sort: Math.random() }))
    .sort((a, b) => a.sort - b.sort)
    .slice(0, count)
    .map(({ item }) => item);
}

function resetResults() {
  summary.textContent = 'No predictions yet.';
  tbody.innerHTML = '<tr class="empty-row"><td colspan="3">Run the classifier to see details.</td></tr>';
  runMeta.hidden = true;
}

function setBadge(text, mode = '') {
  runtimeBadge.textContent = text;
  runtimeBadge.className = `badge ${mode}`.trim();
}

function setRunning(running) {
  state.running = running;
  runButton.disabled = running || !state.dataUrl;
  modelSelect.disabled = running;
  imageInput.disabled = running;
  topK.disabled = running;
  candidateLabels.disabled = running;
  shuffleButton.disabled = running;
  document.querySelectorAll('.sample-card').forEach((card) => { card.disabled = running; });
}

function setSelectedImage(dataUrl, label, readyMessage) {
  const image = new Image();
  image.onload = () => {
    state.dataUrl = dataUrl;
    state.sourceLabel = label;
    previewImage.src = dataUrl;
    previewImage.alt = `${label} selected for classification`;
    previewWrap.classList.remove('empty');
    setRunning(false);
    statusEl.textContent = readyMessage;
    setBadge('Image ready');
    resetResults();
  };
  image.onerror = () => {
    statusEl.textContent = 'Could not read that image.';
    setBadge('Image error', 'error');
  };
  image.src = dataUrl;
}

function blobToDataUrl(blob) {
  return new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.onload = () => resolve(reader.result);
    reader.onerror = reject;
    reader.readAsDataURL(blob);
  });
}

async function loadExample(example, button) {
  try {
    statusEl.textContent = `Loading example: ${example.label}…`;
    setBadge('Loading sample', 'running');
    const response = await fetch(example.src);
    if (!response.ok) throw new Error(`Could not load ${example.src}`);
    const dataUrl = await blobToDataUrl(await response.blob());
    imageInput.value = '';
    document.querySelectorAll('.sample-card').forEach((card) => card.classList.remove('selected'));
    button.classList.add('selected');
    setSelectedImage(dataUrl, example.label, `${example.label} example ready. Click Run AI Classification.`);
  } catch (error) {
    console.error(error);
    statusEl.textContent = 'Could not load that example. Try another image.';
    setBadge('Sample error', 'error');
  }
}

function renderModelChoices() {
  for (const option of MODEL_OPTIONS) {
    const item = document.createElement('option');
    item.value = option.key;
    item.textContent = option.label;
    modelSelect.appendChild(item);
  }
  modelSelect.value = MODEL_OPTIONS[0].key;
  modelNote.textContent = MODEL_OPTIONS[0].note;
}

function renderExampleChoices() {
  sampleGrid.innerHTML = '';
  for (const example of chooseRandom(EXAMPLES, 4)) {
    const button = document.createElement('button');
    button.type = 'button';
    button.className = 'sample-card';
    button.title = `Use sample: ${example.label}`;

    const image = document.createElement('img');
    image.src = example.src;
    image.alt = example.label;
    image.loading = 'lazy';

    const label = document.createElement('span');
    label.textContent = example.label;

    const source = document.createElement('small');
    source.textContent = 'Tested teaching sample';

    button.append(image, label, source);
    button.addEventListener('click', () => loadExample(example, button));
    sampleGrid.appendChild(button);
  }
}

function setProgress(data) {
  let percent = null;
  if (typeof data?.progress === 'number') percent = data.progress;
  else if (data?.total && typeof data?.loaded === 'number') percent = (data.loaded / data.total) * 100;

  if (percent !== null && Number.isFinite(percent)) {
    progressWrap.hidden = false;
    progressBar.style.width = `${Math.max(0, Math.min(100, percent))}%`;
  }
  if (data?.status === 'progress' && data?.file) {
    statusEl.textContent = `Downloading model: ${data.file}`;
  }
}

function parseCandidateLabels() {
  return candidateLabels.value.split(/[,\n]/).map((label) => label.trim()).filter(Boolean);
}

function renderPredictions(predictions) {
  tbody.innerHTML = '';
  const option = getModelOption(modelSelect.value);
  const scoreContext = option.task === 'zero-shot-image-classification'
    ? 'Scores compare only the candidate labels supplied above.'
    : 'Scores are calculated over the 1,000 ImageNet classes.';

  predictions.forEach((prediction, index) => {
    const row = document.createElement('tr');
    const rankCell = document.createElement('td');
    rankCell.textContent = index + 1;

    const labelCell = document.createElement('td');
    labelCell.textContent = prediction.label;

    const score = Math.max(0, Math.min(1, Number(prediction.score)));
    const scoreCell = document.createElement('td');
    const scoreWrap = document.createElement('div');
    scoreWrap.className = 'score-wrap';
    const scoreText = document.createElement('span');
    scoreText.textContent = `${(score * 100).toFixed(2)}%`;
    const scoreTrack = document.createElement('span');
    scoreTrack.className = 'score-track';
    const scoreFill = document.createElement('span');
    scoreFill.className = 'score-fill';
    scoreFill.style.width = `${score * 100}%`;
    scoreTrack.appendChild(scoreFill);
    scoreWrap.append(scoreText, scoreTrack);
    scoreCell.appendChild(scoreWrap);

    row.append(rankCell, labelCell, scoreCell);
    tbody.appendChild(row);
  });

  const best = predictions[0];
  summary.textContent = best
    ? `${option.label} predicts “${best.label}” at ${(Number(best.score) * 100).toFixed(1)}%. ${scoreContext}`
    : `${option.label} returned no predictions.`;
}

function renderHistory() {
  clearHistory.disabled = state.history.length === 0;
  if (!state.history.length) {
    historyBody.innerHTML = '<tr class="empty-row"><td colspan="6">Run two models on the same image to start a comparison.</td></tr>';
    return;
  }

  historyBody.innerHTML = '';
  state.history.forEach((item, index) => {
    const row = document.createElement('tr');
    const values = [
      index + 1,
      item.image,
      item.model,
      item.label,
      `${(item.score * 100).toFixed(1)}%`,
      `${item.seconds.toFixed(1)}s`,
    ];
    values.forEach((value) => {
      const cell = document.createElement('td');
      cell.textContent = value;
      row.appendChild(cell);
    });
    historyBody.appendChild(row);
  });
}

function recordRun(option, predictions, seconds) {
  const best = predictions[0];
  if (!best) return;

  state.history.unshift({
    image: state.sourceLabel || 'Selected image',
    model: option.label,
    label: best.label,
    score: Number(best.score),
    seconds,
  });
  state.history = state.history.slice(0, 12);
  renderHistory();

  metaModel.textContent = option.label;
  metaImage.textContent = state.sourceLabel || 'Selected image';
  metaRuntime.textContent = `${seconds.toFixed(1)}s`;
  runMeta.hidden = false;
}

topK.addEventListener('input', () => { topKValue.value = topK.value; });

modelSelect.addEventListener('change', () => {
  const option = getModelOption(modelSelect.value);
  modelNote.textContent = option.note;
  clipControls.hidden = option.task !== 'zero-shot-image-classification';
  setBadge(`Model: ${option.label}`);
  resetResults();
  if (state.dataUrl) statusEl.textContent = `Model changed to ${option.label}. Run again to compare predictions.`;
});

imageInput.addEventListener('change', () => {
  const file = imageInput.files?.[0];
  if (!file) return;
  if (!file.type.startsWith('image/')) {
    statusEl.textContent = 'Please choose an image file.';
    setBadge('Invalid file', 'error');
    return;
  }
  if (file.size > 12 * 1024 * 1024) {
    statusEl.textContent = 'Please choose an image smaller than 12 MB.';
    setBadge('Image too large', 'error');
    return;
  }

  document.querySelectorAll('.sample-card').forEach((card) => card.classList.remove('selected'));
  const reader = new FileReader();
  reader.onload = () => setSelectedImage(reader.result, 'Uploaded image', 'Your image is ready. Click Run AI Classification.');
  reader.onerror = () => {
    statusEl.textContent = 'Could not read that image.';
    setBadge('Image error', 'error');
  };
  reader.readAsDataURL(file);
});

clearHistory.addEventListener('click', () => {
  state.history = [];
  renderHistory();
});

shuffleButton.addEventListener('click', () => {
  renderExampleChoices();
  statusEl.textContent = state.dataUrl
    ? 'New examples shown. Your current image is still selected.'
    : 'Choose an example, or upload your own image.';
});

runButton.addEventListener('click', async () => {
  if (!state.dataUrl || state.running) return;

  const option = getModelOption(modelSelect.value);
  const labels = parseCandidateLabels();
  if (option.task === 'zero-shot-image-classification' && (labels.length < 2 || labels.length > 20)) {
    statusEl.textContent = 'CLIP needs 2–20 candidate labels.';
    setBadge('Check labels', 'error');
    return;
  }

  setRunning(true);
  setBadge(`Loading ${option.label}`, 'running');
  progressWrap.hidden = false;
  progressBar.style.width = '4%';
  const started = performance.now();

  try {
    statusEl.textContent = `Preparing ${option.label}. The first run may download model files.`;
    const predictions = await classifyImage(state.dataUrl, {
      modelKey: option.key,
      topK: Number(topK.value),
      candidateLabels: labels,
      progressCallback: setProgress,
    });
    renderPredictions(predictions);
    const seconds = (performance.now() - started) / 1000;
    recordRun(option, predictions, seconds);
    statusEl.textContent = `Done in ${seconds.toFixed(1)}s with ${option.label}. Change the model or image to compare.`;
    setBadge(`${predictions.length} predictions`, 'ready');
  } catch (error) {
    console.error(error);
    statusEl.textContent = `${option.label} failed to run. Check the network connection or try ResNet-18.`;
    setBadge('Error', 'error');
  } finally {
    progressWrap.hidden = true;
    setRunning(false);
  }
});

renderModelChoices();
renderExampleChoices();
renderHistory();