"""Scene activities."""

from temporalio import activity

from app.domain.enums import AssetType, LibraryResourceType, OrderStatus, StepStatus
from app.infra.db.models.asset import AssetORM
from app.infra.db.session import get_session_factory
from app.workers.activities.tryon_activities import (
    create_step_record,
    execute_asset_step,
    get_image_generation_service,
    get_order_artifact_storage_service,
    jsonable,
    load_active_library_resource,
    load_order_and_run,
    utc_now,
)
from app.workers.workflows.types import MockActivityResult, StepActivityInput


@activity.defn
async def run_scene_activity(payload: StepActivityInput) -> MockActivityResult:
    """Generate a scene-composited asset, or fall back to mock mode when configured."""

    service = get_image_generation_service()
    if service.__class__.__name__ == "MockImageGenerationService":
        return await execute_asset_step(
            payload,
            AssetType.SCENE,
            extra_metadata={"scene_ref_asset_id": payload.scene_ref_asset_id},
        )

    if payload.source_asset_id is None:
        raise ValueError("run_scene_activity requires source_asset_id")
    if payload.scene_ref_asset_id is None:
        raise ValueError("run_scene_activity requires scene_ref_asset_id")

    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:
            source_asset = await session.get(AssetORM, payload.source_asset_id)
            if source_asset is None:
                raise ValueError(f"Source asset {payload.source_asset_id} not found")
            if source_asset.order_id != payload.order_id:
                raise ValueError(
                    f"Source asset {payload.source_asset_id} does not belong to order {payload.order_id}"
                )

            scene_resource, scene_original = await load_active_library_resource(
                session,
                payload.scene_ref_asset_id,
                resource_type=LibraryResourceType.SCENE,
            )

            generated = await service.generate_scene_image(
                source_image_url=source_asset.uri,
                scene_image_url=scene_original.public_url,
            )
            storage_key, public_url = await get_order_artifact_storage_service().upload_generated_image(
                order_id=payload.order_id,
                step_name=payload.step_name,
                image_bytes=generated.image_bytes,
                mime_type=generated.mime_type,
            )

            metadata = {
                **payload.metadata,
                "source_asset_id": source_asset.id,
                "scene_ref_asset_id": payload.scene_ref_asset_id,
                "scene_resource_id": scene_resource.id,
                "scene_original_file_id": scene_original.id,
                "scene_original_url": scene_original.public_url,
                "provider": generated.provider,
                "model": generated.model,
                "storage_key": storage_key,
                "mime_type": generated.mime_type,
                "prompt": generated.prompt,
            }
            metadata = {key: value for key, value in metadata.items() if value is not None}

            asset = AssetORM(
                order_id=payload.order_id,
                asset_type=AssetType.SCENE,
                step_name=payload.step_name,
                uri=public_url,
                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=1.0,
                passed=True,
                message="scene generated",
                metadata=jsonable(metadata) or {},
            )

            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