import {OBJECT_DETECTION_MODEL_ID} from '../../shared/clusterModels';
import type {ObjectDetector} from '../detection/ObjectDetector';
import type {DetectionResult} from '../detection/workerMessages';
import {AbstractModelTaskHandler} from './AbstractModelTaskHandler';
import type {ClusterTaskHostContext, HostTaskMessage} from './types';

export class DetectTaskHandler extends AbstractModelTaskHandler {
  readonly kind = 'detect' as const;

  constructor(
    host: ClusterTaskHostContext,
    private readonly getDetector: () => ObjectDetector | null,
  ) {
    super(host);
  }

  async process(task: HostTaskMessage): Promise<void> {
    const detector = this.getDetector();
    if (this.host.getSelectedModelId() !== OBJECT_DETECTION_MODEL_ID || !detector) {
      await this.completeTask(
        task.id,
        this.taskError(
          'This host is not running the object detection model. Re-register with "Object detection" selected.',
        ),
      );
      return;
    }

    await this.withVideoFrame(task, async (frame) => {
      const results = await detector.detect(frame, {threshold: task.threshold ?? 0.5});
      return {
        status: 'done',
        kind: 'detect',
        result: DetectTaskHandler.serializeDetectionResults(results),
      };
    });
  }

  private static 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,
      },
    }));
  }
}