"""Single ordered orchestration shell for the completed blind dataset stages."""

from __future__ import annotations

from typing import Any, Protocol

from pydantic import BaseModel

from gcmd_classifier.datasets.errors import DatasetOrchestrationError
from gcmd_classifier.datasets.models import (
    DatasetClassificationOutcome,
    DatasetEvidencePacket,
    ResolvedCMRCollection,
)
from gcmd_classifier.datasets.orchestration import DatasetRunController, DurableInferenceClient
from gcmd_classifier.datasets.runtime_models import (
    ArtifactReference,
    ArtifactType,
    BlindPipelineTerminal,
    CheckpointStatus,
    DatasetRuntimeFailure,
    PersistedExpertSourceBoundary,
    PipelineStage,
    TerminalOutcome,
)
from gcmd_classifier.datasets.runtime_persistence import (
    DatasetArtifactStore,
    canonical_json_bytes,
    persist_expert_source_boundary,
    verify_persisted_blind_result,
)
from gcmd_classifier.llm.base import ModelClient


class BlindStageServices(Protocol):
    """Adapters call the already implemented stage functions; they add no stage logic."""

    def execute(
        self,
        stage: PipelineStage,
        context: dict[PipelineStage, Any],
        model_client: ModelClient,
    ) -> Any:
        """Execute one named stage and return its existing typed stage result."""


_ORDER = (
    PipelineStage.CMR_SOURCE,
    PipelineStage.IDENTITY,
    PipelineStage.BLIND_BOUNDARY,
    PipelineStage.README_DISCOVERY,
    PipelineStage.README_SELECTION,
    PipelineStage.README_RETRIEVAL,
    PipelineStage.SOURCE_PERSISTENCE,
    PipelineStage.EXTRACTION,
    PipelineStage.PRODUCT_RESOLUTION,
    PipelineStage.EVIDENCE_SELECTION,
    PipelineStage.EVIDENCE_PACKET,
    PipelineStage.CLASSIFICATION,
)

_ARTIFACT_TYPE = {
    PipelineStage.CMR_SOURCE: ArtifactType.CMR_RESPONSE,
    PipelineStage.IDENTITY: ArtifactType.COLLECTION_SNAPSHOT,
    PipelineStage.BLIND_BOUNDARY: ArtifactType.COLLECTION_SNAPSHOT,
    PipelineStage.README_DISCOVERY: ArtifactType.README_METADATA,
    PipelineStage.README_SELECTION: ArtifactType.README_METADATA,
    PipelineStage.README_RETRIEVAL: ArtifactType.README_PDF,
    PipelineStage.SOURCE_PERSISTENCE: ArtifactType.README_METADATA,
    PipelineStage.EXTRACTION: ArtifactType.EXTRACTION_MANIFEST,
    PipelineStage.PRODUCT_RESOLUTION: ArtifactType.PRODUCT_RESOLUTION,
    PipelineStage.EVIDENCE_SELECTION: ArtifactType.EVIDENCE_SELECTION,
    PipelineStage.EVIDENCE_PACKET: ArtifactType.EVIDENCE_PACKET,
    PipelineStage.CLASSIFICATION: ArtifactType.BLIND_RESULT,
}


