test-assignment-datawrapper/src/screenshots/screenshots.controller.ts

import { InjectQueue } from '@nestjs/bull';
import { Body, Controller, Get, NotFoundException, Param, Post, Res, StreamableFile } from '@nestjs/common';
import { ApiResponse } from '@nestjs/swagger';
import { Response } from 'express'
import { Readable } from 'stream';
import { CreateJobRequestDto, GetJobResponseDto, JobStatusDto } from './screenshots.dto'
import { QUEUE_NAME, ScreenshotsQueue } from './shared';

@Controller('screenshots')
export class ScreenshotsController {
  constructor(@InjectQueue(QUEUE_NAME) private readonly screenshotsQueue: ScreenshotsQueue) {}

  @Post()
  async createScreenshotJob(@Body() createScreenshotJobDto: CreateJobRequestDto): Promise<{ jobId: string }> {
    const result = await this.screenshotsQueue.add({
      pageUrl: new URL(createScreenshotJobDto.pageUrl),
      imageType: createScreenshotJobDto.imageType,
    })
    return {
      jobId: result.id.toString()
    }
  }

  @Get(':id')
  @ApiResponse({ status: 404 })
  async getJob(@Param('id') id: string): Promise<GetJobResponseDto> {
    const jobInfo = await this.screenshotsQueue.getJob(id)

    if (!jobInfo) {
      throw new NotFoundException()
    }

    switch (await jobInfo.getState()) {
      case 'completed':
        return { status: JobStatusDto.Completed }
      case 'failed':
        return { status: JobStatusDto.Failed }
      default:
        return { status: JobStatusDto.Queued }
    }
  }

  @Get(':id/result')
  @ApiResponse({ status: 404 })
  async getScreenshot(@Param('id') id: string, @Res({ passthrough: true }) res: Response): Promise<StreamableFile> {
    const jobInfo = await this.screenshotsQueue.getJob(id)

    if (!jobInfo) {
      throw new NotFoundException()
    }

    if (!await jobInfo.isCompleted()) {
      throw new NotFoundException()
    }

    switch (jobInfo.data.imageType) {
      case 'jpeg':
        res.setHeader('Content-Type', 'image/jpeg')
        break
      case 'png':
        res.setHeader('Content-Type', 'image/png')
        break
    }

    const buffer = Buffer.from(await jobInfo.returnvalue as string, 'base64')
    return new StreamableFile(buffer)
  }
}