"""ChannelMetadata <-> proto converters."""

from __future__ import annotations

from radiens_core._generated import common_pb2
from radiens_core.models.channels import (
    ChannelMetadata,
    KeyIdxs,
    SignalUnits,
    SitePosition,
)
from radiens_core.models.signals import SignalType

# --- SignalType enum conversion ---

_SIGNAL_TYPE_TO_PROTO: dict[SignalType, common_pb2.SignalType.ValueType] = {
    SignalType.AMP: common_pb2.PRI,
    SignalType.AIN: common_pb2.AUX,
    SignalType.DIN: common_pb2.DIN,
    SignalType.DOUT: common_pb2.DOUT,
}

_SIGNAL_TYPE_FROM_PROTO: dict[int, SignalType] = {
    int(common_pb2.PRI): SignalType.AMP,
    int(common_pb2.AUX): SignalType.AIN,
    int(common_pb2.DIN): SignalType.DIN,
    int(common_pb2.DOUT): SignalType.DOUT,
}


def signal_type_to_proto(st: SignalType) -> common_pb2.SignalType.ValueType:
    """Convert domain SignalType to proto SignalType."""
    return _SIGNAL_TYPE_TO_PROTO[st]


def signal_type_from_proto(pb: int) -> SignalType:
    """Convert proto SignalType int to domain SignalType."""
    return _SIGNAL_TYPE_FROM_PROTO[pb]


def channel_metadata_from_proto(raw: common_pb2.SignalGroup) -> ChannelMetadata:
    """Build domain ChannelMetadata from proto SignalGroup."""
    # Parse signal units
    signal_units: dict[SignalType, SignalUnits] = {}
    for sig_type_pb, sig_unit_pb in raw.signalUnits.units.items():
        try:
            st = signal_type_from_proto(sig_type_pb)
            signal_units[st] = SignalUnits(sig_unit_pb)
        except (KeyError, ValueError):
            pass

    # Parse channels from repeated ChannelRecord
    sig_map: dict[SignalType, KeyIdxs] = {st: KeyIdxs(ntv=[], dset=[], sys=[]) for st in SignalType}
    selected_sig_map: dict[SignalType, KeyIdxs] = {st: KeyIdxs(ntv=[], dset=[], sys=[]) for st in SignalType}
    site_positions: list[SitePosition] = []

    for ch in raw.channels:
        try:
            st = signal_type_from_proto(ch.chanType)
        except KeyError:
            continue
        sig_map[st].ntv.append(ch.ntvChanIdx)
        sig_map[st].dset.append(ch.absIdx)
        sig_map[st].sys.append(ch.sysChanIdx)

        if ch.isSelected:
            selected_sig_map[st].ntv.append(ch.ntvChanIdx)
            selected_sig_map[st].dset.append(ch.absIdx)
            selected_sig_map[st].sys.append(ch.sysChanIdx)

        if st == SignalType.AMP:
            # positions are in 10s of microns in the proto, convert to microns in the domain model
            site_positions.append(
                SitePosition(
                    probe_x=ch.siteCtrX * 10,
                    probe_y=ch.siteCtrY * 10,
                    probe_z=ch.siteCtrZ * 10,
                    tissue_x=ch.siteCtrTcsX * 10,
                    tissue_y=ch.siteCtrTcsY * 10,
                    tissue_z=ch.siteCtrTcsZ * 10,
                    site_shape=ch.siteShape,
                    site_area_um2=ch.siteAreaMicron2,
                    site_lim_x_min=ch.siteLimXMin * 10,
                    site_lim_x_max=ch.siteLimXMax * 10,
                    site_lim_y_min=ch.siteLimYMin * 10,
                    site_lim_y_max=ch.siteLimYMax * 10,
                    site_lim_z_min=ch.siteLimZMin * 10,
                    site_lim_z_max=ch.siteLimZMax * 10,
                    scs_to_tcs_x=ch.scsToTcsX * 10,
                    scs_to_tcs_y=ch.scsToTcsY * 10,
                    scs_to_tcs_z=ch.scsToTcsZ * 10,
                )
            )

    num_channels = {st: len(sig_map[st].ntv) for st in SignalType}
    selected_num_channels = {st: len(selected_sig_map[st].ntv) for st in SignalType}

    return ChannelMetadata(
        dataset_uid=raw.datasetUID,
        fs=raw.sampFreq,
        source_label=raw.sourceLabel,
        neighbor_max_radius_um=raw.neighborMaxRadius,
        signal_units=signal_units,
        has_discrete=raw.hasDiscrete,
        has_continuous_amp=raw.hasContinuousAmp,
        sig_map=sig_map,
        selected_sig_map=selected_sig_map,
        num_channels=num_channels,
        selected_num_channels=selected_num_channels,
        site_positions=site_positions,
    )
