src/pages/detectionNodeHost.ts

import {
  CLUSTER_MODELS,
  DEFAULT_CLUSTER_MODEL_ID,
  getClusterModel,
  isImplementedClusterModel,
} from '../../shared/clusterModels';
import type {ObjectDetector} from '../detection/ObjectDetector';
import type {DetectionResult} from '../detection/workerMessages';
import type {VideoDescriber} from '../videodescription/VideoDescriber';

const API_BASE = import.meta.env.VITE_API_BASE ?? '';

const OBJECT_DETECTION_MODEL_ID = 'rfdetr-medium';
const VIDEO_DESCRIPTION_MODEL_ID = 'smolvlm-500m';

type ClusterTaskKind = 'detect' | 'describe';

interface HostTaskMessage {
  id: string;
  kind: ClusterTaskKind;
  image_base64: string;
  mime_type: string;
  threshold?: number;
  instruction?: string;
  max_new_tokens?: number;
}

const statusEl = document.getElementById('status')!;
const hostIdInput = document.getElementById('host-id') as HTMLInputElement;
const modelSelect = document.getElementById('model-select') as HTMLSelectElement;
const registerBtn = document.getElementById('register-btn') as HTMLButtonElement;
const curlExample = document.getElementById('curl-example')!;

let detector: ObjectDetector | null = null;
let describer: VideoDescriber | null = null;
let hostId = '';
let selectedModelId = DEFAULT_CLUSTER_MODEL_ID;
let processing = false;
let eventSource: EventSource | null = null;
const taskQueue: HostTaskMessage[] = [];

function setStatus(message: string): void {
  statusEl.textContent = message;
  console.log('[detection-host]', message);
}

function populateModelSelect(): void {
  modelSelect.replaceChildren(
    ...CLUSTER_MODELS.map((model) => {
      const option = document.createElement('option');
      option.value = model.id;
      option.textContent = model.implemented
        ? model.label
        : `${model.label} (coming soon)`;
      option.disabled = !model.implemented;
      option.title = model.description;
      return option;
    }),
  );
  modelSelect.value = DEFAULT_CLUSTER_MODEL_ID;
}

function updateCurlExample(): void {
  const origin = window.location.origin;

  if (selectedModelId === OBJECT_DETECTION_MODEL_ID) {
    curlExample.textContent = `curl -X POST '${origin}/v1/detect' \\
  -H 'Content-Type: application/json' \\
  -d '{
    "host": "${hostId}",
    "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png",
    "threshold": 0.5
  }'`;
    return;
  }

  if (selectedModelId === VIDEO_DESCRIPTION_MODEL_ID) {
    curlExample.textContent = `curl -X POST '${origin}/v1/describe' \\
  -H 'Content-Type: application/json' \\
  -d '{
    "host": "${hostId}",
    "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png",
    "instruction": "What do you see?",
    "max_new_tokens": 100
  }'`;
    return;
  }

  const model = getClusterModel(selectedModelId);
  curlExample.textContent = model
    ? `# ${model.label}\n# API endpoint for this model is not wired yet.`
    : '';
}

async function ensureModelLoaded(modelId: string): Promise<void> {
  if (modelId === OBJECT_DETECTION_MODEL_ID) {
    if (!detector) {
      setStatus('Loading RF-DETR (WebGPU)…');
      const {ObjectDetector: Detector} = await import('../detection/ObjectDetector');
      detector = await Detector.create(setStatus);
    }
    return;
  }

  if (modelId === VIDEO_DESCRIPTION_MODEL_ID) {
    if (!describer) {
      setStatus('Loading SmolVLM-500M (WebGPU)…');
      const {VideoDescriber: Describer} = await import('../videodescription/VideoDescriber');
      describer = await Describer.create(setStatus);
    }
    return;
  }

  const model = getClusterModel(modelId);
  throw new Error(
    model
      ? `${model.label} is not available on this host yet.`
      : `Unknown model "${modelId}".`,
  );
}

async function base64ToVideoFrame(base64: string, mimeType: string): Promise<VideoFrame> {
  const binary = atob(base64);
  const bytes = new Uint8Array(binary.length);
  for (let i = 0; i < binary.length; i++) {
    bytes[i] = binary.charCodeAt(i);
  }

  const blob = new Blob([bytes], {type: mimeType});
  const bitmap = await createImageBitmap(blob);
  const frame = new VideoFrame(bitmap, {timestamp: 0});
  bitmap.close();
  return frame;
}

