"""Trace recording helpers for generation harness workflows."""

from typing import TYPE_CHECKING, Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.services.generation_jobs import record_generation_event
from app.services.harness.types import (
    ArtifactKind,
    FailureCategory,
    WorkflowStep,
    normalize_trace_metadata,
)

if TYPE_CHECKING:
    from app.db.models import GenerationJob


class TraceRecorder:
    """Append workflow events with standard harness trace metadata."""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def record_step(
        self,
        *,
        job: "GenerationJob | None",
        event_type: str,
        status: str,
        story_id: int | None = None,
        message: str | None = None,
        metadata: dict[str, Any] | None = None,
        step: WorkflowStep | str | None = None,
        artifact: ArtifactKind | str | None = None,
        failure_category: FailureCategory | str | None = None,
        retryable: bool | None = None,
        blocks_main_result: bool | None = None,
        commit: bool = True,
    ):
        """Append a workflow event when the caller is running under a tracked job."""

        if job is None:
            return None

        return await record_generation_event(
            self.db,
            job=job,
            story_id=story_id,
            event_type=event_type,
            status=status,
            message=message,
            metadata=normalize_trace_metadata(
                event_type,
                metadata,
                step=step,
                artifact=artifact,
                failure_category=failure_category,
                retryable=retryable,
                blocks_main_result=blocks_main_result,
            ),
            commit=commit,
        )