// Use global ort loaded via CDN bundle
const { InferenceSession, Tensor, env } = window.ort;

// Configure WASM fallback paths to CDN (no local wasm binaries required)
env.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/';
env.wasm.simd = true;

// ── COCO Class Names (80 standard classes) ───────────────────────────────────
const COCO = [
  'person','bicycle','car','motorcycle','airplane','bus','train','truck','boat',
  'traffic light','fire hydrant','stop sign','parking meter','bench','bird','cat',
  'dog','horse','sheep','cow','elephant','bear','zebra','giraffe','backpack',
  'umbrella','handbag','tie','suitcase','frisbee','skis','snowboard','sports ball',
  'kite','baseball bat','baseball glove','skateboard','surfboard','tennis racket',
  'bottle','wine glass','cup','fork','knife','spoon','bowl','banana','apple',
  'sandwich','orange','broccoli','carrot','hot dog','pizza','donut','cake','chair',
  'couch','potted plant','bed','dining table','toilet','tv','laptop','mouse',
  'remote','keyboard','cell phone','microwave','oven','toaster','sink','refrigerator',
  'book','clock','vase','scissors','teddy bear','hair drier','toothbrush'
];

// ── DOM References ───────────────────────────────────────────────────────────
const modelPreset = document.getElementById('model-preset');
const devicePreset = document.getElementById('device-preset');
const engineBadge = document.getElementById('engine-badge');
const badgeDot = document.getElementById('badge-dot');
const badgeText = document.getElementById('badge-text');

const loaderBox = document.getElementById('loader-box');
const loaderStatus = document.getElementById('loader-status');
const loaderPercent = document.getElementById('loader-percent');
const loaderBar = document.getElementById('loader-bar');

const confSlider = document.getElementById('conf-slider');
const confVal = document.getElementById('conf-val');
const personFilter = document.getElementById('person-filter');

const metricLatency = document.getElementById('metric-latency');
const metricFps = document.getElementById('metric-fps');
const metricCount = document.getElementById('metric-count');

const tabWebcam = document.getElementById('tab-webcam');
const tabMedia = document.getElementById('tab-media');

const displayCanvas = document.getElementById('display-canvas');
const ctx = displayCanvas.getContext('2d');

const liveVideo = document.getElementById('live-video');
const fileVideo = document.getElementById('file-video');
const fileImage = document.getElementById('file-image');

const camIdleOverlay = document.getElementById('cam-idle-overlay');
const mediaIdleOverlay = document.getElementById('media-idle-overlay');
const startCamBtn = document.getElementById('start-cam-btn');
const stopCamBtn = document.getElementById('stop-cam-btn');

const mediaFileInput = document.getElementById('media-file-input');
const mediaFileInput2 = document.getElementById('media-file-input-2');
const webcamControls = document.getElementById('webcam-controls');
const mediaControls = document.getElementById('media-controls');
const videoPlayBtn = document.getElementById('video-play-btn');
const toastEl = document.getElementById('toast');

// ── State ────────────────────────────────────────────────────────────────────
let session = null;
let activeArchitecture = 'yolo26'; // 'yolo26' (NMS-free [1,300,6]) or 'yolo11' (anchor-free [1,84,N])
let inputSize = 736;               // Resolution exported from notebook
let inTensorName = '';
let outTensorName = '';
let currentMode = 'webcam';        // 'webcam' | 'media'
let isWebcamActive = false;
let isVideoActive = false;
let currentMediaType = null;       // 'image' | 'video'
let animFrameId = null;
let lastTime = performance.now();
let frameCount = 0;

// Preprocessing offscreen canvas
let offscreenCanvas = new OffscreenCanvas(inputSize, inputSize);
let offscreenCtx = offscreenCanvas.getContext('2d', { willReadFrequently: true });

function updateOffscreenSize(sz) {
  if (sz && sz !== inputSize) {
    inputSize = sz;
    offscreenCanvas = new OffscreenCanvas(inputSize, inputSize);
    offscreenCtx = offscreenCanvas.getContext('2d', { willReadFrequently: true });
  }
}

function showToast(msg) {
  toastEl.textContent = msg;
  toastEl.classList.add('show');
  setTimeout(() => toastEl.classList.remove('show'), 3000);
}

