"""DatasetMetadata <-> proto converters."""

from __future__ import annotations

from radiens_core._converters._channels import channel_metadata_from_proto
from radiens_core._generated import common_pb2, datasource_pb2
from radiens_core.models.dataset import (
    BandParams,
    BulkSinkParams,
    BulkSourceParams,
    CARParams,
    DatasetMetadata,
    DownsampleParams,
    FileInfo,
    HighLowPassParams,
    NotchParams,
    PairedRefParams,
    ProtocolSpec,
    RadiensFileType,
    SinkParams,
    SliceChannelsParams,
    SliceTimeParams,
    SourceParams,
    TransformEdge,
    TransformNode,
    TransformNodeType,
    VirtualRefParams,
)
from radiens_core.models.time_range import TimeRange

# --- RadiensFileType enum conversion ---

_FILE_TYPE_FROM_PROTO: dict[int, RadiensFileType] = {
    int(common_pb2.RHD): RadiensFileType.RHD,
    int(common_pb2.XDAT): RadiensFileType.XDAT,
    int(common_pb2.CSV): RadiensFileType.CSV,
    int(common_pb2.HDF5): RadiensFileType.HDF5,
    int(common_pb2.NEX5): RadiensFileType.NEX5,
    int(common_pb2.NWB): RadiensFileType.NWB,
    int(common_pb2.BIOINTERFACE): RadiensFileType.BIOINTERFACE,
    int(common_pb2.KILOSORT2): RadiensFileType.KILOSORT2,
    int(common_pb2.ANC): RadiensFileType.ANC,
    int(common_pb2.SPKTRAIN_MU): RadiensFileType.SPKTRAIN_MU,
    int(common_pb2.BIO_RADIENS_EXPORT_HDF5): RadiensFileType.BIO_RADIENS_EXPORT_HDF5,
    int(common_pb2.BIO_RADIENS_EXPORT_MAT5): RadiensFileType.BIO_RADIENS_EXPORT_MAT5,
    int(common_pb2.NSX): RadiensFileType.NSX,
    int(common_pb2.PL2): RadiensFileType.PL2,
    int(common_pb2.TDT): RadiensFileType.TDT,
    int(common_pb2.SPIKES): RadiensFileType.SPIKES,
    int(common_pb2.KPI): RadiensFileType.KPI,
    int(common_pb2.UNKNOWN_FILE_TYPE): RadiensFileType.UNKNOWN_FILE_TYPE,
    int(common_pb2.MEAREC_HDF5): RadiensFileType.MEAREC_HDF5,
    int(common_pb2.OFFLINE_SORTER_SPKS_MAT5): RadiensFileType.OFFLINE_SORTER_SPKS_MAT5,
    int(common_pb2.SPKS_NPY): RadiensFileType.SPKS_NPY,
    int(common_pb2.OE_DAT): RadiensFileType.OE_DAT,
    int(common_pb2.SPKS_PHY2): RadiensFileType.SPKS_PHY2,
}

