from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from cloud.internal_api.models import AssignmentModel
from core.models import Task
from host_agent.execution import ExecutionFactories


@dataclass(frozen=True)
class AssignmentExecutionResult:
    status: str
    failure_reason: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


class AssignmentExecutor:
    def __init__(self, factories: ExecutionFactories) -> None:
        self.factories = factories

    def execute(self, assignment: AssignmentModel) -> AssignmentExecutionResult:
        if assignment.workflow_definition_id is not None:
            return self._execute_workflow(assignment)
        if assignment.goal is not None:
            return self._execute_goal(assignment)
        return AssignmentExecutionResult(
            status="failed",
            failure_reason="assignment has neither goal nor workflow definition",
        )

    def _execute_goal(
        self,
        assignment: AssignmentModel,
    ) -> AssignmentExecutionResult:
        task = Task(goal=assignment.goal or "", device_id=assignment.device_id)
        completed = self.factories.task_runner_factory().run(task)
        return AssignmentExecutionResult(
            status="done" if completed.status == "completed" else "failed",
            failure_reason=completed.failure_reason,
            metadata={
                "runtime_task_id": completed.id,
                "runtime_status": completed.status,
            },
        )

    def _execute_workflow(
        self,
        assignment: AssignmentModel,
    ) -> AssignmentExecutionResult:
        definition_id = assignment.workflow_definition_id or ""
        definition = self.factories.workflow_store.get_definition(definition_id)
        if definition is None:
            return AssignmentExecutionResult(
                status="failed",
                failure_reason=f"unknown workflow definition {definition_id!r}",
            )
        run = self.factories.workflow_runner_factory().run(
            definition,
            device_id=assignment.device_id,
        )
        return AssignmentExecutionResult(
            status="done" if run.status == "completed" else "failed",
            failure_reason=(
                None if run.status == "completed" else f"workflow ended as {run.status}"
            ),
            metadata={
                "workflow_run_id": run.id,
                "workflow_status": run.status,
            },
        )