app/routers/screenshot/screenshot.py

from __future__ import annotations

import logging

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

router = APIRouter()
playwright_context = AsyncPlaywrightContext()


@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()
    async with playwright_context as sc:
        try:
            response = await sc.screenshot(screenshot_model=screenshot_model, page_model=page_model)
            return Response(content=response, media_type="image/png")
        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