async function processDetectTask(task: HostTaskMessage): Promise<void> {
  if (!detector) {
    return;
  }

  let frame: VideoFrame | null = null;

  try {
    frame = await base64ToVideoFrame(task.image_base64, task.mime_type);
    const results = await detector.detect(frame, {threshold: task.threshold ?? 0.5});

    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({
        status: 'done',
        results: serializeDetectionResults(results),
      }),
    });
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({status: 'error', error: message}),
    });
    throw error;
  } finally {
    frame?.close();
  }
}

async function processDescribeTask(task: HostTaskMessage): Promise<void> {
  if (!describer) {
    return;
  }

  let frame: VideoFrame | null = null;

  try {
    frame = await base64ToVideoFrame(task.image_base64, task.mime_type);
    const description = await describer.describe(frame, {
      instruction: task.instruction ?? 'What do you see?',
      maxNewTokens: task.max_new_tokens ?? 100,
    });

    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({
        status: 'done',
        description,
      }),
    });
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    await fetch(`${API_BASE}/api/tasks/${task.id}/complete`, {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({status: 'error', error: message}),
    });
    throw error;
  } finally {
    frame?.close();
  }
}

async function processTask(task: HostTaskMessage): Promise<void> {
  processing = true;
  setStatus(`Processing task ${task.id.slice(0, 8)}…`);

  try {
    if (task.kind === 'describe') {
      await processDescribeTask(task);
    } else {
      await processDetectTask(task);
    }
    setStatus(`Task ${task.id.slice(0, 8)} complete. Waiting for requests…`);
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    setStatus(`Task failed: ${message}`);
  } finally {
    processing = false;
    void drainTaskQueue();
  }
}

async function drainTaskQueue(): Promise<void> {
  if (processing || taskQueue.length === 0) {
    return;
  }

  const task = taskQueue.shift();
  if (task) {
    await processTask(task);
  }
}

function enqueueTask(task: HostTaskMessage): void {
  taskQueue.push(task);
  void drainTaskQueue();
}

function serializeDetectionResults(results: DetectionResult[]): unknown[] {
  return results.map((result) => ({
    label: result.label,
    score: result.score,
    box: {
      xmin: result.box.xmin,
      ymin: result.box.ymin,
      xmax: result.box.xmax,
      ymax: result.box.ymax,
    },
  }));
}

function connectStream(): void {
  eventSource?.close();

  const url = `${API_BASE}/api/hosts/stream?host_id=${encodeURIComponent(hostId)}`;
  eventSource = new EventSource(url);

  eventSource.addEventListener('ready', () => {
    const model = getClusterModel(selectedModelId);
    const modelLabel = model?.label ?? selectedModelId;
    setStatus(`Hosting "${hostId}" (${modelLabel}). Waiting for requests…`);
  });

  eventSource.addEventListener('task', (event) => {
    const raw = JSON.parse(event.data) as HostTaskMessage & {kind?: ClusterTaskKind};
    const task: HostTaskMessage = {
      ...raw,
      kind: raw.kind ?? 'detect',
    };
    enqueueTask(task);
  });

  eventSource.onerror = () => {
    if (eventSource?.readyState === EventSource.CLOSED) {
      setStatus('Broker connection lost. Reconnecting in 2s…');
      setTimeout(connectStream, 2000);
    }
  };
}

async function registerHost(): Promise<void> {
  hostId = hostIdInput.value.trim();
  selectedModelId = modelSelect.value;

  if (!hostId) {
    setStatus('Enter a host id (e.g. my-gpu-node).');
    return;
  }

  if (!isImplementedClusterModel(selectedModelId)) {
    setStatus('Choose an available model from the list.');
    return;
  }

  registerBtn.disabled = true;
  modelSelect.disabled = true;
  hostIdInput.disabled = true;

  try {
    await ensureModelLoaded(selectedModelId);

    const res = await fetch(`${API_BASE}/api/hosts/register`, {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({id: hostId, model: selectedModelId}),
    });

    if (!res.ok) {
      const json = await res.json().catch(() => ({}));
      throw new Error(json.error ?? `Register failed (${res.status})`);
    }

    updateCurlExample();
    connectStream();
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    setStatus(`Failed to start host:\n${message}`);
    registerBtn.disabled = false;
    modelSelect.disabled = false;
    hostIdInput.disabled = false;
  }
}

registerBtn.addEventListener('click', () => {
  void registerHost();
});

modelSelect.addEventListener('change', () => {
  selectedModelId = modelSelect.value;
  updateCurlExample();
});

populateModelSelect();
hostIdInput.value = `node-${Math.random().toString(36).slice(2, 8)}`;
selectedModelId = modelSelect.value;
updateCurlExample();

setStatus('Pick a model, enter a host id, and click “Start hosting”.');