src/pages/addNode.ts

import {
  CLUSTER_MODELS,
  DEFAULT_CLUSTER_MODEL_ID,
  getClusterModel,
  isImplementedClusterModel,
  OBJECT_DETECTION_MODEL_ID,
  VIDEO_DESCRIPTION_MODEL_ID,
} from '../../shared/clusterModels';
import {
  ClusterTaskHandler,
  type ClusterTaskKind,
  type HostTaskMessage,
} from '../clusterTaskHandler';
import type {ObjectDetector} from '../detection/ObjectDetector';
import type {VideoFrameDescriber} from '../videodescription/VideoFrameDescriber';

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

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: VideoFrameDescriber | null = null;
let hostId = '';
let selectedModelId = DEFAULT_CLUSTER_MODEL_ID;
let eventSource: EventSource | null = null;

const taskHandler = new ClusterTaskHandler({
  apiBase: API_BASE,
  getSelectedModelId: () => selectedModelId,
  getDetector: () => detector,
  getDescriber: () => describer,
  setStatus,
});

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 {VideoFrameDescriber: Describer} = await import('../videodescription/VideoFrameDescriber');
      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}".`,
  );
}

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 defaultKind: ClusterTaskKind =
      selectedModelId === VIDEO_DESCRIPTION_MODEL_ID ? 'describe' : 'detect';
    taskHandler.enqueueTask(taskHandler.parseIncomingTask(raw, defaultKind));
  });

  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”.');