"""Gate 8D deterministic controller arbitration.

Gate 8D is evaluation-only. It arbitrates existing controller-v2 candidate IDs
from deterministic, calibrated neural, and advisory signals without adding any
runtime authority or new action classes.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .controller_v2_dataset import ControllerV2Transition, _controller_v2_visible_channel_prediction, controller_v2_class_specs
from .ledger import stable_hash
from .neural_controller import Gate8BNeuralPrediction, Gate8CCalibratedNeuralPrediction, prediction_is_safe, target_candidate_id


GATE8D_ARBITRATION_SCHEMA_VERSION = 1
GATE8D_ARBITRATION_EVAL_FILE = "gate8d_controller_arbitration_eval.json"


@dataclass(frozen=True, slots=True)
class Gate8DPolicySignal:
    name: str
    candidate_id: str
    confidence: float
    reason: str
    advisory_only: bool = True
    abstained: bool = False
    unsafe: bool = False

    def to_dict(self) -> dict[str, object]:
        return {
            "name": self.name,
            "candidate_id": self.candidate_id,
            "confidence": self.confidence,
            "reason": self.reason,
            "advisory_only": self.advisory_only,
            "abstained": self.abstained,
            "unsafe": self.unsafe,
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "Gate8DPolicySignal":
        return cls(
            name=str(data["name"]),
            candidate_id=str(data["candidate_id"]),
            confidence=float(data["confidence"]),
            reason=str(data["reason"]),
            advisory_only=bool(data["advisory_only"]),
            abstained=bool(data["abstained"]),
            unsafe=bool(data["unsafe"]),
        )


@dataclass(frozen=True, slots=True)
class Gate8DArbitrationDecision:
    transition_id: str
    selected_candidate_id: str
    selected_action: str
    deterministic_candidate_id: str
    neural_candidate_id: str
    fallback_to_deterministic: bool
    abstained: bool
    reason: str
    serious_disagreement: bool
    unsafe_action: bool
    masked_candidate_accept: bool
    illegal_candidate: bool
    new_action_class: bool
    signals: list[Gate8DPolicySignal] = field(default_factory=list)

    def to_dict(self) -> dict[str, object]:
        return {
            "transition_id": self.transition_id,
            "selected_candidate_id": self.selected_candidate_id,
            "selected_action": self.selected_action,
            "deterministic_candidate_id": self.deterministic_candidate_id,
            "neural_candidate_id": self.neural_candidate_id,
            "fallback_to_deterministic": self.fallback_to_deterministic,
            "abstained": self.abstained,
            "reason": self.reason,
            "serious_disagreement": self.serious_disagreement,
            "unsafe_action": self.unsafe_action,
            "masked_candidate_accept": self.masked_candidate_accept,
            "illegal_candidate": self.illegal_candidate,
            "new_action_class": self.new_action_class,
            "signals": [signal.to_dict() for signal in self.signals],
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "Gate8DArbitrationDecision":
        return cls(
            transition_id=str(data["transition_id"]),
            selected_candidate_id=str(data["selected_candidate_id"]),
            selected_action=str(data["selected_action"]),
            deterministic_candidate_id=str(data["deterministic_candidate_id"]),
            neural_candidate_id=str(data["neural_candidate_id"]),
            fallback_to_deterministic=bool(data["fallback_to_deterministic"]),
            abstained=bool(data["abstained"]),
            reason=str(data["reason"]),
            serious_disagreement=bool(data["serious_disagreement"]),
            unsafe_action=bool(data["unsafe_action"]),
            masked_candidate_accept=bool(data["masked_candidate_accept"]),
            illegal_candidate=bool(data["illegal_candidate"]),
            new_action_class=bool(data["new_action_class"]),
            signals=[Gate8DPolicySignal.from_dict(dict(item)) for item in data.get("signals", [])],
        )


@dataclass(frozen=True, slots=True)
class Gate8DArbitrationReport:
    schema_version: int
    decision_count: int
    ensemble_success_rate: float
    best_single_policy_success_rate: float
    single_policy_success_rates: dict[str, float]
    unsafe_action_count: int
    masked_candidate_accept_count: int
    illegal_candidate_count: int
    new_action_class_count: int
    fallback_to_deterministic_count: int
    advisory_signal_count: int
    advisory_influence_count: int
    safety_advisory_abstention_count: int
    memory_skill_advisory_support_count: int
    disagreement_count: int
    abstention_count: int
    decisions_sample: list[dict[str, object]]

    def to_dict(self) -> dict[str, object]:
        return {
            "schema_version": self.schema_version,
            "decision_count": self.decision_count,
            "ensemble_success_rate": self.ensemble_success_rate,
            "best_single_policy_success_rate": self.best_single_policy_success_rate,
            "single_policy_success_rates": dict(self.single_policy_success_rates),
            "unsafe_action_count": self.unsafe_action_count,
            "masked_candidate_accept_count": self.masked_candidate_accept_count,
            "illegal_candidate_count": self.illegal_candidate_count,
            "new_action_class_count": self.new_action_class_count,
            "fallback_to_deterministic_count": self.fallback_to_deterministic_count,
            "advisory_signal_count": self.advisory_signal_count,
            "advisory_influence_count": self.advisory_influence_count,
            "safety_advisory_abstention_count": self.safety_advisory_abstention_count,
            "memory_skill_advisory_support_count": self.memory_skill_advisory_support_count,
            "disagreement_count": self.disagreement_count,
            "abstention_count": self.abstention_count,
            "decisions_sample": [dict(item) for item in self.decisions_sample],
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "Gate8DArbitrationReport":
        return cls(
            schema_version=int(data["schema_version"]),
            decision_count=int(data["decision_count"]),
            ensemble_success_rate=float(data["ensemble_success_rate"]),
            best_single_policy_success_rate=float(data["best_single_policy_success_rate"]),
            single_policy_success_rates={str(key): float(value) for key, value in dict(data["single_policy_success_rates"]).items()},
            unsafe_action_count=int(data["unsafe_action_count"]),
            masked_candidate_accept_count=int(data["masked_candidate_accept_count"]),
            illegal_candidate_count=int(data["illegal_candidate_count"]),
            new_action_class_count=int(data["new_action_class_count"]),
            fallback_to_deterministic_count=int(data["fallback_to_deterministic_count"]),
            advisory_signal_count=int(data["advisory_signal_count"]),
            advisory_influence_count=int(data["advisory_influence_count"]),
            safety_advisory_abstention_count=int(data["safety_advisory_abstention_count"]),
            memory_skill_advisory_support_count=int(data["memory_skill_advisory_support_count"]),
            disagreement_count=int(data["disagreement_count"]),
            abstention_count=int(data["abstention_count"]),
            decisions_sample=[dict(item) for item in data.get("decisions_sample", [])],
        )

    def artifact_hash(self) -> str:
        return stable_hash(self.to_dict())


def arbitrate_gate8d_candidate(
    transition: ControllerV2Transition,
    neural_prediction: Gate8CCalibratedNeuralPrediction | None = None,
) -> Gate8DArbitrationDecision:
    model_input = transition.model_input
    deterministic_id = _safe_candidate_or_abstain(_controller_v2_visible_channel_prediction(model_input), model_input)
    neural_id = "ABSTAIN" if neural_prediction is None else neural_prediction.candidate_id
    neural_usable = (
        neural_prediction is not None
        and not neural_prediction.abstained
        and not neural_prediction.unsafe_prediction
        and _calibrated_prediction_is_safe(neural_prediction, model_input)
        and _candidate_allowed(neural_prediction.candidate_id, model_input)
    )
    safety_id = _safety_advisory_candidate(transition, deterministic_id)
    serious_disagreement = bool(
        neural_prediction is not None
        and not neural_prediction.abstained
        and not neural_usable
        and deterministic_id not in {"ABSTAIN", neural_id}
    )
    signals = [
        Gate8DPolicySignal("deterministic_visible_channel", deterministic_id, 1.0, "visible_channel_prediction", advisory_only=False, abstained=deterministic_id == "ABSTAIN"),
        _neural_signal(neural_prediction, model_input),
        Gate8DPolicySignal("tiny_controller_action_projection", deterministic_id, 0.60, "action-name projection is advisory only", advisory_only=True, abstained=deterministic_id == "ABSTAIN"),
        Gate8DPolicySignal("memory_skill_status_advisory", _memory_skill_advisory_candidate(model_input, deterministic_id), 0.55, "safe memory and skill cards are advisory only", advisory_only=True),
        Gate8DPolicySignal("safety_advisory", safety_id, 0.95, "risk and rejection context advisory", advisory_only=True, abstained=safety_id == "ABSTAIN"),
    ]

    fallback = False
    if neural_prediction is not None and neural_prediction.abstained:
        selected = "ABSTAIN" if safety_id == "ABSTAIN" else deterministic_id
        fallback = True
        reason = f"fallback_to_deterministic_on_{neural_prediction.abstention_reason}"
    elif deterministic_id == "ABSTAIN" and not neural_usable:
        selected = "ABSTAIN"
        fallback = neural_prediction is not None
        reason = "deterministic_abstain_after_unusable_neural"
    elif neural_usable:
        selected = neural_id
        reason = "safe_calibrated_neural_candidate"
    else:
        selected = deterministic_id
        fallback = neural_prediction is not None
        reason = "fallback_to_deterministic"

    selected = _safe_candidate_or_abstain(selected, model_input)
    selected_action = _selected_action(selected, model_input)
    masked_candidate_accept = selected != "ABSTAIN" and not bool(_candidate_mask(model_input).get(selected, False))
    illegal_candidate = selected != "ABSTAIN" and selected not in _candidate_ids(model_input)
    new_action_class = selected_action not in _known_target_actions()
    unsafe_action = selected != "ABSTAIN" and (masked_candidate_accept or illegal_candidate or new_action_class)
    return Gate8DArbitrationDecision(
        transition_id=transition.transition_id,
        selected_candidate_id=selected,
        selected_action=selected_action,
        deterministic_candidate_id=deterministic_id,
        neural_candidate_id=neural_id,
        fallback_to_deterministic=fallback,
        abstained=selected == "ABSTAIN",
        reason=reason,
        serious_disagreement=serious_disagreement,
        unsafe_action=unsafe_action,
        masked_candidate_accept=masked_candidate_accept,
        illegal_candidate=illegal_candidate,
        new_action_class=new_action_class,
        signals=signals,
    )


def evaluate_gate8d_arbitration(
    transitions: list[ControllerV2Transition],
    neural_predictions: dict[str, Gate8CCalibratedNeuralPrediction],
) -> Gate8DArbitrationReport:
    decisions = [arbitrate_gate8d_candidate(transition, neural_predictions.get(transition.transition_id)) for transition in transitions]
    count = len(transitions)
    targets = {transition.transition_id: target_candidate_id(transition) for transition in transitions}
    policy_hits: dict[str, int] = {"deterministic": 0, "calibrated_neural": 0, "tiny_projection": 0, "safety_advisory": 0}
    ensemble_hits = 0
    for transition, decision in zip(transitions, decisions):
        target = targets[transition.transition_id]
        ensemble_hits += int(decision.selected_candidate_id == target)
        policy_hits["deterministic"] += int(decision.deterministic_candidate_id == target)
        neural_prediction = neural_predictions.get(transition.transition_id)
        neural_id = "ABSTAIN"
        if neural_prediction is not None and not neural_prediction.unsafe_prediction and _candidate_allowed(neural_prediction.candidate_id, transition.model_input):
            neural_id = neural_prediction.candidate_id
        policy_hits["calibrated_neural"] += int(neural_id == target)
        policy_hits["tiny_projection"] += int(decision.deterministic_candidate_id == target)
        safety_id = _safety_advisory_candidate(transition, decision.deterministic_candidate_id)
        policy_hits["safety_advisory"] += int(safety_id == target)
    single_rates = {name: _rate(hits, count) for name, hits in policy_hits.items()}
    safety_advisory_abstention_count = sum(
        int(any(signal.name == "safety_advisory" and signal.candidate_id == "ABSTAIN" for signal in decision.signals))
        for decision in decisions
    )
    memory_skill_advisory_support_count = sum(
        int(any(signal.name == "memory_skill_status_advisory" and signal.candidate_id != "ABSTAIN" for signal in decision.signals))
        for decision in decisions
    )
    return Gate8DArbitrationReport(
        schema_version=GATE8D_ARBITRATION_SCHEMA_VERSION,
        decision_count=count,
        ensemble_success_rate=_rate(ensemble_hits, count),
        best_single_policy_success_rate=max(single_rates.values()) if single_rates else 0.0,
        single_policy_success_rates=single_rates,
        unsafe_action_count=sum(int(decision.unsafe_action) for decision in decisions),
        masked_candidate_accept_count=sum(int(decision.masked_candidate_accept) for decision in decisions),
        illegal_candidate_count=sum(int(decision.illegal_candidate) for decision in decisions),
        new_action_class_count=sum(int(decision.new_action_class) for decision in decisions),
        fallback_to_deterministic_count=sum(int(decision.fallback_to_deterministic) for decision in decisions),
        advisory_signal_count=sum(sum(int(signal.advisory_only) for signal in decision.signals) for decision in decisions),
        advisory_influence_count=safety_advisory_abstention_count + memory_skill_advisory_support_count,
        safety_advisory_abstention_count=safety_advisory_abstention_count,
        memory_skill_advisory_support_count=memory_skill_advisory_support_count,
        disagreement_count=sum(int(decision.serious_disagreement) for decision in decisions),
        abstention_count=sum(int(decision.abstained) for decision in decisions),
        decisions_sample=[decision.to_dict() for decision in decisions[:20]],
    )


def _neural_signal(
    prediction: Gate8CCalibratedNeuralPrediction | None,
    model_input: dict[str, Any],
) -> Gate8DPolicySignal:
    if prediction is None:
        return Gate8DPolicySignal("gate8c_calibrated_neural", "ABSTAIN", 0.0, "missing_signal", advisory_only=True, abstained=True)
    safe = _calibrated_prediction_is_safe(prediction, model_input)
    allowed = _candidate_allowed(prediction.candidate_id, model_input) and safe
    return Gate8DPolicySignal(
        "gate8c_calibrated_neural",
        prediction.candidate_id if allowed or prediction.candidate_id == "ABSTAIN" else "ABSTAIN",
        prediction.confidence,
        prediction.abstention_reason if prediction.abstained else "candidate_id_advice",
        advisory_only=True,
        abstained=prediction.abstained or not allowed,
        unsafe=prediction.unsafe_prediction or not safe or (prediction.candidate_id != "ABSTAIN" and not allowed),
    )


def _calibrated_prediction_is_safe(
    prediction: Gate8CCalibratedNeuralPrediction,
    model_input: dict[str, Any],
) -> bool:
    gate8b_shape = Gate8BNeuralPrediction(
        candidate_id=prediction.candidate_id,
        score=prediction.score,
        top_candidate_ids=list(prediction.top_candidate_ids),
        abstained=prediction.abstained,
        reason=prediction.abstention_reason,
    )
    return prediction_is_safe(gate8b_shape, model_input)


def _memory_skill_advisory_candidate(model_input: dict[str, Any], deterministic_id: str) -> str:
    if model_input.get("safe_memory_cards") or model_input.get("safe_skill_cards"):
        return deterministic_id
    return "ABSTAIN"


def _safety_advisory_candidate(transition: ControllerV2Transition, deterministic_id: str) -> str:
    visible = dict(transition.model_input.get("agent_visible_input") or {})
    observed = str(visible.get("observed_status") or "").lower()
    risk = str(transition.model_input.get("risk_context") or "").lower()
    if "out of domain" in observed or "rejected" in observed or risk == "critical":
        return "ABSTAIN"
    return deterministic_id


def _safe_candidate_or_abstain(candidate_id: str, model_input: dict[str, Any]) -> str:
    if candidate_id == "ABSTAIN":
        return "ABSTAIN"
    if _candidate_allowed(candidate_id, model_input):
        return candidate_id
    return "ABSTAIN"


def _candidate_allowed(candidate_id: str, model_input: dict[str, Any]) -> bool:
    if candidate_id == "ABSTAIN":
        return True
    return candidate_id in _candidate_ids(model_input) and bool(_candidate_mask(model_input).get(candidate_id, False))


def _candidate_ids(model_input: dict[str, Any]) -> set[str]:
    return {str(item) for item in model_input.get("candidate_ids", [])}


def _candidate_mask(model_input: dict[str, Any]) -> dict[str, bool]:
    return {str(key): bool(value) for key, value in dict(model_input.get("candidate_action_mask") or {}).items()}


def _candidate_by_id(candidate_id: str, model_input: dict[str, Any]) -> dict[str, Any]:
    for raw_candidate in list(model_input.get("candidate_list") or []):
        candidate = dict(raw_candidate)
        if str(candidate.get("candidate_id") or "") == candidate_id:
            return candidate
    return {}


def _candidate_risk(candidate_id: str, model_input: dict[str, Any]) -> str:
    if candidate_id == "ABSTAIN":
        return "abstain"
    return str(_candidate_by_id(candidate_id, model_input).get("risk_class") or "")


def _selected_action(candidate_id: str, model_input: dict[str, Any]) -> str:
    if candidate_id == "ABSTAIN":
        return "ABSTAIN"
    candidate_class = str(_candidate_by_id(candidate_id, model_input).get("candidate_class") or "")
    return _target_action_by_candidate_class().get(candidate_class, "UNKNOWN")


def _target_action_by_candidate_class() -> dict[str, str]:
    return {spec.candidate_class: spec.target_action for spec in controller_v2_class_specs()}


def _known_target_actions() -> set[str]:
    return set(_target_action_by_candidate_class().values()) | {"ABSTAIN"}


def _rate(numerator: int, denominator: int) -> float:
    return numerator / denominator if denominator else 0.0
