server/store.ts

import {randomUUID} from 'node:crypto';
import type {
  DetectionResultJson,
  DetectionTask,
  HostRecord,
  HostStream,
} from './types.js';

type TaskWaiter = {
  resolve: (task: DetectionTask) => void;
  reject: (error: Error) => void;
  timeout: ReturnType<typeof setTimeout>;
};

export class BrokerStore {
  private readonly hosts = new Map<string, HostRecord>();
  private readonly tasks = new Map<string, DetectionTask>();
  private readonly waiters = new Map<string, TaskWaiter>();

  registerHost(id: string, model: string): HostRecord {
    const existing = this.hosts.get(id);
    const host: HostRecord = {
      id,
      model,
      registeredAt: existing?.registeredAt ?? Date.now(),
      lastHeartbeat: Date.now(),
      stream: existing?.stream,
    };
    this.hosts.set(id, host);
    return host;
  }

  attachHostStream(id: string, stream: HostStream): void {
    const host = this.hosts.get(id);
    if (!host) {
      throw new Error(`Host '${id}' is not registered`);
    }

    host.stream?.close();
    host.stream = stream;
    host.lastHeartbeat = Date.now();
    this.flushPendingTasks(id);
  }

  detachHostStream(id: string): void {
    const host = this.hosts.get(id);
    if (host) {
      host.stream = undefined;
    }
  }

  heartbeat(id: string): void {
    const host = this.hosts.get(id);
    if (host) {
      host.lastHeartbeat = Date.now();
    }
  }

  listHosts(): Array<{
    id: string;
    model: string;
    online: boolean;
    registeredAt: number;
    lastHeartbeat: number;
  }> {
    return [...this.hosts.values()].map((host) => ({
      id: host.id,
      model: host.model,
      online: Boolean(host.stream),
      registeredAt: host.registeredAt,
      lastHeartbeat: host.lastHeartbeat,
    }));
  }

  getHost(id: string): HostRecord | undefined {
    return this.hosts.get(id);
  }

  isHostOnline(id: string): boolean {
    const host = this.hosts.get(id);
    return Boolean(host?.stream);
  }

  createDetectTask(
    hostId: string,
    payload: {imageBase64: string; mimeType: string; threshold: number},
  ): DetectionTask {
    return this.createTask(hostId, {
      kind: 'detect',
      imageBase64: payload.imageBase64,
      mimeType: payload.mimeType,
      threshold: payload.threshold,
    });
  }

  createDescribeTask(
    hostId: string,
    payload: {
      imageBase64: string;
      mimeType: string;
      instruction: string;
      maxNewTokens: number;
    },
  ): DetectionTask {
    return this.createTask(hostId, {
      kind: 'describe',
      imageBase64: payload.imageBase64,
      mimeType: payload.mimeType,
      instruction: payload.instruction,
      maxNewTokens: payload.maxNewTokens,
    });
  }

  private createTask(
    hostId: string,
    payload: {
      kind: DetectionTask['kind'];
      imageBase64: string;
      mimeType: string;
      threshold?: number;
      instruction?: string;
      maxNewTokens?: number;
    },
  ): DetectionTask {
    const host = this.hosts.get(hostId);
    if (!host) {
      throw new Error(`Host '${hostId}' not found. Register the browser host first.`);
    }

    const task: DetectionTask = {
      id: randomUUID(),
      hostId,
      kind: payload.kind,
      status: 'pending',
      imageBase64: payload.imageBase64,
      mimeType: payload.mimeType,
      threshold: payload.threshold,
      instruction: payload.instruction,
      maxNewTokens: payload.maxNewTokens,
      createdAt: Date.now(),
    };

    this.tasks.set(task.id, task);
    this.dispatchTask(task);
    return task;
  }

  waitForTask(taskId: string, timeoutMs: number): Promise<DetectionTask> {
    const task = this.tasks.get(taskId);
    if (!task) {
      return Promise.reject(new Error(`Task '${taskId}' not found`));
    }

    if (task.status === 'done' || task.status === 'error') {
      return Promise.resolve(task);
    }

    return new Promise<DetectionTask>((resolve, reject) => {
      const timeout = setTimeout(() => {
        this.waiters.delete(taskId);
        reject(new Error(`Task '${taskId}' timed out after ${timeoutMs}ms`));
      }, timeoutMs);

      this.waiters.set(taskId, {
        resolve,
        reject,
        timeout,
      });
    });
  }

  markProcessing(taskId: string): void {
    const task = this.tasks.get(taskId);
    if (!task) {
      return;
    }

    task.status = 'processing';
    task.startedAt = Date.now();
  }

  completeTask(
    taskId: string,
    outcome:
      | {status: 'done'; kind: 'detect'; results: DetectionResultJson[]}
      | {status: 'done'; kind: 'describe'; description: string}
      | {status: 'error'; error: string},
  ): DetectionTask {
    const task = this.tasks.get(taskId);
    if (!task) {
      throw new Error(`Task '${taskId}' not found`);
    }

    task.completedAt = Date.now();

    if (outcome.status === 'done') {
      task.status = 'done';
      if (outcome.kind === 'detect') {
        task.results = outcome.results;
        delete task.description;
      } else {
        task.description = outcome.description;
        delete task.results;
      }
    } else {
      task.status = 'error';
      task.error = outcome.error;
    }

    const waiter = this.waiters.get(taskId);
    if (waiter) {
      clearTimeout(waiter.timeout);
      this.waiters.delete(taskId);
      if (task.status === 'error') {
        waiter.reject(new Error(task.error ?? 'Detection failed'));
      } else {
        waiter.resolve(task);
      }
    }

    return task;
  }

  getTask(taskId: string): DetectionTask | undefined {
    return this.tasks.get(taskId);
  }

  private dispatchTask(task: DetectionTask): void {
    const host = this.hosts.get(task.hostId);
    if (!host?.stream) {
      return;
    }

    host.stream.write('task', {
      id: task.id,
      kind: task.kind,
      threshold: task.threshold,
      instruction: task.instruction,
      max_new_tokens: task.maxNewTokens,
      image_base64: task.imageBase64,
      mime_type: task.mimeType,
    });
  }

  private flushPendingTasks(hostId: string): void {
    for (const task of this.tasks.values()) {
      if (task.hostId === hostId && task.status === 'pending') {
        this.dispatchTask(task);
      }
    }
  }
}

export const store = new BrokerStore();