"""Deterministic schema-6 protocol pairs for fixed validation banks."""

from __future__ import annotations

import hashlib
from dataclasses import dataclass, replace
from typing import Callable

from pff_pk.data.studies import Study, StudyFamily
from pff_pk.tasks.batch import FlowTask
from pff_pk.tasks.config import TaskConfig
from pff_pk.tasks.generation import sample_task


class ProtocolPairUnavailable(ValueError):
    """Raised when one family cannot instantiate a requested validation case."""


ArmPredicate = Callable[[Study], bool]


def _kind(name: str) -> ArmPredicate:
    return lambda arm: arm.metadata.get("protocol_arm_kind") == name


def _global_direction(direction: str) -> ArmPredicate:
    def predicate(arm: Study) -> bool:
        if arm.metadata.get("protocol_arm_kind") != "global_scale":
            return False
        scale = arm.metadata.get("relative_global_scale")
        if scale is None:
            return False
        return float(scale) < 1.0 if direction == "lower" else float(scale) > 1.0

    return predicate


def _add_on_position(position: str) -> ArmPredicate:
    def predicate(arm: Study) -> bool:
        if arm.metadata.get("protocol_arm_kind") != "add_on":
            return False
        event = arm.metadata.get("added_dose")
        if not isinstance(event, dict) or event.get("time") is None:
            return False
        return float(event["time"]) < 0.5 if position == "early" else float(event["time"]) >= 0.5

    return predicate


@dataclass(frozen=True)
class ProtocolPairCase:
    name: str
    context: ArmPredicate
    target: ArmPredicate
    same_arm: bool = False
    distinct_arms: bool = False


PROTOCOL_PAIR_CASES = (
    ProtocolPairCase("reference_identity", _kind("reference"), _kind("reference"), same_arm=True),
    ProtocolPairCase("global_identity", _kind("global_scale"), _kind("global_scale"), same_arm=True),
    ProtocolPairCase("add_on_identity", _kind("add_on"), _kind("add_on"), same_arm=True),
    ProtocolPairCase("reference_to_lower_scale", _kind("reference"), _global_direction("lower")),
    ProtocolPairCase("reference_to_higher_scale", _kind("reference"), _global_direction("higher")),
    ProtocolPairCase("lower_scale_to_reference", _global_direction("lower"), _kind("reference")),
    ProtocolPairCase("higher_scale_to_reference", _global_direction("higher"), _kind("reference")),
    ProtocolPairCase(
        "global_scale_cross",
        _kind("global_scale"),
        _kind("global_scale"),
        distinct_arms=True,
    ),
    ProtocolPairCase("reference_to_early_add_on", _kind("reference"), _add_on_position("early")),
    ProtocolPairCase("reference_to_late_add_on", _kind("reference"), _add_on_position("late")),
    ProtocolPairCase("add_on_to_reference", _kind("add_on"), _kind("reference")),
    ProtocolPairCase("add_on_cross", _kind("add_on"), _kind("add_on"), distinct_arms=True),
    ProtocolPairCase("global_to_add_on", _kind("global_scale"), _kind("add_on")),
    ProtocolPairCase("add_on_to_global", _kind("add_on"), _kind("global_scale")),
)


def _choice(values: list[int], token: str) -> int:
    digest = hashlib.blake2b(token.encode(), digest_size=8).digest()
    return values[int.from_bytes(digest, "little") % len(values)]


def select_protocol_pair(
    family: StudyFamily,
    case: ProtocolPairCase,
    *,
    draw: int = 0,
) -> tuple[int, int]:
    """Select a reproducible ordered arm pair satisfying ``case``."""
    contexts = [index for index, arm in enumerate(family.arms) if case.context(arm)]
    targets = [index for index, arm in enumerate(family.arms) if case.target(arm)]
    if case.same_arm:
        common = sorted(set(contexts) & set(targets))
        if not common:
            raise ProtocolPairUnavailable(f"{family.identifier} lacks {case.name}")
        selected = _choice(common, f"{family.identifier}:{case.name}:{draw}:same")
        return selected, selected
    pairs = [
        (context, target)
        for context in contexts
        for target in targets
        if not case.distinct_arms or context != target
    ]
    if not pairs:
        raise ProtocolPairUnavailable(f"{family.identifier} lacks {case.name}")
    return pairs[
        _choice(list(range(len(pairs))), f"{family.identifier}:{case.name}:{draw}:pair")
    ]


@dataclass(frozen=True)
class FixedProtocolValidationSampler:
    """Build one reproducible protocol-pair/acquisition validation stratum."""

    case: ProtocolPairCase
    acquisition: str

    def __call__(
        self,
        study: Study | StudyFamily,
        config: TaskConfig,
        *,
        draw: int = 0,
    ) -> FlowTask:
        if not isinstance(study, StudyFamily):
            raise ProtocolPairUnavailable("v6 protocol validation requires schema-6 families")
        pair = select_protocol_pair(study, self.case, draw=draw)
        fixed = replace(config, acquisition_weights={self.acquisition: 1.0})
        return sample_task(study, fixed, draw=draw, protocol_arm_indices=pair)


def protocol_validation_definitions() -> tuple[tuple[str, float, FixedProtocolValidationSampler], ...]:
    """Return the balanced 14-pair by 3-acquisition validation design."""
    acquisitions = ("unscheduled", "exact", "pseudo_scheduled")
    count = len(PROTOCOL_PAIR_CASES) * len(acquisitions)
    return tuple(
        (
            f"{case.name}__{acquisition}",
            1.0 / count,
            FixedProtocolValidationSampler(case, acquisition),
        )
        for case in PROTOCOL_PAIR_CASES
        for acquisition in acquisitions
    )