class BlindDatasetPipeline:
    """Execute the approved blind order with durable state and artifact boundaries."""

    def __init__(
        self,
        *,
        controller: DatasetRunController,
        store: DatasetArtifactStore,
        services: BlindStageServices,
        model_client: ModelClient,
        vocabulary_hash: str,
    ) -> None:
        self.controller = controller
        self.store = store
        self.services = services
        self.model_client = DurableInferenceClient(model_client, controller)
        self.vocabulary_hash = vocabulary_hash
        self._source_boundary: PersistedExpertSourceBoundary | None = None

    def run(self) -> BlindPipelineTerminal:
        context: dict[PipelineStage, Any] = {}
        references: dict[PipelineStage, ArtifactReference] = {}
        try:
            for stage in _ORDER:
                self.controller.begin(stage, f"Begin {stage.value}.")
                value = self.services.execute(stage, context, self.model_client)
                if stage is PipelineStage.IDENTITY:
                    if not isinstance(value, ResolvedCMRCollection):
                        raise DatasetOrchestrationError(
                            "EXPERT_SOURCE_BINDING_MISSING",
                            "Identity stage must return the resolved CMR collection.",
                        )
                    self._source_boundary = persist_expert_source_boundary(
                        store=self.store,
                        run=self.controller.state.run,
                        session_id=self.controller.state.session.session_id,
                        collection=value,
                        cmr_source_artifact=references[PipelineStage.CMR_SOURCE],
                    )
                    context.clear()
                    context[PipelineStage.CMR_SOURCE] = references[PipelineStage.CMR_SOURCE]
                    value = self._source_boundary.blind_view
                if stage is PipelineStage.README_DISCOVERY and value is None:
                    self.controller.transition(
                        stage, CheckpointStatus.SKIPPED, "No README candidate was discovered."
                    )
                    return self._terminal(TerminalOutcome.COMPLETED_SKIPPED)
                if stage is PipelineStage.SOURCE_PERSISTENCE:
                    # The stage exists as an explicit service checkpoint; exact source bytes
                    # were already published at their acquisition boundaries.
                    reference = references[PipelineStage.README_RETRIEVAL]
                else:
                    reference = self._persist(stage, value, context)
                references[stage] = reference
                context[stage] = value
                checkpoint_artifacts = (reference,)
                if stage is PipelineStage.IDENTITY and self._source_boundary is not None:
                    checkpoint_artifacts = (
                        reference,
                        self._source_boundary.cmr_source_artifact,
                        self._source_boundary.sealed_expert_artifact,
                        self._source_boundary.expert_source_binding_artifact,
                    )
                self.controller.complete(
                    stage, f"Completed {stage.value}.", artifacts=checkpoint_artifacts
                )
                if stage is PipelineStage.CLASSIFICATION:
                    self.controller.mark_blind_result(produced=True)
                    self.controller.begin(
                        PipelineStage.BLIND_RESULT_PERSISTENCE,
                        "Persist blind result.",
                    )
                    self.controller.complete(
                        PipelineStage.BLIND_RESULT_PERSISTENCE,
                        "Blind result atomically persisted.",
                        artifacts=(reference,),
                    )
                    self.controller.mark_blind_result(persisted=True)
            result = context[PipelineStage.CLASSIFICATION]
            packet = self._packet(context[PipelineStage.EVIDENCE_PACKET])
            verified = verify_persisted_blind_result(
                store=self.store,
                run=self.controller.state.run,
                session_id=self.controller.state.session.session_id,
                reference=references[PipelineStage.CLASSIFICATION],
                expected_packet_sha256=packet.packet_sha256,
                expected_native_id=result.identity.native_id,
                vocabulary_hash=self.vocabulary_hash,
                source_boundary=self._required_source_boundary(),
            )
            self.controller.begin(
                PipelineStage.BLIND_RESULT_VERIFICATION, "Verify persisted blind result."
            )
            self.controller.complete(
                PipelineStage.BLIND_RESULT_VERIFICATION,
                "Persisted blind-result hash verified.",
                artifacts=(verified.artifact,),
            )
            self.controller.mark_blind_result(verified=True)
            terminal = (
                TerminalOutcome.COMPLETED_NOT_CLASSIFIED
                if result.classification_outcome is DatasetClassificationOutcome.NOT_CLASSIFIED
                else TerminalOutcome.COMPLETED_CLASSIFIED
            )
            return self._terminal(terminal, verified=verified)
        except Exception as exc:
            attempted = self.controller.state.classification_attempted
            failure = DatasetRuntimeFailure(
                code=getattr(exc, "code", "BLIND_PIPELINE_FAILED"),
                stage=self.controller.state.checkpoints[-1].stage,
                message="Blind dataset pipeline stage failed.",
                retryable=False,
                classification_attempted=attempted,
            )
            return self._terminal(
                TerminalOutcome.POST_INFERENCE_FAILURE
                if attempted
                else TerminalOutcome.PRE_INFERENCE_FAILURE,
                failure=failure,
            )

    def _required_source_boundary(self) -> PersistedExpertSourceBoundary:
        if self._source_boundary is None:
            raise DatasetOrchestrationError(
                "COMPARISON_PROVENANCE_INCOMPLETE",
                "Verified blind results require persisted expert source provenance.",
            )
        return self._source_boundary

    def _persist(
        self, stage: PipelineStage, value: Any, context: dict[PipelineStage, Any]
    ) -> ArtifactReference:
        if stage is PipelineStage.CMR_SOURCE and hasattr(value, "response_bytes"):
            data = value.response_bytes
        elif stage is PipelineStage.README_RETRIEVAL and hasattr(value, "content"):
            data = value.content
        elif isinstance(value, bytes):
            data = value
        elif isinstance(value, BaseModel):
            data = canonical_json_bytes(value)
        elif isinstance(value, dict):
            data = canonical_json_bytes(value)
        else:
            raise DatasetOrchestrationError(
                "UNPERSISTABLE_STAGE_RESULT", "Stage did not return a typed persistable value."
            )
        if stage not in {PipelineStage.CMR_SOURCE} and b'"ScienceKeywords"' in data:
            raise DatasetOrchestrationError(
                "BLIND_BOUNDARY_VIOLATION", "Expert keywords entered a blind artifact."
            )
        upstream = tuple(reference.sha256 for reference in self._upstream_references(context))
        name = (
            "blind_result.json"
            if stage is PipelineStage.CLASSIFICATION
            else f"artifacts/{stage.value}.json"
        )
        if stage is PipelineStage.README_RETRIEVAL:
            name = "selected_readme.pdf"
        return self.store.persist_bytes(
            run=self.controller.state.run,
            session_id=self.controller.state.session.session_id,
            relative_name=name,
            data=data,
            artifact_type=_ARTIFACT_TYPE[stage],
            upstream_sha256=upstream,
        )

    def _upstream_references(
        self, context: dict[PipelineStage, Any]
    ) -> tuple[ArtifactReference, ...]:
        refs: list[ArtifactReference] = []
        for checkpoint in self.controller.state.checkpoints:
            if checkpoint.status is CheckpointStatus.COMPLETED:
                refs.extend(checkpoint.artifact_references)
        return tuple(refs)

    @staticmethod
    def _packet(value: Any) -> DatasetEvidencePacket:
        if isinstance(value, DatasetEvidencePacket):
            return value
        packet = getattr(value, "packet", None)
        if isinstance(packet, DatasetEvidencePacket):
            return packet
        raise DatasetOrchestrationError(
            "PACKET_RESULT_INVALID", "Evidence packet stage returned no typed packet."
        )

    def _terminal(
        self,
        outcome: TerminalOutcome,
        *,
        verified: Any = None,
        failure: DatasetRuntimeFailure | None = None,
    ) -> BlindPipelineTerminal:
        status = (
            CheckpointStatus.SKIPPED
            if outcome is TerminalOutcome.COMPLETED_SKIPPED
            else CheckpointStatus.FAILED
            if failure is not None
            else CheckpointStatus.COMPLETED
        )
        self.controller.transition(
            PipelineStage.TERMINAL,
            status,
            f"Terminal outcome: {outcome.value}.",
            terminal_outcome=outcome,
        )
        return BlindPipelineTerminal(
            state=self.controller.state,
            terminal_outcome=outcome,
            verified_blind_result=verified,
            failure=failure,
        )