"""CPU scalar reference backend for canonical Wave-4 parity."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from copy import deepcopy
from typing import Any

from src.runtime.backend import BackendCapabilities, BackendState, RunResult, StepResult
from src.verification.parity import config_fingerprint, execution_fingerprint

from .backend import cpu_recurrent_reference
from .state import (
    canonical_int,
    prefix_inputs,
    recurrent_inputs_from_mapping,
    recurrent_inputs_to_mapping,
    set_external_tick_config,
    stable_digest,
    step_payload,
    validate_checkpoint,
)


class CPUReferenceBackend:
    backend_name = "cpu-reference"
    backend_version = "wave4-scalar-reference-v1"

    def __init__(self) -> None:
        self._config: dict[str, object] | None = None
        self._seed = 0
        self._tick = 0

    def initialize(self, config: Mapping[str, object], seed: int) -> None:
        inputs = recurrent_inputs_from_mapping(config, seed=seed)
        self._config = recurrent_inputs_to_mapping(inputs)
        self._seed = seed
        self._tick = 0

    def _require_config(self) -> dict[str, object]:
        if self._config is None:
            raise RuntimeError("backend is not initialized")
        return self._config

    def _execute_prefix(self, total_ticks: int) -> Mapping[str, object]:
        inputs = recurrent_inputs_from_mapping(self._require_config(), seed=self._seed)
        return cpu_recurrent_reference(prefix_inputs(inputs, total_ticks))

    def _step_from_outputs(
        self, outputs: Mapping[str, object], tick: int
    ) -> StepResult:
        config = self._require_config()
        payload = step_payload(
            outputs,
            tick=tick,
            n_neurons=canonical_int(config["n_neurons"], field="n_neurons"),
        )
        spikes_raw = payload["spikes"]
        if not isinstance(spikes_raw, list):
            raise ValueError("spike payload must be a list")
        spikes = tuple(
            index for index, value in enumerate(spikes_raw) if int(value) == 1
        )
        return StepResult(
            tick=tick + 1,
            spikes=spikes,
            state_digest=stable_digest(payload),
            metrics={
                "voltage": payload["voltage"],
                "adaptation": payload["adaptation"],
                "spike_count": len(spikes),
            },
        )

    def set_external_tick(self, tick: int, currents: Sequence[float]) -> None:
        if tick != self._tick:
            raise ValueError("live external input tick must equal continuation cursor")
        set_external_tick_config(self._require_config(), tick=tick, currents=currents)

    def step(self, tick: int) -> StepResult:
        if tick != self._tick:
            raise ValueError("tick must equal the backend continuation cursor")
        config = self._require_config()
        if tick >= canonical_int(config["ticks"], field="ticks"):
            raise ValueError("configured tick limit reached")
        outputs = self._execute_prefix(tick + 1)
        self._tick = tick + 1
        return self._step_from_outputs(outputs, tick)

    def run(self, ticks: int) -> RunResult:
        if type(ticks) is not int or ticks < 1:
            raise ValueError("ticks must be a positive int")
        config = self._require_config()
        total = self._tick + ticks
        if total > canonical_int(config["ticks"], field="ticks"):
            raise ValueError("requested run exceeds configured tick limit")
        start = self._tick
        outputs = self._execute_prefix(total)
        steps = tuple(
            self._step_from_outputs(outputs, tick) for tick in range(start, total)
        )
        self._tick = total
        return RunResult(
            ticks_requested=ticks,
            steps=steps,
            final_state=self.snapshot(),
            execution_fingerprint=execution_fingerprint(
                seed=self._seed,
                config_hash=config_fingerprint(config),
                backend_name=self.backend_name,
                backend_version=self.backend_version,
                ticks=total,
            ),
        )

    def snapshot(self) -> BackendState:
        config = self._require_config()
        payload: dict[str, Any] = {
            "backend": self.backend_version,
            "seed": self._seed,
            "tick": self._tick,
            "config": deepcopy(config),
            "replay_required": True,
        }
        return BackendState(
            tick=self._tick,
            payload=payload,
            state_digest=stable_digest(payload),
        )

    def restore(self, state: BackendState) -> None:
        raw_config, raw_seed = validate_checkpoint(state, self.backend_version)
        self.initialize(raw_config, raw_seed)
        self._tick = state.tick

    def capabilities(self) -> BackendCapabilities:
        return BackendCapabilities(
            supports_recurrent=True,
            supports_plasticity=True,
            supports_pan_hyperstate=False,
            supports_structural_plasticity=False,
            max_neurons=4096,
            max_ticks=2000,
            max_edges=65536,
            deterministic=True,
            plasticity_semantics="REFERENCE_MATCH_TARGET_DRAFT",
            execution_mode="BOUNDED_REPLAY_REFERENCE",
            supports_live_external_input=True,
        )