// ── Model Downloader & Loader ────────────────────────────────────────────────
async function loadSelectedModel() {
  const modelName = modelPreset.value;
  const isWebGPU = devicePreset.value === 'webgpu';
  
  badgeDot.className = 'badge-dot loading';
  badgeText.textContent = 'Compiling...';
  loaderBox.style.display = 'block';
  loaderStatus.textContent = `Streaming ${modelName}...`;
  loaderBar.style.width = '10%';
  loaderPercent.textContent = '10%';

  try {
    // Path resolution: try local directory first, then ../models/
    let modelPath = `./${modelName}.onnx`;
    let res = await fetch(modelPath);
    if (!res.ok) {
      modelPath = `../models/${modelName}.onnx`;
      res = await fetch(modelPath);
    }
    if (!res.ok) throw new Error(`HTTP ${res.status}: File not found at ${modelPath}`);

    const contentLength = +res.headers.get('content-length') || 0;
    const reader = res.body.getReader();
    const chunks = [];
    let received = 0;

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
      received += value.length;
      if (contentLength > 0) {
        const pct = Math.round((received / contentLength) * 80);
        loaderBar.style.width = `${pct}%`;
        loaderPercent.textContent = `${pct}%`;
        loaderStatus.textContent = `Downloaded ${(received / 1024 / 1024).toFixed(1)} MB`;
      }
    }

    loaderStatus.textContent = 'Allocating GPU Buffers...';
    loaderBar.style.width = '90%';
    loaderPercent.textContent = '90%';

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

    // Check if WebGPU adapter is actually available
    let useWebGPU = isWebGPU;
    if (useWebGPU) {
      try {
        if (!navigator.gpu || !(await navigator.gpu.requestAdapter())) {
          useWebGPU = false;
        }
      } catch (e) {
        useWebGPU = false;
      }
    }

    const providers = useWebGPU ? ['webgpu', 'wasm'] : ['wasm'];
    try {
      session = await InferenceSession.create(modelBuffer.buffer, {
        executionProviders: providers,
        graphOptimizationLevel: 'all'
      });
    } catch (createErr) {
      console.warn('Primary EP init fallback to WASM:', createErr);
      session = await InferenceSession.create(modelBuffer.buffer, {
        executionProviders: ['wasm'],
        graphOptimizationLevel: 'all'
      });
      useWebGPU = false;
    }

    inTensorName = session.inputNames[0];
    outTensorName = session.outputNames[0];

    // Determine architecture type & dimensions
    const isYOLO26 = modelName.startsWith('yolo26');
    activeArchitecture = isYOLO26 ? 'yolo26' : 'yolo11';

    // Adapt to model input dimensions if known
    try {
      const inputMeta = session._inputsMeta || {};
      const dims = inputMeta[inTensorName]?.dimensions;
      if (dims && dims[2] && typeof dims[2] === 'number') {
        updateOffscreenSize(dims[2]);
      }
    } catch (e) {}

    // Warmup probe
    try {
      const dummy = new Tensor('float32', new Float32Array(1 * 3 * inputSize * inputSize), [1, 3, inputSize, inputSize]);
      await session.run({ [inTensorName]: dummy });
    } catch (probeErr) {
      console.warn('Probe notice:', probeErr);
    }

    badgeDot.className = 'badge-dot';
    badgeText.textContent = isWebGPU ? 'WebGPU Active' : 'WASM CPU';
    loaderBar.style.width = '100%';
    loaderPercent.textContent = '100%';
    loaderStatus.textContent = 'Neural Engine Ready';
    setTimeout(() => { loaderBox.style.display = 'none'; }, 900);
    showToast(`Loaded ${modelName.toUpperCase()} on ${isWebGPU ? 'WebGPU' : 'CPU'}`);

    if (currentMode === 'media' && currentMediaType === 'image' && fileImage.src) {
      detectAndRender(fileImage);
    }
  } catch (err) {
    console.error('Model load error:', err);
    badgeDot.className = 'badge-dot loading';
    badgeText.textContent = 'Load Failed';
    loaderStatus.textContent = 'Model not available';
    showToast(`Notice: Run notebook export to generate ${modelName}.onnx`);
  }
}

