"""Stimulation parameter converter — proto ↔ domain."""

from __future__ import annotations

from radiens_core._generated import allegoserver_pb2
from radiens_core.models.stim import (
    StimParams,
    StimPolarity,
    StimPulseMode,
    StimShape,
    StimTriggerEdgeOrLevel,
    StimTriggerHighOrLow,
)

# --- StimShape ---

_STIM_SHAPE_TO_PROTO: dict[StimShape, allegoserver_pb2.StimShape.ValueType] = {
    StimShape.BIPHASIC: allegoserver_pb2.BIPHASIC,
    StimShape.BIPHASIC_INTERPHASE_DELAY: allegoserver_pb2.BIPHASIC_INTERPHASE_DELAY,
    StimShape.TRIPHASIC: allegoserver_pb2.TRIPHASIC,
    StimShape.MONOPHASIC: allegoserver_pb2.MONOPHASIC,
}

_STIM_SHAPE_FROM_PROTO: dict[int, StimShape] = {
    int(allegoserver_pb2.BIPHASIC): StimShape.BIPHASIC,
    int(allegoserver_pb2.BIPHASIC_INTERPHASE_DELAY): StimShape.BIPHASIC_INTERPHASE_DELAY,
    int(allegoserver_pb2.TRIPHASIC): StimShape.TRIPHASIC,
    int(allegoserver_pb2.MONOPHASIC): StimShape.MONOPHASIC,
}


def stim_shape_to_proto(ss: StimShape) -> allegoserver_pb2.StimShape.ValueType:
    """Convert domain StimShape to proto StimShape."""
    return _STIM_SHAPE_TO_PROTO[ss]


def stim_shape_from_proto(pb: int) -> StimShape:
    """Convert proto StimShape int to domain StimShape."""
    return _STIM_SHAPE_FROM_PROTO[pb]


# --- StimPolarity ---

_STIM_POLARITY_TO_PROTO: dict[StimPolarity, allegoserver_pb2.StimPolarity.ValueType] = {
    StimPolarity.CATHODIC_FIRST: allegoserver_pb2.CATHODIC_FIRST,
    StimPolarity.ANODIC_FIRST: allegoserver_pb2.ANODIC_FIRST,
}

_STIM_POLARITY_FROM_PROTO: dict[int, StimPolarity] = {
    int(allegoserver_pb2.CATHODIC_FIRST): StimPolarity.CATHODIC_FIRST,
    int(allegoserver_pb2.ANODIC_FIRST): StimPolarity.ANODIC_FIRST,
}


def stim_polarity_to_proto(sp: StimPolarity) -> allegoserver_pb2.StimPolarity.ValueType:
    """Convert domain StimPolarity to proto StimPolarity."""
    return _STIM_POLARITY_TO_PROTO[sp]


def stim_polarity_from_proto(pb: int) -> StimPolarity:
    """Convert proto StimPolarity int to domain StimPolarity."""
    return _STIM_POLARITY_FROM_PROTO[pb]


# --- StimTriggerEdgeOrLevel ---

_STIM_TRIGGER_EDGE_OR_LEVEL_TO_PROTO: dict[
    StimTriggerEdgeOrLevel, allegoserver_pb2.StimTriggerEdgeOrLevel.ValueType
] = {
    StimTriggerEdgeOrLevel.EDGE: allegoserver_pb2.ST_EDGE,
    StimTriggerEdgeOrLevel.LEVEL: allegoserver_pb2.ST_LEVEL,
}

_STIM_TRIGGER_EDGE_OR_LEVEL_FROM_PROTO: dict[int, StimTriggerEdgeOrLevel] = {
    int(allegoserver_pb2.ST_EDGE): StimTriggerEdgeOrLevel.EDGE,
    int(allegoserver_pb2.ST_LEVEL): StimTriggerEdgeOrLevel.LEVEL,
}


