"""KPI domain model <-> proto converters."""

from __future__ import annotations

from typing import cast

import numpy as np
import numpy.typing as npt

from radiens_core._generated import common_pb2
from radiens_core.models.kpi import KpiMetric, KpiMetricId, KpiMetricsResult, KpiMode, KpiStatus
from radiens_core.models.signals import SignalType
from radiens_core.models.time_range import TimeRange

# --- KpiMode ---

_KPI_MODE_TO_PROTO: dict[KpiMode, common_pb2.KpiMetricsMode.ValueType] = {
    KpiMode.BASE: common_pb2.Kpi_MetricsMode_Base,
    KpiMode.AVG: common_pb2.Kpi_MetricsMode_Avg,
    KpiMode.STREAM: common_pb2.Kpi_MetricsMode_Stream,
}


def kpi_mode_to_proto(mode: KpiMode) -> common_pb2.KpiMetricsMode.ValueType:
    """Convert domain KpiMode to proto KpiMetricsMode."""
    return _KPI_MODE_TO_PROTO[mode]


# --- KpiMetric ---


def kpi_metric_to_proto(metric: KpiMetric) -> common_pb2.KpiMetricsEnum.ValueType:
    """Convert domain KpiMetric to proto KpiMetricsEnum."""
    return cast("common_pb2.KpiMetricsEnum.ValueType", metric.value)


# --- KpiMetricId ---


def kpi_metric_id_to_proto(mid: KpiMetricId) -> common_pb2.KpiMetricID:
    """Convert domain KpiMetricId to proto KpiMetricID."""
    return common_pb2.KpiMetricID(
        mode=kpi_mode_to_proto(mid.mode),
        name=kpi_metric_to_proto(mid.name),
    )


def kpi_metric_id_from_proto(pb: common_pb2.KpiMetricID) -> KpiMetricId:
    """Convert proto KpiMetricID to domain KpiMetricId."""
    return KpiMetricId(mode=KpiMode(pb.mode), name=KpiMetric(pb.name))


# --- SignalType (internal helpers) ---

_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(sig_type: SignalType) -> common_pb2.SignalType.ValueType:
    return _SIGNAL_TYPE_TO_PROTO[sig_type]


def _signal_type_from_proto(pb: int) -> SignalType:
    return _SIGNAL_TYPE_FROM_PROTO[pb]


# --- DenseMatrix ---


def _dense_matrix_to_numpy(matrix: common_pb2.DenseMatrix) -> npt.NDArray[np.float64]:
    """Unpack a DenseMatrix proto to a float64 numpy array of shape (rows, cols)."""
    rows = matrix.rows
    cols = matrix.cols
    if rows == 0 or cols == 0:
        return np.empty((rows, cols), dtype=np.float64)
    return np.frombuffer(matrix.data, dtype=np.float64).reshape(rows, cols)


# --- TimeRange ---


def _time_range_from_kpi_proto(pb: common_pb2.TimeRange) -> TimeRange:
    """Build a domain TimeRange from a proto TimeRange as used in KPI responses."""
    sec_start = pb.timeRangeSec[0] if len(pb.timeRangeSec) >= 1 else 0.0
    sec_end = pb.timeRangeSec[1] if len(pb.timeRangeSec) >= 2 else sec_start
    ts_start = pb.timestamp[0] if len(pb.timestamp) >= 1 else 0
    ts_end = pb.timestamp[1] if len(pb.timestamp) >= 2 else ts_start
    fs = pb.fs if pb.HasField("fs") else 0.0
    walltime = pb.wallTimeStart if pb.HasField("wallTimeStart") else ""
    return TimeRange(
        sec=(sec_start, sec_end),
        timestamp=(ts_start, ts_end),
        fs=fs,
        dur_sec=sec_end - sec_start,
        walltime_str=walltime,
        n=ts_end - ts_start,
    )


# --- KpiMetricsResult ---


def kpi_metrics_result_from_proto(pb: common_pb2.SignalMetricsPacketReply) -> KpiMetricsResult:
    """Convert SignalMetricsPacketReply proto to domain KpiMetricsResult."""
    return KpiMetricsResult(
        channel_indices=tuple(pb.ntvIdxs),
        metrics=tuple(kpi_metric_id_from_proto(m) for m in pb.metrics),
        signal_type=_signal_type_from_proto(int(pb.sigType)),
        values=_dense_matrix_to_numpy(pb.metricVals),
        time_range=_time_range_from_kpi_proto(pb.tR),
    )


# --- KpiStatus ---


def kpi_status_from_proto(pb: common_pb2.KpiStatusReply) -> KpiStatus:
    """Convert KpiStatusReply proto to domain KpiStatus."""
    tr: TimeRange | None = None
    if pb.HasField("tR"):
        tr = _time_range_from_kpi_proto(pb.tR)
    num_sigs: int | None = pb.numSigs if pb.HasField("numSigs") else None
    return KpiStatus(
        stream_group_id=pb.streamGroupId,
        num_packets_memory=pb.numPacketsMemory,
        packet_dur=pb.packetDur,
        update_period=pb.updatePeriod,
        persistence=pb.persistence,
        beta=pb.beta,
        is_tracking_signal_cache=pb.isTrackingSignalCache,
        num_packets=pb.numPackets,
        memory_bytes=pb.memoryBytes,
        num_sigs=num_sigs,
        time_range=tr,
    )


# --- BundleReq builder ---


def build_bundle_req(
    channel_indices: list[int],
    time_range_sec: list[float],
    metrics: list[KpiMetricId],
) -> common_pb2.BundleReq:
    """Build a BundleReq proto from domain types.

    Parameters
    ----------
    channel_indices
        Native channel indices (ntvIdxs).
    time_range_sec
        Time range as ``[start_sec, end_sec]`` floats.
    metrics
        List of KPI metric identifiers.
    """
    return common_pb2.BundleReq(
        ntvIdxs=channel_indices,
        tR=common_pb2.TimeRange(timeRangeSec=time_range_sec),
        metrics=[kpi_metric_id_to_proto(m) for m in metrics],
    )