// ── Preprocessing (Letterbox with Aspect-Ratio Padding) ───────────────────────
function preprocess(source) {
  const sw = source.videoWidth || source.naturalWidth || source.width || inputSize;
  const sh = source.videoHeight || source.naturalHeight || source.height || inputSize;
  const scale = Math.min(inputSize / sw, inputSize / sh);
  const nw = Math.round(sw * scale);
  const nh = Math.round(sh * scale);
  const dx = Math.floor((inputSize - nw) / 2);
  const dy = Math.floor((inputSize - nh) / 2);

  offscreenCtx.fillStyle = 'rgb(114, 114, 114)';
  offscreenCtx.fillRect(0, 0, inputSize, inputSize);
  offscreenCtx.drawImage(source, dx, dy, nw, nh);

  const imgData = offscreenCtx.getImageData(0, 0, inputSize, inputSize);
  const data = imgData.data;
  const channelSize = inputSize * inputSize;
  const floatData = new Float32Array(3 * channelSize);

  for (let i = 0; i < channelSize; i++) {
    floatData[i] = data[4 * i] / 255.0;                         // R
    floatData[channelSize + i] = data[4 * i + 1] / 255.0;       // G
    floatData[2 * channelSize + i] = data[4 * i + 2] / 255.0;   // B
  }

  return {
    tensor: new Tensor('float32', floatData, [1, 3, inputSize, inputSize]),
    scale,
    dx,
    dy
  };
}

// ── Greedy NMS for YOLO11 Anchor-Free Head ────────────────────────────────────
function runNms(boxes, scores, iouThreshold = 0.45) {
  const indices = scores.map((_, i) => i).sort((a, b) => scores[b] - scores[a]);
  const selected = [];
  const suppressed = new Uint8Array(indices.length);

  for (let i = 0; i < indices.length; i++) {
    if (suppressed[i]) continue;
    selected.push(indices[i]);
    const [x1A, y1A, x2A, y2A] = boxes[indices[i]];
    const areaA = (x2A - x1A) * (y2A - y1A);

    for (let j = i + 1; j < indices.length; j++) {
      if (suppressed[j]) continue;
      const [x1B, y1B, x2B, y2B] = boxes[indices[j]];
      const interX = Math.max(0, Math.min(x2A, x2B) - Math.max(x1A, x1B));
      const interY = Math.max(0, Math.min(y2A, y2B) - Math.max(y1A, y1B));
      const interArea = interX * interY;
      const unionArea = areaA + ((x2B - x1B) * (y2B - y1B)) - interArea;
      if (interArea / (unionArea + 1e-6) > iouThreshold) {
        suppressed[j] = 1;
      }
    }
  }
  return selected;
}

// ── Apple Monotone Drawing Renderer ──────────────────────────────────────────
function renderDetections(detections) {
  const count = detections.length;
  metricCount.textContent = count;

  ctx.lineWidth = 1.5;
  ctx.font = '500 11px -apple-system, BlinkMacSystemFont, "SF Mono", monospace';

  for (const det of detections) {
    const { x1, y1, x2, y2, score, cls } = det;
    const w = x2 - x1;
    const h = y2 - y1;

    // Sleek monotone bounding box
    ctx.strokeStyle = 'rgba(0, 0, 0, 0.75)';
    ctx.strokeRect(x1, y1, w, h);

    // Apple corner accent marks
    const corner = Math.min(8, w / 4, h / 4);
    ctx.lineWidth = 2.5;
    ctx.strokeStyle = '#000000';
    ctx.beginPath();
    ctx.moveTo(x1, y1 + corner); ctx.lineTo(x1, y1); ctx.lineTo(x1 + corner, y1);
    ctx.moveTo(x2 - corner, y1); ctx.lineTo(x2, y1); ctx.lineTo(x2, y1 + corner);
    ctx.moveTo(x1, y2 - corner); ctx.lineTo(x1, y2); ctx.lineTo(x1 + corner, y2);
    ctx.moveTo(x2 - corner, y2); ctx.lineTo(x2, y2); ctx.lineTo(x2, y2 - corner);
    ctx.stroke();
    ctx.lineWidth = 1.5;

    // Label tag
    const name = COCO[cls] || `ID ${cls}`;
    const tagText = `${name} ${(score * 100).toFixed(0)}%`;
    const tagWidth = ctx.measureText(tagText).width + 8;
    const tagHeight = 16;
    const tagY = Math.max(0, y1 - tagHeight - 2);

    ctx.fillStyle = 'rgba(0, 0, 0, 0.9)';
    ctx.fillRect(x1, tagY, tagWidth, tagHeight);
    ctx.fillStyle = '#ffffff';
    ctx.fillText(tagText, x1 + 4, tagY + 12);
  }
}