def stim_trigger_edge_or_level_to_proto(
    ste: StimTriggerEdgeOrLevel,
) -> allegoserver_pb2.StimTriggerEdgeOrLevel.ValueType:
    """Convert domain StimTriggerEdgeOrLevel to proto StimTriggerEdgeOrLevel."""
    return _STIM_TRIGGER_EDGE_OR_LEVEL_TO_PROTO[ste]


def stim_trigger_edge_or_level_from_proto(pb: int) -> StimTriggerEdgeOrLevel:
    """Convert proto StimTriggerEdgeOrLevel int to domain StimTriggerEdgeOrLevel."""
    return _STIM_TRIGGER_EDGE_OR_LEVEL_FROM_PROTO[pb]


# --- StimTriggerHighOrLow ---

_STIM_TRIGGER_HIGH_OR_LOW_TO_PROTO: dict[StimTriggerHighOrLow, allegoserver_pb2.StimTriggerHighOrLow.ValueType] = {
    StimTriggerHighOrLow.HIGH: allegoserver_pb2.ST_HIGH,
    StimTriggerHighOrLow.LOW: allegoserver_pb2.ST_LOW,
}

_STIM_TRIGGER_HIGH_OR_LOW_FROM_PROTO: dict[int, StimTriggerHighOrLow] = {
    int(allegoserver_pb2.ST_HIGH): StimTriggerHighOrLow.HIGH,
    int(allegoserver_pb2.ST_LOW): StimTriggerHighOrLow.LOW,
}


def stim_trigger_level_to_proto(
    stl: StimTriggerHighOrLow,
) -> allegoserver_pb2.StimTriggerHighOrLow.ValueType:
    """Convert domain StimTriggerLevel to proto StimTriggerHighOrLow."""
    return _STIM_TRIGGER_HIGH_OR_LOW_TO_PROTO[stl]


def stim_trigger_level_from_proto(pb: int) -> StimTriggerHighOrLow:
    """Convert proto StimTriggerHighOrLow int to domain StimTriggerLevel."""
    return _STIM_TRIGGER_HIGH_OR_LOW_FROM_PROTO[pb]


# --- StimPulseMode ---

_STIM_PULSE_MODE_TO_PROTO: dict[StimPulseMode, allegoserver_pb2.StimPulseOrTrain.ValueType] = {
    StimPulseMode.SINGLE_PULSE: allegoserver_pb2.SINGLE_PULSE,
    StimPulseMode.PULSE_TRAIN: allegoserver_pb2.PULSE_TRAIN,
}

_STIM_PULSE_MODE_FROM_PROTO: dict[int, StimPulseMode] = {
    int(allegoserver_pb2.SINGLE_PULSE): StimPulseMode.SINGLE_PULSE,
    int(allegoserver_pb2.PULSE_TRAIN): StimPulseMode.PULSE_TRAIN,
}


def stim_pulse_mode_to_proto(
    spm: StimPulseMode,
) -> allegoserver_pb2.StimPulseOrTrain.ValueType:
    """Convert domain StimPulseMode to proto StimPulseOrTrain."""
    return _STIM_PULSE_MODE_TO_PROTO[spm]


def stim_pulse_mode_from_proto(pb: int) -> StimPulseMode:
    """Convert proto StimPulseOrTrain int to domain StimPulseMode."""
    return _STIM_PULSE_MODE_FROM_PROTO[pb]


def stim_params_from_proto(
    reply: allegoserver_pb2.StimParamsReply,
) -> dict[int, StimParams]:
    """Convert ``StimParamsReply`` proto → dict keyed by system channel index."""
    result: dict[int, StimParams] = {}
    for sys_chan_idx, pb in reply.params.items():
        result[sys_chan_idx] = _single_from_proto(pb)
    return result


