import {VIDEO_DESCRIPTION_MODEL_ID} from '../../shared/clusterModels';
import type {VideoFrameDescriber} from '../videodescription/VideoFrameDescriber';
import {AbstractModelTaskHandler} from './AbstractModelTaskHandler';
import type {ClusterTaskHostContext, HostTaskMessage} from './types';

export class DescribeTaskHandler extends AbstractModelTaskHandler {
  readonly kind = 'describe' as const;

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

  async process(task: HostTaskMessage): Promise<void> {
    const describer = this.getDescriber();
    if (this.host.getSelectedModelId() !== VIDEO_DESCRIPTION_MODEL_ID || !describer) {
      await this.completeTask(
        task.id,
        this.taskError(
          'This host is not running the video description model. Re-register with "Video description" selected.',
        ),
      );
      return;
    }

    await this.withVideoFrame(task, async (frame) => {
      const description = await describer.describe(frame, {
        instruction: task.instruction ?? 'What do you see?',
        maxNewTokens: task.max_new_tokens ?? 100,
      });

      if (!description.trim()) {
        throw new Error('SmolVLM returned an empty description');
      }

      return {status: 'done', kind: 'describe', result: {description}};
    });
  }
}