// ── Inference Dispatcher ─────────────────────────────────────────────────────
async function detectAndRender(source) {
  if (!session) return;

  const tStart = performance.now();
  const { tensor, scale, dx, dy } = preprocess(source);

  // Synchronize canvas dimensions with source element
  const sw = source.videoWidth || source.naturalWidth || source.width || inputSize;
  const sh = source.videoHeight || source.naturalHeight || source.height || inputSize;
  if (displayCanvas.width !== sw || displayCanvas.height !== sh) {
    displayCanvas.width = sw;
    displayCanvas.height = sh;
  }

  const output = await session.run({ [inTensorName]: tensor });
  const latency = performance.now() - tStart;
  metricLatency.textContent = latency.toFixed(0);

  // Compute FPS
  frameCount++;
  const now = performance.now();
  if (now - lastTime >= 1000) {
    metricFps.textContent = ((frameCount * 1000) / (now - lastTime)).toFixed(0);
    frameCount = 0;
    lastTime = now;
  }

  // Draw source frame
  ctx.drawImage(source, 0, 0, sw, sh);

  const confThreshold = parseFloat(confSlider.value);
  const onlyPerson = personFilter.checked;
  const raw = output[outTensorName];
  const data = raw.data;
  const dims = raw.dims;
  const detections = [];

  if (activeArchitecture === 'yolo26' && dims[2] === 6) {
    // End-to-End NMS-Free format [1, 300, 6]
    for (let i = 0; i < dims[1]; i++) {
      const idx = i * 6;
      const score = data[idx + 4];
      if (score < confThreshold) continue;
      const cls = Math.round(data[idx + 5]);
      if (onlyPerson && cls !== 0) continue;

      detections.push({
        x1: (data[idx] - dx) / scale,
        y1: (data[idx + 1] - dy) / scale,
        x2: (data[idx + 2] - dx) / scale,
        y2: (data[idx + 3] - dy) / scale,
        score,
        cls
      });
    }
  } else {
    // YOLO11 Anchor-Free format [1, 84, N]
    const numCandidates = dims[2];
    const candidateBoxes = [];
    const candidateScores = [];
    const candidateClasses = [];

    for (let i = 0; i < numCandidates; i++) {
      let maxScore = -1;
      let maxCls = -1;

      // Classes start at offset 4
      for (let c = 0; c < 80; c++) {
        const score = data[(4 + c) * numCandidates + i];
        if (score > maxScore) {
          maxScore = score;
          maxCls = c;
        }
      }

      if (maxScore >= confThreshold) {
        if (onlyPerson && maxCls !== 0) continue;
        const xc = data[0 * numCandidates + i];
        const yc = data[1 * numCandidates + i];
        const w = data[2 * numCandidates + i];
        const h = data[3 * numCandidates + i];

        candidateBoxes.push([xc - w / 2, yc - h / 2, xc + w / 2, yc + h / 2]);
        candidateScores.push(maxScore);
        candidateClasses.push(maxCls);
      }
    }

    const nmsWinners = runNms(candidateBoxes, candidateScores, 0.45);
    for (const k of nmsWinners) {
      const [bx1, by1, bx2, by2] = candidateBoxes[k];
      detections.push({
        x1: (bx1 - dx) / scale,
        y1: (by1 - dy) / scale,
        x2: (bx2 - dx) / scale,
        y2: (by2 - dy) / scale,
        score: candidateScores[k],
        cls: candidateClasses[k]
      });
    }
  }

  renderDetections(detections);
}

// ── Loops ────────────────────────────────────────────────────────────────────
async function webcamLoop() {
  if (!isWebcamActive) return;
  if (liveVideo.readyState >= 2) {
    await detectAndRender(liveVideo);
  }
  animFrameId = requestAnimationFrame(webcamLoop);
}

async function videoLoop() {
  if (!isVideoActive || currentMediaType !== 'video') return;
  if (fileVideo.readyState >= 2 && !fileVideo.paused) {
    await detectAndRender(fileVideo);
  }
  requestAnimationFrame(videoLoop);
}