def _single_from_proto(pb: allegoserver_pb2.StimParams) -> StimParams:
    """Convert a single ``StimParams`` proto message to a domain model."""
    return StimParams(
        stim_shape=stim_shape_from_proto(pb.stimShape),
        stim_polarity=stim_polarity_from_proto(pb.stimPolarity),
        first_phase_duration=pb.firstPhaseDuration,
        second_phase_duration=pb.secondPhaseDuration,
        interphase_delay=pb.interphaseDelay,
        first_phase_amplitude=pb.firstPhaseAmplitude,
        second_phase_amplitude=pb.secondPhaseAmplitude,
        baseline_voltage=pb.baselineVoltage,
        trigger_edge_or_level=stim_trigger_edge_or_level_from_proto(pb.triggerEdgeOrLevel),
        trigger_high_or_low=stim_trigger_level_from_proto(pb.triggerHighOrLow),
        enabled=pb.enabled,
        post_trigger_delay=pb.postTriggerDelay,
        pulse_or_train=stim_pulse_mode_from_proto(pb.pulseOrTrain),
        number_of_stim_pulses=pb.numberOfStimPulses,
        pulse_train_period=pb.pulseTrainPeriod,
        refractory_period=pb.refractoryPeriod,
        pre_stim_amp_settle=pb.preStimAmpSettle,
        post_stim_amp_settle=pb.postStimAmpSettle,
        maintain_amp_settle=pb.maintainAmpSettle,
        enable_amp_settle=pb.enableAmpSettle,
        headstage_global_amp_settle=pb.headstageGlobalAmpSettle,
        post_stim_charge_recov_on=pb.postStimChargeRecovOn,
        post_stim_charge_recov_off=pb.postStimChargeRecovOff,
        enable_charge_recovery=pb.enableChargeRecovery,
        trigger_source_is_keypress=pb.triggerSourceIsKeypress,
        trigger_source_idx=pb.triggerSourceIdx,
        stim_sys_chan_idx=pb.stimSysChanIdx,
    )


def stim_params_to_proto(
    stim_params: StimParams,
) -> allegoserver_pb2.StimParams:
    """Convert ``StimParams`` domain model → proto message."""
    return allegoserver_pb2.StimParams(
        stimShape=stim_shape_to_proto(stim_params.stim_shape),
        stimPolarity=stim_polarity_to_proto(stim_params.stim_polarity),
        firstPhaseDuration=stim_params.first_phase_duration,
        secondPhaseDuration=stim_params.second_phase_duration,
        interphaseDelay=stim_params.interphase_delay,
        firstPhaseAmplitude=stim_params.first_phase_amplitude,
        secondPhaseAmplitude=stim_params.second_phase_amplitude,
        baselineVoltage=stim_params.baseline_voltage,
        triggerEdgeOrLevel=stim_trigger_edge_or_level_to_proto(stim_params.trigger_edge_or_level),
        triggerHighOrLow=stim_trigger_level_to_proto(stim_params.trigger_high_or_low),
        enabled=stim_params.enabled,
        postTriggerDelay=stim_params.post_trigger_delay,
        pulseOrTrain=stim_pulse_mode_to_proto(stim_params.pulse_or_train),
        numberOfStimPulses=stim_params.number_of_stim_pulses,
        pulseTrainPeriod=stim_params.pulse_train_period,
        refractoryPeriod=stim_params.refractory_period,
        preStimAmpSettle=stim_params.pre_stim_amp_settle,
        postStimAmpSettle=stim_params.post_stim_amp_settle,
        maintainAmpSettle=stim_params.maintain_amp_settle,
        enableAmpSettle=stim_params.enable_amp_settle,
        headstageGlobalAmpSettle=stim_params.headstage_global_amp_settle,
        postStimChargeRecovOn=stim_params.post_stim_charge_recov_on,
        postStimChargeRecovOff=stim_params.post_stim_charge_recov_off,
        enableChargeRecovery=stim_params.enable_charge_recovery,
        triggerSourceIsKeypress=stim_params.trigger_source_is_keypress,
        triggerSourceIdx=stim_params.trigger_source_idx,
        stimSysChanIdx=stim_params.stim_sys_chan_idx,
    )
