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

from __future__ import annotations

from typing import TYPE_CHECKING

from radiens_core._converters._enums import signal_type_from_proto
from radiens_core.models.channels import (
    ChannelMetadata,
    KeyIdxs,
    SitePosition,
)
from radiens_core.models.enums import SignalType, SignalUnits

if TYPE_CHECKING:
    from radiens_core._generated import common_pb2


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:
            site_positions.append(SitePosition(
                probe_x=ch.siteCtrX,
                probe_y=ch.siteCtrY,
                probe_z=ch.siteCtrZ,
                tissue_x=ch.siteCtrTcsX,
                tissue_y=ch.siteCtrTcsY,
                tissue_z=ch.siteCtrTcsZ,
            ))

    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,
    )