_FILE_TYPE_TO_PROTO: dict[RadiensFileType, common_pb2.RadixFileTypes.ValueType] = {
    RadiensFileType.RHD: common_pb2.RHD,
    RadiensFileType.XDAT: common_pb2.XDAT,
    RadiensFileType.CSV: common_pb2.CSV,
    RadiensFileType.HDF5: common_pb2.HDF5,
    RadiensFileType.NEX5: common_pb2.NEX5,
    RadiensFileType.NWB: common_pb2.NWB,
    RadiensFileType.BIOINTERFACE: common_pb2.BIOINTERFACE,
    RadiensFileType.KILOSORT2: common_pb2.KILOSORT2,
    RadiensFileType.ANC: common_pb2.ANC,
    RadiensFileType.SPKTRAIN_MU: common_pb2.SPKTRAIN_MU,
    RadiensFileType.BIO_RADIENS_EXPORT_HDF5: common_pb2.BIO_RADIENS_EXPORT_HDF5,
    RadiensFileType.BIO_RADIENS_EXPORT_MAT5: common_pb2.BIO_RADIENS_EXPORT_MAT5,
    RadiensFileType.NSX: common_pb2.NSX,
    RadiensFileType.PL2: common_pb2.PL2,
    RadiensFileType.TDT: common_pb2.TDT,
    RadiensFileType.SPIKES: common_pb2.SPIKES,
    RadiensFileType.KPI: common_pb2.KPI,
    RadiensFileType.UNKNOWN_FILE_TYPE: common_pb2.UNKNOWN_FILE_TYPE,
    RadiensFileType.MEAREC_HDF5: common_pb2.MEAREC_HDF5,
    RadiensFileType.OFFLINE_SORTER_SPKS_MAT5: common_pb2.OFFLINE_SORTER_SPKS_MAT5,
    RadiensFileType.SPKS_NPY: common_pb2.SPKS_NPY,
    RadiensFileType.OE_DAT: common_pb2.OE_DAT,
    RadiensFileType.SPKS_PHY2: common_pb2.SPKS_PHY2,
}


def file_type_from_proto(pb: int) -> RadiensFileType:
    """Convert proto RadixFileTypes int to domain RadiensFileType."""
    return _FILE_TYPE_FROM_PROTO[pb]


def file_type_to_proto(ft: RadiensFileType) -> common_pb2.RadixFileTypes.ValueType:
    """Convert domain RadiensFileType to proto RadixFileTypes."""
    return _FILE_TYPE_TO_PROTO[ft]


def guess_file_type(path: str) -> RadiensFileType:
    """Guess file type from file extension."""
    ext = path.rsplit(".", maxsplit=1)[-1].lower() if "." in path else ""
    ext_map: dict[str, RadiensFileType] = {
        "xdat": RadiensFileType.XDAT,
        "rhd": RadiensFileType.RHD,
        "csv": RadiensFileType.CSV,
        "hdf5": RadiensFileType.HDF5,
        "h5": RadiensFileType.HDF5,
        "nex5": RadiensFileType.NEX5,
        "nwb": RadiensFileType.NWB,
        "bin": RadiensFileType.KILOSORT2,
        "nsx": RadiensFileType.NSX,
        "tdt": RadiensFileType.TDT,
        "oebin": RadiensFileType.OE_DAT,
    }
    return ext_map.get(ext, RadiensFileType.XDAT)


def dataset_metadata_from_proto(
    raw: datasource_pb2.DataSourceSetSaveReply,
) -> DatasetMetadata:
    """Build domain DatasetMetadata from proto DataSourceSetSaveReply."""
    return dataset_metadata_from_status(raw.status, raw.dsourceID, raw.status.signalGroup.datasetUID)


def dataset_metadata_from_status(
    status: common_pb2.DataSourceStatus, dsource_id: str, probe_uid: str = ""
) -> DatasetMetadata:
    """Build domain DatasetMetadata from proto DataSourceStatus."""
    sig_group = status.signalGroup
    pb_tr = status.tR

    channel_metadata = channel_metadata_from_proto(sig_group)

    # Build TimeRange from the proto TimeRange message
    sec_list = list(pb_tr.timeRangeSec)
    ts_list = list(pb_tr.timestamp)
    sec_start = sec_list[0] if len(sec_list) > 0 else 0.0
    sec_end = sec_list[1] if len(sec_list) > 1 else 0.0
    ts_start = ts_list[0] if len(ts_list) > 0 else 0
    ts_end = ts_list[1] if len(ts_list) > 1 else 0

    time_range = TimeRange(
        sec=(sec_start, sec_end),
        timestamp=(ts_start, ts_end),
        fs=pb_tr.fs,
        dur_sec=sec_end - sec_start,
        walltime_str=pb_tr.wallTimeStart,
        n=ts_end - ts_start,
    )

    return DatasetMetadata(
        id=dsource_id,
        base_name=status.baseFileName,
        path=status.path,
        file_type=file_type_from_proto(status.fileType),
        time_range=time_range,
        channel_metadata=channel_metadata,
        parent_dsource_id="",
        probe_uid=probe_uid or sig_group.datasetUID,
    )


