auto-virtual-tryon/app/workers/activities/tryon_activities.py

"""Prepare-model and try-on mock activities plus shared helpers."""

from __future__ import annotations

from dataclasses import asdict, is_dataclass
from datetime import datetime, timezone
from enum import Enum
from typing import Any
from uuid import uuid4

from temporalio import activity

from app.domain.enums import AssetType, OrderStatus, StepStatus
from app.infra.db.models.asset import AssetORM
from app.infra.db.models.order import OrderORM
from app.infra.db.models.workflow_run import WorkflowRunORM
from app.infra.db.models.workflow_step import WorkflowStepORM
from app.infra.db.session import get_session_factory
from app.workers.workflows.types import MockActivityResult, StepActivityInput


def utc_now() -> datetime:
    """Return the current UTC timestamp."""

    return datetime.now(timezone.utc)


def jsonable(value: Any) -> Any:
    """Convert enums, dataclasses, and nested values to JSON-safe structures."""

    if value is None:
        return None
    if isinstance(value, Enum):
        return value.value
    if isinstance(value, datetime):
        return value.isoformat()
    if is_dataclass(value):
        return jsonable(asdict(value))
    if isinstance(value, dict):
        return {key: jsonable(item) for key, item in value.items() if item is not None}
    if isinstance(value, (list, tuple, set)):
        return [jsonable(item) for item in value]
    return value


def mock_uri(order_id: int, step_name: str, filename: str = "result.png") -> str:
    """Build a deterministic-looking mock URI for an order step."""

    return f"mock://orders/{order_id}/{step_name}/{uuid4().hex[:8]}-{filename}"


async def load_order_and_run(session, order_id: int, workflow_run_id: int) -> tuple[OrderORM, WorkflowRunORM]:
    """Load the order and workflow run required by an activity."""

    order = await session.get(OrderORM, order_id)
    workflow_run = await session.get(WorkflowRunORM, workflow_run_id)
    if order is None or workflow_run is None:
        raise ValueError("Order or workflow run not found for activity execution")
    return order, workflow_run


def create_step_record(payload: StepActivityInput) -> WorkflowStepORM:
    """Create a running workflow step row for an activity execution."""

    return WorkflowStepORM(
        workflow_run_id=payload.workflow_run_id,
        step_name=payload.step_name,
        step_status=StepStatus.RUNNING,
        input_json=jsonable(payload),
        started_at=utc_now(),
    )


async def execute_asset_step(
    payload: StepActivityInput,
    asset_type: AssetType,
    *,
    score: float = 0.95,
    filename: str = "result.png",
    message: str = "mock success",
    extra_metadata: dict[str, Any] | None = None,
    finalize: bool = False,
) -> MockActivityResult:
    """Persist a mock asset-producing step and return its result."""

    async with get_session_factory()() as session:
        order, workflow_run = await load_order_and_run(session, payload.order_id, payload.workflow_run_id)
        step = create_step_record(payload)
        session.add(step)

        order.status = OrderStatus.RUNNING
        workflow_run.status = OrderStatus.RUNNING
        workflow_run.current_step = payload.step_name
        await session.flush()

        try:
            metadata = {
                **payload.metadata,
                "source_asset_id": payload.source_asset_id,
                "selected_asset_id": payload.selected_asset_id,
                **(extra_metadata or {}),
            }
            metadata = {key: value for key, value in metadata.items() if value is not None}

            asset = AssetORM(
                order_id=payload.order_id,
                asset_type=asset_type,
                step_name=payload.step_name,
                uri=mock_uri(payload.order_id, payload.step_name.value, filename),
                metadata_json=jsonable(metadata),
            )
            session.add(asset)
            await session.flush()

            result = MockActivityResult(
                step_name=payload.step_name,
                success=True,
                asset_id=asset.id,
                uri=asset.uri,
                score=score,
                passed=True,
                message=message,
                metadata=jsonable(metadata) or {},
            )

            if finalize:
                order.final_asset_id = asset.id
                order.status = OrderStatus.SUCCEEDED
                workflow_run.status = OrderStatus.SUCCEEDED

            step.step_status = StepStatus.SUCCEEDED
            step.output_json = jsonable(result)
            step.ended_at = utc_now()
            await session.commit()
            return result
        except Exception as exc:
            step.step_status = StepStatus.FAILED
            step.error_message = str(exc)
            step.ended_at = utc_now()
            order.status = OrderStatus.FAILED
            workflow_run.status = OrderStatus.FAILED
            await session.commit()
            raise


@activity.defn
async def prepare_model_activity(payload: StepActivityInput) -> MockActivityResult:
    """Mock model preparation for the pipeline."""

    return await execute_asset_step(
        payload,
        AssetType.PREPARED_MODEL,
        extra_metadata={
            "model_id": payload.model_id,
            "pose_id": payload.pose_id,
            "garment_asset_id": payload.garment_asset_id,
            "scene_ref_asset_id": payload.scene_ref_asset_id,
        },
    )


@activity.defn
async def run_tryon_activity(payload: StepActivityInput) -> MockActivityResult:
    """Mock try-on rendering."""

    return await execute_asset_step(
        payload,
        AssetType.TRYON,
        extra_metadata={"prepared_asset_id": payload.source_asset_id},
    )