"""RadiensCore RPC wrappers — thin service layer (proto only)."""

from __future__ import annotations

import os
from typing import TYPE_CHECKING

import grpc

from radiens_core._converters._dataset import dataset_metadata_from_proto, file_type_to_proto, guess_file_type
from radiens_core._generated import (
    biointerface_pb2,
    common_pb2,
    datasource_pb2,
    radiensserver_pb2_grpc,
    spikesorter_pb2,
)
from radiens_core._transport._errors import handle_grpc_error
from radiens_core.models.common import ClientType
from radiens_core.models.dataset import RadiensFileType

if TYPE_CHECKING:
    from collections.abc import Iterable

    from radiens_core._transport._channel_pool import ChannelPool
    from radiens_core.models.dataset import DatasetMetadata


def get_hd_snapshot_py_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.HDSnapshotRequest2,
) -> common_pb2.HDSnapshot2:
    """Send ``GetHDSnapshotPy`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.GetHDSnapshotPy(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def set_datasource(
    pool: ChannelPool,
    addr: str,
    path: str,
) -> DatasetMetadata:
    """Link a data file and register any associated spikes files.

    Calls ``SetDataSourceFromFile`` with ``findAssociated=True``, then calls
    ``SpikesSetFromFile`` for each discovered ``.spikes`` file.  Returns a
    domain ``DatasetMetadata`` with ``associated_spikes_ids`` populated.
    """
    ft = guess_file_type(path)
    base_name = os.path.splitext(os.path.basename(path))[0]

    # Normalize XDAT base name by stripping standard suffixes if present
    if ft == RadiensFileType.XDAT:
        if base_name.endswith("_data"):
            base_name = base_name[:-5]
        elif base_name.endswith("_timestamp"):
            base_name = base_name[:-10]

    dir_path = os.path.dirname(path) or "."

    req = datasource_pb2.DataSourceSetSaveRequest(
        path=dir_path,
        baseName=base_name,
        fileType=file_type_to_proto(ft),
        findAssociated=True,
    )

    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        res = stub.SetDataSourceFromFile(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy

    meta = dataset_metadata_from_proto(res)

    # Register any associated spikes files discovered by the server.
    # When the datasource is already loaded the server returns early and
    # associatedDsrcs is empty — in that case fall back to SpikesGetIDs.
    spikes_ids: list[str] = []
    for assoc in res.associatedDsrcs:
        if assoc.fileType == common_pb2.SPIKES:
            spikes_reply = _spikes_set_from_file(pool, addr, assoc.path, assoc.baseName, meta.id)
            spikes_ids.append(spikes_reply.dsourceID)

    if not spikes_ids:
        prefix = meta.id + "_"
        spikes_ids = [sid for sid in _spikes_get_ids(pool, addr) if sid.startswith(prefix)]

    if spikes_ids:
        return meta.model_copy(update={"associated_spikes_ids": spikes_ids})
    return meta


def _spikes_get_ids(pool: ChannelPool, addr: str) -> list[str]:
    """Return all currently registered spikes datasource IDs (internal helper)."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        reply = stub.SpikesGetIDs(datasource_pb2.DataSourceRequest())
        return list(reply.sortedID)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def _spikes_set_from_file(
    pool: ChannelPool,
    addr: str,
    path: str,
    base_name: str,
    parent_dsource_id: str,
) -> datasource_pb2.DataSourceSetSaveReply:
    """Register a .spikes file via SetDataSourceFromFile RPC (internal helper).

    Uses the same RPC as the TypeScript frontend — SpikesSetFromFile is unimplemented
    on the server side.  The SPIKES file type and parentDsourceID are what cause the
    server to treat this as a biointerface registration.
    """
    req = datasource_pb2.DataSourceSetSaveRequest(
        path=path,
        baseName=base_name,
        fileType=common_pb2.SPIKES,
        parentDsourceID=parent_dsource_id,
    )
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.SetDataSourceFromFile(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def set_protocol_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.Protocol,
) -> common_pb2.Protocol:
    """Send ``SetProtocol`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.SetProtocol(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.CURATE)
        raise  # unreachable, but satisfies mypy


def apply_protocol_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.ProtocolRequest,
) -> Iterable[common_pb2.ApplyProtocolProgress]:
    """Send ``ApplyProtocol`` RPC, return stream of ApplyProtocolProgress."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.ApplyProtocol(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.CURATE)
        raise  # unreachable, but satisfies mypy


