freepik-api/app/routers/image_generation.py

import asyncio
import logging
from datetime import datetime, timezone

from fastapi import APIRouter, HTTPException, Query

from app.schemas.common import TaskDetail, TaskResponse, TaskStatus
from app.schemas.image_generation import (
    FluxDevRequest,
    FluxProRequest,
    MysticRequest,
    SeedreamRequest,
)
from app.services import task_tracker
from app.services.freepik_client import TaskResult, _extract_task_id

logger = logging.getLogger(__name__)

router = APIRouter(prefix='/api/v1/generate/image', tags=['image-generation'])


async def _submit_and_respond(
    result: TaskResult,
    sync: bool,
    metadata: dict | None = None,
) -> TaskResponse | TaskDetail:
    """Extract task_id from Freepik response, track it, optionally wait."""
    freepik_task_id = _extract_task_id(result.data)
    if not freepik_task_id:
        raise HTTPException(status_code=502, detail='No task_id in Freepik response')

    internal_id = task_tracker.submit(freepik_task_id, result.status_path, metadata)

    if not sync:
        return TaskResponse(
            task_id=internal_id,
            status=TaskStatus.pending,
            created_at=datetime.now(timezone.utc),
        )

    # Sync mode: wait for completion
    from app.config import settings
    elapsed = 0
    while elapsed < settings.task_poll_timeout_seconds:
        await asyncio.sleep(2)
        elapsed += 2
        task = task_tracker.get_task(internal_id)
        if not task:
            raise HTTPException(status_code=500, detail='Task disappeared')
        if task['status'] == TaskStatus.completed:
            return TaskDetail(
                task_id=internal_id,
                status=TaskStatus.completed,
                created_at=task['created_at'],
                updated_at=task['updated_at'],
                progress=1.0,
                result_url=f'/api/v1/tasks/{internal_id}/result',
            )
        if task['status'] == TaskStatus.failed:
            raise HTTPException(status_code=502, detail=task.get('error', 'Task failed'))

    raise HTTPException(status_code=504, detail='Task did not complete in time')


@router.post('/mystic', response_model=TaskResponse)
async def generate_mystic(request: MysticRequest, sync: bool = Query(False)):
    from app.services import freepik_client
    result = await freepik_client.generate_mystic(
        prompt=request.prompt,
        resolution=request.resolution,
        aspect_ratio=request.aspect_ratio,
        model=request.model,
        seed=request.seed,
        styling=request.styling,
        structure_reference=request.structure_reference,
        style_reference=request.style_reference,
    )
    return await _submit_and_respond(result, sync, {'model': 'mystic'})


@router.post('/flux-dev', response_model=TaskResponse)
async def generate_flux_dev(request: FluxDevRequest, sync: bool = Query(False)):
    from app.services import freepik_client
    result = await freepik_client.generate_flux_dev(
        prompt=request.prompt,
        aspect_ratio=request.aspect_ratio,
        styling=request.styling,
        seed=request.seed,
    )
    return await _submit_and_respond(result, sync, {'model': 'flux-dev'})


@router.post('/flux-pro', response_model=TaskResponse)
async def generate_flux_pro(request: FluxProRequest, sync: bool = Query(False)):
    from app.services import freepik_client
    result = await freepik_client.generate_flux_pro(
        prompt=request.prompt,
        aspect_ratio=request.aspect_ratio,
        styling=request.styling,
        seed=request.seed,
    )
    return await _submit_and_respond(result, sync, {'model': 'flux-pro'})


@router.post('/seedream', response_model=TaskResponse)
async def generate_seedream(request: SeedreamRequest, sync: bool = Query(False)):
    from app.services import freepik_client
    result = await freepik_client.generate_seedream(
        prompt=request.prompt,
        aspect_ratio=request.aspect_ratio,
        seed=request.seed,
    )
    return await _submit_and_respond(result, sync, {'model': 'seedream'})