from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal, Protocol

from cloud.internal_api.models import AssignmentModel
from host_agent.assignment import AssignmentExecutionResult
from host_agent.client import HostAgentClient
from host_agent.status import AgentStatusTracker

if TYPE_CHECKING:
    from collections.abc import Callable


class ActiveAssignmentExecutor(Protocol):
    async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult: ...

    def request_stop(self) -> None: ...


@dataclass(frozen=True)
class AssignmentProcessingResult:
    execution: AssignmentExecutionResult
    report_status: Literal["recorded", "already_recorded"]


class AssignmentProcessor:
    def __init__(
        self,
        client: HostAgentClient,
        active_executor: ActiveAssignmentExecutor,
        *,
        status_tracker: AgentStatusTracker | None = None,
        on_result: Callable[[AssignmentModel, AssignmentProcessingResult], None]
        | None = None,
    ) -> None:
        self.client = client
        self.active_executor = active_executor
        self.status_tracker = status_tracker
        self.on_result = on_result

    def request_stop(self) -> None:
        self.active_executor.request_stop()

    async def process(self, assignment: AssignmentModel) -> AssignmentProcessingResult:
        if self.status_tracker is not None:
            self.status_tracker.mark_assignment_started(assignment)
        try:
            execution = await self.active_executor.run(assignment)
            if execution.status in {"done", "cancelled"}:
                status = execution.status
            else:
                status = "failed"
            failure_reason = execution.failure_reason if status != "done" else None
            response = await self.client.report_result(
                assignment,
                status=status,
                failure_reason=failure_reason,
                result=dict(execution.metadata),
            )
            result = AssignmentProcessingResult(
                execution=execution,
                report_status=response.status,
            )
        finally:
            if self.status_tracker is not None:
                self.status_tracker.mark_assignment_finished()
        if self.on_result is not None:
            try:
                self.on_result(assignment, result)
            except Exception:
                pass
        return result