def dataset_metadata_map_from_proto(
    proto: datasource_pb2.DataSourceStatusMap,
) -> dict[str, DatasetMetadata]:
    """Convert DataSourceStatusMap proto to domain map."""
    return {dsid: dataset_metadata_from_status(status, dsid) for dsid, status in proto.dsource.items()}


def protocol_spec_from_proto(proto: common_pb2.Protocol) -> ProtocolSpec:
    """Convert proto Protocol to domain ProtocolSpec."""
    return ProtocolSpec(id=proto.id, node_count=len(proto.transformNodes))


def protocol_spec_list_from_proto(reply: common_pb2.GetAllProtocolsReply) -> list[ProtocolSpec]:
    """Convert GetAllProtocolsReply proto to list of domain ProtocolSpec."""
    return [protocol_spec_from_proto(p) for p in reply.protocols]


def file_info_list_from_proto(reply: datasource_pb2.CpRmMvLsReply) -> list[FileInfo]:
    """Convert CpRmMvLsReply.dsource list to list[FileInfo]."""
    result: list[FileInfo] = []
    for info in reply.dsource:
        result.append(
            FileInfo(
                path=info.desc.path,
                base_name=info.desc.baseName,
                file_type=file_type_from_proto(info.desc.fileType),
                dataset_uid=info.stat.datasetUID,
                num_channels=info.stat.numChannels,
                duration_sec=info.stat.durationSec,
                sample_rate=info.stat.sampleRate,
                num_bytes=info.stat.numBytes,
            )
        )
    return result


# ===================================================================
# TransformNode / TransformEdge / Protocol converters
# ===================================================================

_TRANSFORM_NODE_TYPE_TO_PROTO: dict[TransformNodeType, common_pb2.TransformNodeType.ValueType] = {
    TransformNodeType.NOTCH_FILTER: common_pb2.NOTCH_FILTER,
    TransformNodeType.BANDPASS_FILTER: common_pb2.BANDPASS_FILTER,
    TransformNodeType.BANDSTOP_FILTER: common_pb2.BANDSTOP_FILTER,
    TransformNodeType.HIGHPASS_FILTER: common_pb2.HIGHPASS_FILTER,
    TransformNodeType.LOWPASS_FILTER: common_pb2.LOWPASS_FILTER,
    TransformNodeType.PAIRED_REF: common_pb2.PAIRED_REF,
    TransformNodeType.VIRTUAL_REF: common_pb2.VIRTUAL_REF,
    TransformNodeType.CAR_REF: common_pb2.CAR_REF,
    TransformNodeType.SLICE_TIME: common_pb2.SLICE_TIME,
    TransformNodeType.SLICE_CHANNELS: common_pb2.SLICE_CHANNELS,
    TransformNodeType.SOURCE: common_pb2.SOURCE,
    TransformNodeType.SINK: common_pb2.SINK,
    TransformNodeType.DOWNSAMPLE: common_pb2.DOWNSAMPLE,
    TransformNodeType.REMAP_SENSOR: common_pb2.REMAP_SENSOR,
    TransformNodeType.BULK_SOURCE: common_pb2.BULK_SOURCE,
    TransformNodeType.BULK_SINK: common_pb2.BULK_SINK,
}


