from __future__ import annotations

from dataclasses import dataclass
from typing import Literal, Protocol

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


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,
    ) -> None:
        self.client = client
        self.active_executor = active_executor

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

    async def process(self, assignment: AssignmentModel) -> AssignmentProcessingResult:
        execution = await self.active_executor.run(assignment)
        status = "done" if execution.status == "done" else "failed"
        failure_reason = execution.failure_reason if status == "failed" else None
        response = await self.client.report_result(
            assignment,
            status=status,
            failure_reason=failure_reason,
            result=dict(execution.metadata),
        )
        return AssignmentProcessingResult(
            execution=execution,
            report_status=response.status,
        )