app/routers/screenshot/screenshot.py

from __future__ import annotations

import logging

from core.service.playwright import PlaywrightInstance, TimeoutError
from core.service.playwright.models import PageModel, ScreenshotModel
from fastapi import APIRouter, HTTPException
from fastapi.responses import Response

router = APIRouter()

playwright_instance = PlaywrightInstance()


@router.post("/screenshot", tags=["V1"])
async def screenshot(
    screenshot_model: ScreenshotModel,
    page_model: PageModel | None = None,
) -> Response:
    page_model = page_model if page_model else PageModel()

    await playwright_instance

    try:
        response = await playwright_instance.screenshot(
            screenshot_model=screenshot_model,
            page_model=page_model,
        )
        return Response(content=response, media_type=f"image/{screenshot_model.image_type}")
    except TimeoutError as e:
        status_code = 504 if isinstance(e, TimeoutError) else 400
        detail = f"An error occurred while generating the screenshot: {e}"
        raise HTTPException(status_code=status_code, detail=detail) from e
    except Exception as e:
        logging.exception("Unhandled error in screenshot generation")
        raise HTTPException(status_code=500, detail="An unexpected error occurred") from e