"""Small-step workflow executor helpers for generation harness adoption."""

from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.services.harness.artifacts import AssetCompletionResult
from app.services.harness.plans import WorkflowPlan
from app.services.harness.trace import TraceRecorder
from app.services.harness.types import ArtifactKind, WorkflowStep

if TYPE_CHECKING:
    from app.db.models import GenerationJob

AssetTask = Callable[[], Awaitable[AssetCompletionResult]]


@dataclass(frozen=True)
class AssetPlanRunResult:
    """Result of executing asset-producing tasks from one workflow plan."""

    task_results: tuple[AssetCompletionResult, ...]
    executed_task_keys: tuple[str, ...]
    ignored_task_keys: tuple[str, ...]

    @property
    def result_assets(self) -> tuple[str, ...]:
        """Assets returned by executed task handlers."""

        return tuple(result.asset for result in self.task_results)

    def to_metadata(self, plan: WorkflowPlan) -> dict[str, Any]:
        """Return internal executor coverage metadata for admin-only analytics."""

        return {
            "plan_mode": plan.mode.value,
            "planned_task_count": len(plan.tasks),
            "executed_task_count": len(self.executed_task_keys),
            "ignored_task_count": len(self.ignored_task_keys),
            "result_count": len(self.task_results),
            "executed_task_keys": list(self.executed_task_keys),
            "ignored_task_keys": list(self.ignored_task_keys),
            "result_assets": list(self.result_assets),
        }


async def record_workflow_plan(
    db: AsyncSession,
    *,
    job: "GenerationJob | None",
    plan: WorkflowPlan,
) -> None:
    """Persist a workflow plan snapshot for a tracked job."""

    await TraceRecorder(db).record_step(
        job=job,
        event_type="workflow_planned",
        status="succeeded",
        message="Workflow plan selected for this generation request.",
        metadata={"plan": plan.to_snapshot()},
        step=WorkflowStep.REQUEST_ACCEPTANCE,
        artifact=ArtifactKind.NONE,
        blocks_main_result=True,
    )


async def record_evaluation_result(
    db: AsyncSession,
    *,
    job: "GenerationJob | None",
    story_id: int | None = None,
    metadata: dict[str, Any],
    status: str,
    artifact: ArtifactKind | str = ArtifactKind.STORY_TEXT,
) -> None:
    """Persist a deterministic evaluation result for a tracked job."""

    await TraceRecorder(db).record_step(
        job=job,
        story_id=story_id,
        event_type="evaluation_completed",
        status=status,
        message="Generated content evaluation completed.",
        metadata=metadata,
        step=WorkflowStep.EVALUATION,
        artifact=artifact,
        blocks_main_result=status != "succeeded",
    )


async def record_executor_result(
    db: AsyncSession,
    *,
    job: "GenerationJob | None",
    plan: WorkflowPlan,
    result: AssetPlanRunResult,
) -> None:
    """Persist internal executor coverage metadata for a tracked job."""

    await TraceRecorder(db).record_step(
        job=job,
        event_type="executor_completed",
        status="succeeded",
        message="Workflow executor completed planned asset tasks.",
        metadata=result.to_metadata(plan),
        step=WorkflowStep.UNKNOWN,
        artifact=ArtifactKind.NONE,
        blocks_main_result=False,
    )


async def run_asset_plan(
    plan: WorkflowPlan,
    *,
    image_task: AssetTask | None = None,
    audio_task: AssetTask | None = None,
) -> AssetPlanRunResult:
    """Execute asset-producing tasks in the order declared by a workflow plan."""

    if plan.mode.value not in {"asset_generation", "asset_retry"}:
        raise ValueError("run_asset_plan only supports asset workflow plans")

    task_results: list[AssetCompletionResult] = []
    executed_task_keys: list[str] = []
    ignored_task_keys: list[str] = []

    for task in plan.tasks:
        if task.key == "complete_image_asset":
            if image_task is None:
                raise ValueError("Asset workflow plan requires an image task handler")
            task_results.append(await image_task())
            executed_task_keys.append(task.key)
            continue

        if task.key == "complete_audio_asset":
            if audio_task is None:
                raise ValueError("Asset workflow plan requires an audio task handler")
            task_results.append(await audio_task())
            executed_task_keys.append(task.key)
            continue

        ignored_task_keys.append(task.key)

    return AssetPlanRunResult(
        task_results=tuple(task_results),
        executed_task_keys=tuple(executed_task_keys),
        ignored_task_keys=tuple(ignored_task_keys),
    )