"""VidereClient — offline file-based analysis client."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, cast

from pydantic import ConfigDict, validate_call

from radiens_core._base_client import BaseClient
from radiens_core._converters._signals import sig_select_to_proto, signal_arrays_from_signals2
from radiens_core._converters._time_range import time_range_to_proto, trs_mode_to_proto
from radiens_core._converters._validation import validate_proto
from radiens_core._generated import common_pb2
from radiens_core._services._radiens_core import get_hd_snapshot_py_proto, set_datasource
from radiens_core._services._radiens_lifecycle import healthcheck_proto
from radiens_core._transport._channel_pool import ChannelPool
from radiens_core._transport._discovery import discover_radiens_addresses
from radiens_core.exceptions import RadiensError
from radiens_core.models.common import ClientType, ServiceEndpoint
from radiens_core.models.dataset import DatasetMetadata, FlexTimeRange
from radiens_core.models.signals import FlexSignalSpec, SignalArrays

if TYPE_CHECKING:
    from radiens_core.models.dataset import TimeRangeSpec
    from radiens_core.models.signals import SignalSpec
    from radiens_core.models.types import DatasetRef


class VidereClient(BaseClient):
    """Client for offline file-based analysis.

    Auto-discovers a running ``radiensserver`` on the local machine.

    Usage::

        with VidereClient() as client:
            meta = client.link_data_file("/path/to/recording.xdat")
            sigs = client.get_signals(meta, time_range=[0, 1], signals=[0, 1, 2])
    """

    def __init__(self) -> None:
        # 1. Discover addresses
        tmp_pool = ChannelPool()
        try:
            addresses = discover_radiens_addresses(tmp_pool)
        finally:
            tmp_pool.close()

        # 2. Base init
        core_addr = addresses.get(ServiceEndpoint.RADIENS_CORE, "")
        super().__init__(
            addresses=addresses,
            auth_server_address=core_addr,
            fallback_service=ServiceEndpoint.RADIENS_CORE,
        )

        self._datasets: dict[str, DatasetMetadata] = {}

    def _resolve_dataset(self, ref: DatasetRef) -> DatasetMetadata:
        """Resolve a dataset reference to DatasetMetadata."""
        if isinstance(ref, DatasetMetadata):
            return ref
        ref_str = str(ref)
        if ref_str in self._datasets:
            return self._datasets[ref_str]
        raise ValueError(f"Unknown dataset: {ref_str}. Link it first with link_data_file().")

    # --- Public API ---

    def healthcheck(self, service: ServiceEndpoint = ServiceEndpoint.RADIENS_LIFECYCLE) -> None:
        """Verify the radiensserver is responsive.

        Parameters
        ----------
        service
            Which gRPC service to probe (default ``RADIENS_LIFECYCLE``).
            ``RADIENS_CORE`` is also supported.

        Raises
        ------
        RadiensError
            If the health check fails or the server is unresponsive.
        """

        addr = self._server_address(service)
        response = healthcheck_proto(self._pool, addr, ClientType.VIDERE, service=service)
        if response.errMsg:
            raise RadiensError(f"Health check failed: {response.errMsg}")

    def link_data_file(
        self,
        path: str | Path,
    ) -> DatasetMetadata:
        """Link a recording file and return its metadata.

        Parameters
        ----------
        path
            Path to the recording file (e.g., .xdat, .rhd).

        Returns
        -------
        DatasetMetadata
            Metadata for the linked dataset.
        """

        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        meta = set_datasource(self._pool, addr, str(path))
        self._datasets[meta.id] = meta
        return meta

    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
    def get_signals(
        self,
        dataset: str | Path | DatasetMetadata,
        time_range: FlexTimeRange,
        signals: FlexSignalSpec,
    ) -> SignalArrays:
        """Retrieve signal data for a specific time range and channel selection.

        Parameters
        ----------
        dataset
            Dataset to query (metadata object or ID).
        time_range
            Time range specification (e.g., [start, end] in seconds).
        signals
            Signal selection (e.g., list of amplifier channel indices).

        Returns
        -------
        SignalArrays
            The requested signal data as numpy arrays.
        """

        meta = self._resolve_dataset(dataset)

        # Coerce flex types to specs (redundant due to @validate_call but satisfies mypy)
        tr_spec = cast("TimeRangeSpec", time_range)
        sig_spec = cast("SignalSpec", signals)

        # Resolve specs into concrete domain models
        concrete_tr = tr_spec.resolve(dataset_metadata=meta)
        if concrete_tr is None:
            # FROM_HEAD (lookback) mode is not meaningful for offline files
            raise ValueError(
                f"lookback() time range is not supported for offline files; "
                f"use TimeRangeSpec.subset(start, end) or TimeRangeSpec.full() instead. Got: {time_range}"
            )

        sig_select = sig_spec.to_sig_select(meta.channel_metadata)

        # Build proto request
        req = common_pb2.HDSnapshotRequest2(
            dsourceID=meta.id,
            tR=time_range_to_proto(concrete_tr),
            selMode=trs_mode_to_proto(tr_spec.mode),
            sigSel=sig_select_to_proto(sig_select),
        )
        validate_proto(req)

        # Execute RPC
        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        response = get_hd_snapshot_py_proto(self._pool, addr, req)

        return signal_arrays_from_signals2(response.sigs)