// ── Camera Controls ──────────────────────────────────────────────────────────
startCamBtn.onclick = async () => {
  try {
    const stream = await navigator.mediaDevices.getUserMedia({
      video: { width: { ideal: 1280 }, height: { ideal: 720 } }
    });
    liveVideo.srcObject = stream;
    await liveVideo.play();
    isWebcamActive = true;
    camIdleOverlay.style.display = 'none';
    stopCamBtn.disabled = false;
    startCamBtn.disabled = true;
    webcamLoop();
    showToast('FaceTime Camera Active');
  } catch (e) {
    alert('Camera permission denied or camera unavailable.');
  }
};

stopCamBtn.onclick = () => {
  isWebcamActive = false;
  if (liveVideo.srcObject) {
    liveVideo.srcObject.getTracks().forEach(t => t.stop());
    liveVideo.srcObject = null;
  }
  if (animFrameId) cancelAnimationFrame(animFrameId);
  camIdleOverlay.style.display = 'flex';
  stopCamBtn.disabled = true;
  startCamBtn.disabled = false;
  metricLatency.textContent = '--';
  metricFps.textContent = '--';
};

// ── Media File Handling ──────────────────────────────────────────────────────
function handleUploadedFile(file) {
  if (!file) return;
  const isVideo = file.type.startsWith('video');
  const fileUrl = URL.createObjectURL(file);

  mediaIdleOverlay.style.display = 'none';
  mediaControls.style.display = isVideo ? 'flex' : 'none';

  if (isVideo) {
    currentMediaType = 'video';
    fileImage.style.display = 'none';
    fileVideo.src = fileUrl;
    fileVideo.onloadedmetadata = () => {
      displayCanvas.width = fileVideo.videoWidth;
      displayCanvas.height = fileVideo.videoHeight;
      fileVideo.play();
      isVideoActive = true;
      videoPlayBtn.textContent = 'Pause Video';
      videoLoop();
    };
  } else {
    currentMediaType = 'image';
    isVideoActive = false;
    fileVideo.pause();
    fileVideo.src = '';
    fileImage.src = fileUrl;
    fileImage.onload = async () => {
      displayCanvas.width = fileImage.naturalWidth;
      displayCanvas.height = fileImage.naturalHeight;
      await detectAndRender(fileImage);
    };
  }
}

mediaFileInput.onchange = (e) => handleUploadedFile(e.target.files[0]);
mediaFileInput2.onchange = (e) => handleUploadedFile(e.target.files[0]);

videoPlayBtn.onclick = () => {
  if (fileVideo.paused) {
    fileVideo.play();
    isVideoActive = true;
    videoPlayBtn.textContent = 'Pause Video';
    videoLoop();
  } else {
    fileVideo.pause();
    isVideoActive = false;
    videoPlayBtn.textContent = 'Play Video';
  }
};

// ── Tab Navigation ───────────────────────────────────────────────────────────
function switchTab(mode) {
  currentMode = mode;
  if (mode === 'webcam') {
    tabWebcam.classList.add('active');
    tabMedia.classList.remove('active');
    webcamControls.style.display = 'flex';
    mediaControls.style.display = 'none';
    mediaIdleOverlay.style.display = 'none';
    if (!isWebcamActive) camIdleOverlay.style.display = 'flex';
    if (isVideoActive) {
      fileVideo.pause();
      isVideoActive = false;
    }
  } else {
    tabMedia.classList.add('active');
    tabWebcam.classList.remove('active');
    webcamControls.style.display = 'none';
    camIdleOverlay.style.setProperty('display', 'none', 'important');
    if (!fileImage.src && !fileVideo.src) {
      mediaIdleOverlay.style.setProperty('display', 'flex', 'important');
    } else if (currentMediaType === 'video') {
      mediaControls.style.display = 'flex';
    }
  }
}

tabWebcam.addEventListener('click', () => switchTab('webcam'));
tabMedia.addEventListener('click', () => switchTab('media'));

// ── Control Listeners ────────────────────────────────────────────────────────
modelPreset.addEventListener('change', () => loadSelectedModel());
devicePreset.addEventListener('change', () => loadSelectedModel());

confSlider.addEventListener('input', (e) => {
  confVal.textContent = parseFloat(e.target.value).toFixed(2);
  if (currentMode === 'media' && currentMediaType === 'image' && fileImage.src) {
    detectAndRender(fileImage);
  }
});

personFilter.addEventListener('change', () => {
  if (currentMode === 'media' && currentMediaType === 'image' && fileImage.src) {
    detectAndRender(fileImage);
  }
});

// Initialize default model on start
loadSelectedModel();