def transform_node_to_proto(node: TransformNode) -> common_pb2.TransformNode:
    """Convert domain TransformNode to proto TransformNode."""
    node_type = _TRANSFORM_NODE_TYPE_TO_PROTO[node.type]
    params = node.params

    if isinstance(params, HighLowPassParams):
        hlp_pb = common_pb2.HighLowPassTransformParams(frequency=params.frequency)
        if params.order is not None:
            hlp_pb.order = params.order
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            highLowPassParams=hlp_pb,
        )

    if isinstance(params, BandParams):
        band_pb = common_pb2.BandTransformParams(
            lowFrequency=params.low_frequency,
            highFrequency=params.high_frequency,
        )
        if params.order is not None:
            band_pb.order = params.order
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            bandParams=band_pb,
        )

    if isinstance(params, NotchParams):
        notch_pb = common_pb2.NotchTransformParams(
            notchFrequency=params.notch_frequency,
            notchBandwidth=params.bandwidth,
        )
        if params.order is not None:
            notch_pb.order = params.order
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            notchParams=notch_pb,
        )

    if isinstance(params, CARParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            carParams=common_pb2.CARTransformParams(),
        )

    if isinstance(params, VirtualRefParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            virtualRefParams=common_pb2.VirtualRefTransformParams(
                refNtvChanIdx=params.ref_channel,
            ),
        )

    if isinstance(params, PairedRefParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            pairedRefParams=common_pb2.PairedRefTransformParams(
                refNtvChanIdx=params.ref_channel,
                targetNtvChanIdx=params.target_channel,
            ),
        )

    if isinstance(params, SourceParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            datasourceSinkParams=common_pb2.DataSourceSinkTransformParams(
                id=params.dataset_id,
                dsrcName=params.base_name,
                path=params.path,
                fileType=file_type_to_proto(params.file_type),
            ),
        )

    if isinstance(params, SinkParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            datasourceSinkParams=common_pb2.DataSourceSinkTransformParams(
                dsrcName=params.base_name,
                path=params.path,
                fileType=file_type_to_proto(params.file_type),
            ),
        )

    if isinstance(params, SliceTimeParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            sliceTimeParams=common_pb2.SliceTimeTransformParams(
                timeStart=params.start_sec,
                timeEnd=params.end_sec,
            ),
        )

    if isinstance(params, SliceChannelsParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            sliceChannelsParams=common_pb2.SliceChannelsTransformParams(
                sysChanIdxs=params.channels,
            ),
        )

    if isinstance(params, DownsampleParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            downsampleParams=common_pb2.DownsampleTransformParams(
                sampleFactor=params.factor,
            ),
        )

    if isinstance(params, BulkSourceParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            bulkDataSourceParams=common_pb2.BulkDataSourceTransformParams(
                srcParams=[
                    common_pb2.DataSourceSinkTransformParams(
                        id=s.dataset_id,
                        dsrcName=s.base_name,
                        path=s.path,
                        fileType=file_type_to_proto(s.file_type),
                    )
                    for s in params.sources
                ]
            ),
        )

    if isinstance(params, BulkSinkParams):
        return common_pb2.TransformNode(
            id=node.id,
            type=node_type,
            invalid=False,
            bulkDataSinkParams=common_pb2.BulkDataSinkTransformParams(
                suffix=params.suffix,
                path=params.path,
                fileType=file_type_to_proto(params.file_type),
            ),
        )

    # params is None (e.g. REMAP_SENSOR or bare node)
    return common_pb2.TransformNode(id=node.id, type=node_type, invalid=False)


def transform_edge_to_proto(edge: TransformEdge) -> common_pb2.TransformEdge:
    """Convert domain TransformEdge to proto TransformEdge."""
    return common_pb2.TransformEdge(id=edge.id, source=edge.source, target=edge.target)


def protocol_to_proto(
    protocol_id: str,
    nodes: list[TransformNode],
    edges: list[TransformEdge],
) -> common_pb2.Protocol:
    """Build a proto Protocol from domain nodes, edges, and a protocol ID."""
    return common_pb2.Protocol(
        id=protocol_id,
        transformNodes=[transform_node_to_proto(n) for n in nodes],
        transformEdges=[transform_edge_to_proto(e) for e in edges],
    )
