"""Stimulation step control converters — domain ↔ proto."""

from __future__ import annotations

from radiens_core._generated import allegoserver_pb2
from radiens_core.models.stim import (
    StimStep,
    StimStepRequest,
    StimStepState,
)

# --- StimStep enum conversion ---

_STIM_STEP_TO_PROTO: dict[StimStep, allegoserver_pb2.StimStep.ValueType] = {
    StimStep.STIM_STEP_SIZE_MIN: allegoserver_pb2.StimStepSizeMin,
    StimStep.STIM_STEP_SIZE_10_NA: allegoserver_pb2.StimStepSize10nA,
    StimStep.STIM_STEP_SIZE_20_NA: allegoserver_pb2.StimStepSize20nA,
    StimStep.STIM_STEP_SIZE_50_NA: allegoserver_pb2.StimStepSize50nA,
    StimStep.STIM_STEP_SIZE_100_NA: allegoserver_pb2.StimStepSize100nA,
    StimStep.STIM_STEP_SIZE_200_NA: allegoserver_pb2.StimStepSize200nA,
    StimStep.STIM_STEP_SIZE_500_NA: allegoserver_pb2.StimStepSize500nA,
    StimStep.STIM_STEP_SIZE_1_UA: allegoserver_pb2.StimStepSize1uA,
    StimStep.STIM_STEP_SIZE_2_UA: allegoserver_pb2.StimStepSize2uA,
    StimStep.STIM_STEP_SIZE_5_UA: allegoserver_pb2.StimStepSize5uA,
    StimStep.STIM_STEP_SIZE_10_UA: allegoserver_pb2.StimStepSize10uA,
    StimStep.STIM_STEP_SIZE_MAX: allegoserver_pb2.StimStepSizeMax,
}

_STIM_STEP_FROM_PROTO: dict[int, StimStep] = {
    int(allegoserver_pb2.StimStepSizeMin): StimStep.STIM_STEP_SIZE_MIN,
    int(allegoserver_pb2.StimStepSize10nA): StimStep.STIM_STEP_SIZE_10_NA,
    int(allegoserver_pb2.StimStepSize20nA): StimStep.STIM_STEP_SIZE_20_NA,
    int(allegoserver_pb2.StimStepSize50nA): StimStep.STIM_STEP_SIZE_50_NA,
    int(allegoserver_pb2.StimStepSize100nA): StimStep.STIM_STEP_SIZE_100_NA,
    int(allegoserver_pb2.StimStepSize200nA): StimStep.STIM_STEP_SIZE_200_NA,
    int(allegoserver_pb2.StimStepSize500nA): StimStep.STIM_STEP_SIZE_500_NA,
    int(allegoserver_pb2.StimStepSize1uA): StimStep.STIM_STEP_SIZE_1_UA,
    int(allegoserver_pb2.StimStepSize2uA): StimStep.STIM_STEP_SIZE_2_UA,
    int(allegoserver_pb2.StimStepSize5uA): StimStep.STIM_STEP_SIZE_5_UA,
    int(allegoserver_pb2.StimStepSize10uA): StimStep.STIM_STEP_SIZE_10_UA,
    int(allegoserver_pb2.StimStepSizeMax): StimStep.STIM_STEP_SIZE_MAX,
}


def stim_step_to_proto(ss: StimStep) -> allegoserver_pb2.StimStep.ValueType:
    """Convert domain StimStep to proto StimStep."""
    return _STIM_STEP_TO_PROTO[ss]


def stim_step_from_proto(pb: int) -> StimStep:
    """Convert proto StimStep int to domain StimStep."""
    return _STIM_STEP_FROM_PROTO[pb]


def stim_step_request_to_proto(
    request: StimStepRequest,
) -> allegoserver_pb2.StimStepMessage:
    """Convert ``StimStepRequest`` domain model → proto message."""
    return allegoserver_pb2.StimStepMessage(
        stimStep=stim_step_to_proto(request.stim_step),
    )


def stim_step_state_from_proto(
    response: allegoserver_pb2.StimStepMessage,
) -> StimStepState:
    """Convert ``StimStepMessage`` proto response → domain model."""
    return StimStepState(
        stim_step=stim_step_from_proto(response.stimStep),
    )


def build_stim_step_proto(step: StimStep) -> allegoserver_pb2.StimStepMessage:
    """Build a StimStepMessage proto directly from a StimStep."""
    return allegoserver_pb2.StimStepMessage(
        stimStep=stim_step_to_proto(step),
    )
