from __future__ import annotations

import hashlib
from dataclasses import dataclass
from pathlib import Path, PurePosixPath
from typing import Protocol

from k1link.simulation.contracts import QualificationRun
from k1link.simulation.orchestrator import WorkerStartResult, WorkerStopResult
from k1link.simulation.process_supervisor import (
    PosixProcessSupervisor,
    ProcessSpec,
)
from k1link.simulation.s0 import S0Profile, load_s0_profile


class WorkerAdmissionError(RuntimeError):
    """The worker no longer matches the accepted S0 generation or D boundary."""


@dataclass(frozen=True, slots=True)
class WorkerAdmission:
    profile: S0Profile
    profile_sha256: str
    artifact_run_root: PurePosixPath
    process_run_root: PurePosixPath


class SimulationWorldControl(Protocol):
    def pause(self, run_id: str) -> None: ...

    def resume(self, run_id: str) -> None: ...

    def step(self, run_id: str, step_count: int) -> int: ...


class S0WorkerGuard:
    """Bind every S1 start to the exact accepted S0 profile and D artifact root."""

    def __init__(self, profile_path: Path) -> None:
        self.profile_path = profile_path.expanduser().resolve()
        self.profile = load_s0_profile(self.profile_path)
        self.profile_sha256 = hashlib.sha256(self.profile_path.read_bytes()).hexdigest()

    def admit(
        self,
        run: QualificationRun,
        artifact_run_root: PurePosixPath,
        process_run_root: PurePosixPath,
    ) -> WorkerAdmission:
        if run.host_profile_id != self.profile.profile_id:
            raise WorkerAdmissionError("run host profile id does not match accepted S0")
        if run.host_profile_sha256 != self.profile_sha256:
            raise WorkerAdmissionError("run host profile digest does not match accepted S0")
        artifact_roots = tuple(
            item.wsl_path
            for item in self.profile.storage.mutable_paths
            if item.identifier == "artifacts"
        )
        if len(artifact_roots) != 1:
            raise WorkerAdmissionError("accepted S0 profile has no unique artifact root")
        runtime_roots = tuple(
            item.wsl_path
            for item in self.profile.storage.mutable_paths
            if item.identifier == "runtime"
        )
        if len(runtime_roots) != 1:
            raise WorkerAdmissionError("accepted S0 profile has no unique runtime root")
        expected_artifact_parent = artifact_roots[0] / "s1" / "runs"
        expected_runtime_parent = runtime_roots[0] / "s1" / "runs"
        if artifact_run_root != expected_artifact_parent / run.run_id:
            raise WorkerAdmissionError("S1 run root escapes the accepted D-only layout")
        if process_run_root != expected_runtime_parent / run.run_id:
            raise WorkerAdmissionError("S1 process root escapes the accepted D-only layout")
        if not all(
            component.qualification_state == "accepted" for component in self.profile.components
        ):
            raise WorkerAdmissionError("S0 component generation is no longer accepted")
        return WorkerAdmission(
            profile=self.profile,
            profile_sha256=self.profile_sha256,
            artifact_run_root=artifact_run_root,
            process_run_root=process_run_root,
        )


class LocalProcessWorkerAdapter:
    """Worker-local S1B adapter; transport and PX4 mapping remain separate ports."""

    def __init__(
        self,
        guard: S0WorkerGuard,
        supervisor: PosixProcessSupervisor,
        specs: tuple[ProcessSpec, ...],
        world_control: SimulationWorldControl,
    ) -> None:
        self.guard = guard
        self.supervisor = supervisor
        self.specs = specs
        self.world_control = world_control

    def start(self, run: QualificationRun) -> WorkerStartResult:
        artifact_root = next(
            item.wsl_path
            for item in self.guard.profile.storage.mutable_paths
            if item.identifier == "artifacts"
        )
        artifact_run_root = artifact_root / "s1" / "runs" / run.run_id
        process_run_root = PurePosixPath(self.supervisor.runtime_root.as_posix()) / run.run_id
        admission = self.guard.admit(run, artifact_run_root, process_run_root)
        records = self.supervisor.start(run.run_id, self.specs)
        return WorkerStartResult(
            provider_ids=tuple(record.process_id for record in records),
            profile_sha256=admission.profile_sha256,
        )

    def pause(self, run: QualificationRun) -> None:
        self.world_control.pause(run.run_id)

    def resume(self, run: QualificationRun) -> None:
        self.world_control.resume(run.run_id)

    def step(self, run: QualificationRun, step_count: int) -> int:
        return self.world_control.step(run.run_id, step_count)

    def stop(self, run: QualificationRun) -> WorkerStopResult:
        result = self.supervisor.stop()
        return WorkerStopResult(
            stopped_provider_ids=result.stopped_process_ids,
            residue_provider_ids=result.residue_process_ids,
        )

    def reconcile(self, run: QualificationRun) -> WorkerStopResult:
        return self.stop(run)