def get_protocol_proto(
    pool: ChannelPool,
    addr: str,
    protocol_id: str,
) -> common_pb2.Protocol:
    """Send ``GetProtocol`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.GetProtocol(common_pb2.ProtocolRequest(protocolID=protocol_id))
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.CURATE)
        raise  # unreachable, but satisfies mypy


def get_all_protocols_proto(
    pool: ChannelPool,
    addr: str,
) -> common_pb2.GetAllProtocolsReply:
    """Send ``GetAllProtocols`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.GetAllProtocols(common_pb2.StandardRequest())
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.CURATE)
        raise  # unreachable, but satisfies mypy


def get_dsp_group_proto(
    pool: ChannelPool,
    addr: str,
    stream_group_id: str,
) -> common_pb2.DSPGroup:
    """Send ``GetDSPGroup`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    req = common_pb2.SignalGroupIDRequest(streamGroupId=stream_group_id)
    try:
        return stub.GetDSPGroup(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def set_dsp_group_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.SetDSPGroupRequest,
) -> common_pb2.StandardReply:
    """Send ``SetDSPGroup`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.SetDSPGroup(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def list_data_sources_proto(
    pool: ChannelPool,
    addr: str,
    dsource_ids: list[str] | None = None,
) -> datasource_pb2.DataSourceStatusMap:
    """Send ``ListDataSource`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    req = datasource_pb2.DataSourceRequest(dsourceID=dsource_ids if dsource_ids else [])
    try:
        return stub.ListDataSource(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def list_data_source_ids_proto(
    pool: ChannelPool,
    addr: str,
) -> datasource_pb2.DataSourceIDReply:
    """Send ``ListDataSourceIDs`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.ListDataSourceIDs(datasource_pb2.DataSourceRequest())
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def clear_data_source_proto(
    pool: ChannelPool,
    addr: str,
    dsource_ids: list[str],
) -> datasource_pb2.DataSourceIDReply:
    """Send ``ClearDataSource`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    req = datasource_pb2.DataSourceRequest(dsourceID=dsource_ids)
    try:
        return stub.ClearDataSource(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def export_data_source_proto(
    pool: ChannelPool,
    addr: str,
    request: datasource_pb2.ExportDataSourceFileRequest,
) -> common_pb2.StandardReply:
    """Send ``ExportDataSource`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.ExportDataSource(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def list_directory_proto(
    pool: ChannelPool,
    addr: str,
    directory: str,
) -> datasource_pb2.CpRmMvLsReply:
    """Send ``ListDirectory`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    req = common_pb2.ListDataSourcesRequest(directory=directory)
    try:
        return stub.ListDirectory(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def copy_data_source_file_proto(
    pool: ChannelPool,
    addr: str,
    req: datasource_pb2.CopyRemoveDataSourceFileRequest,
) -> datasource_pb2.CpRmMvLsReply:
    """Send ``CopyDataSourceFile`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.CopyDataSourceFile(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def move_data_source_file_proto(
    pool: ChannelPool,
    addr: str,
    req: datasource_pb2.MoveDataSourceFileRequest,
) -> datasource_pb2.CpRmMvLsReply:
    """Send ``MoveDataSourceFile`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.MoveDataSourceFile(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


def remove_data_source_file_proto(
    pool: ChannelPool,
    addr: str,
    req: datasource_pb2.CopyRemoveDataSourceFileRequest,
) -> datasource_pb2.CpRmMvLsReply:
    """Send ``RemoveDataSourceFile`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.RemoveDataSourceFile(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise


# --- KPI ---


def kpi_get_metrics_radiens_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.KpiMetricsReq,
) -> common_pb2.SignalMetricsPacketReply:
    """Send ``KpiGetMetrics`` RPC to RadiensCore, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.KpiGetMetrics(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def get_kpi_status_radiens_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.GetKpiStatusRequest,
) -> common_pb2.KpiStatusReply:
    """Send ``GetKpiStatus`` RPC to RadiensCore, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.GetKpiStatus(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def kpi_calculate_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.KpiStandardRequest,
) -> common_pb2.StandardReply:
    """Send ``KpiCalculate`` RPC to RadiensCore, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.KpiCalculate(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def kpi_clear_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.KpiStandardRequest,
) -> common_pb2.StandardReply:
    """Send ``KpiClear`` RPC to RadiensCore, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.KpiClear(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def set_kpi_packet_dur_radiens_proto(
    pool: ChannelPool,
    addr: str,
    request: common_pb2.SetKpiParamRequest,
) -> common_pb2.StandardReply:
    """Send ``SetKpiPacketDur`` RPC to RadiensCore, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.SetKpiPacketDur(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


# --- Spikes ---


def spikes_get_spec_proto(
    pool: ChannelPool,
    addr: str,
    dsource_id: str,
) -> biointerface_pb2.SpikesSpecReply:
    """Send ``SpikesGetSpec`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    req = datasource_pb2.DataSourceRequest(dsourceID=[dsource_id])
    try:
        return stub.SpikesGetSpec(req)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def spikes_get_raster_data_proto(
    pool: ChannelPool,
    addr: str,
    request: spikesorter_pb2.SpikeSorterGetRasterDataRequest,
) -> spikesorter_pb2.SpikeSorterRasterDataReply:
    """Send ``SpikesGetRasterData`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.SpikesGetRasterData(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def spikes_get_dense_proto(
    pool: ChannelPool,
    addr: str,
    request: biointerface_pb2.SpikesGetSpikesRequest,
) -> biointerface_pb2.SpikesSpikeDataDenseReply:
    """Send ``SpikesGetSpikesDense`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.SpikesGetSpikesDense(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy


def biointerface_get_neurons_proto(
    pool: ChannelPool,
    addr: str,
    request: biointerface_pb2.BiointerfaceGetNeuronsRequest,
) -> biointerface_pb2.BiointerfaceGetNeuronsReply:
    """Send ``BiointerfaceGetNeurons`` RPC, return raw proto response."""
    channel = pool.get_channel(addr)
    stub = radiensserver_pb2_grpc.RadiensCoreStub(channel)
    try:
        return stub.BiointerfaceGetNeurons(request)
    except grpc.RpcError as ex:
        handle_grpc_error(ex, ClientType.VIDERE)
        raise  # unreachable, but satisfies mypy
