"""CurateClient — data curation client."""

from __future__ import annotations

import os
import uuid
from typing import TYPE_CHECKING

from tqdm import tqdm

from radiens_core._base_client import BaseClient
from radiens_core._converters._dataset import file_type_to_proto, guess_file_type
from radiens_core._generated import common_pb2
from radiens_core._services._radiens_core import apply_protocol_proto, set_datasource, set_protocol_proto
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

if TYPE_CHECKING:
    from pathlib import Path

    from radiens_core.models.types import DatasetRef


class CurateClient(BaseClient):
    """Client for data curation operations.

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

    Usage::

        with CurateClient() as client:
            meta = client.link_data_file("/path/to/recording.xdat")
            new_meta = client.slice_time(meta, 0, 10, "sliced.xdat")
    """

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

        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] = {}

    # --- Internal helpers ---

    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().")

    def _build_linear_protocol(
        self,
        input_dataset: DatasetMetadata,
        transform_node: common_pb2.TransformNode,
        output_path: str | Path,
    ) -> tuple[list[common_pb2.TransformNode], list[common_pb2.TransformEdge]]:
        """Build a standard 3-node linear protocol: Source -> Transform -> Sink."""
        output_path_str = str(output_path)
        base_name = os.path.splitext(os.path.basename(output_path_str))[0]
        dir_path = os.path.dirname(output_path_str) or "."

        source_id = str(uuid.uuid4())
        sink_id = str(uuid.uuid4())

        # NOTE: 'invalid' is a 'required' field in the proto definition.
        # We must set it to False to pass client-side serialization.
        source_node = common_pb2.TransformNode(
            id=source_id,
            type=common_pb2.SOURCE,
            invalid=False,
            datasourceSinkParams=common_pb2.DataSourceSinkTransformParams(
                id=input_dataset.id,
                dsrcName=input_dataset.base_name,
                path=input_dataset.path,
                fileType=file_type_to_proto(input_dataset.file_type),
            ),
        )

        sink_node = common_pb2.TransformNode(
            id=sink_id,
            type=common_pb2.SINK,
            invalid=False,
            datasourceSinkParams=common_pb2.DataSourceSinkTransformParams(
                dsrcName=base_name,
                path=dir_path,
                fileType=file_type_to_proto(guess_file_type(output_path_str)),
            ),
        )

        transform_node.invalid = False

        edges = [
            common_pb2.TransformEdge(id=str(uuid.uuid4()), source=source_id, target=transform_node.id),
            common_pb2.TransformEdge(id=str(uuid.uuid4()), source=transform_node.id, target=sink_id),
        ]

        return [source_node, transform_node, sink_node], edges

    def _execute_protocol(
        self,
        nodes: list[common_pb2.TransformNode],
        edges: list[common_pb2.TransformEdge],
        desc: str = "Processing",
    ) -> None:
        """Submit and execute a protocol with a progress bar."""
        protocol_id = str(uuid.uuid4())
        protocol = common_pb2.Protocol(
            id=protocol_id,
            transformNodes=nodes,
            transformEdges=edges,
        )

        addr = self._server_address(ServiceEndpoint.RADIENS_CORE)
        validated_protocol = set_protocol_proto(self._pool, addr, protocol)

        # Check for validation errors from the server
        for node in validated_protocol.transformNodes:
            if node.invalid:
                raise RadiensError(f"Curation task failed validation at node '{node.id}': {node.errorMessage}")

        req = common_pb2.ProtocolRequest(protocolID=protocol_id)
        stream = apply_protocol_proto(self._pool, addr, req)

        with tqdm(total=100, desc=desc, unit="%") as pbar:
            last_frac = 0.0
            for progress in stream:
                # progress.fracComplete is 0.0 to 1.0
                inc = (progress.fracComplete - last_frac) * 100
                if inc > 0:
                    pbar.update(inc)
                    last_frac = progress.fracComplete
            # Ensure it reaches 100% if the stream finished successfully
            if last_frac < 1.0:
                pbar.update((1.0 - last_frac) * 100)

    # --- 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.CURATE, 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

    def slice_time(
        self,
        dataset: DatasetRef,
        start_sec: float,
        end_sec: float,
        output_path: str | Path,
    ) -> DatasetMetadata:
        """Slice a dataset in time.

        Parameters
        ----------
        dataset
            Dataset to slice (metadata or ID).
        start_sec
            Start time in seconds.
        end_sec
            End time in seconds.
        output_path
            Where to save the result.

        Returns
        -------
        DatasetMetadata
            Metadata for the new (sliced) dataset.
        """
        meta = self._resolve_dataset(dataset)

        transform_id = str(uuid.uuid4())
        transform_node = common_pb2.TransformNode(
            id=transform_id,
            type=common_pb2.SLICE_TIME,
            invalid=False,
            sliceTimeParams=common_pb2.SliceTimeTransformParams(
                timeStart=start_sec,
                timeEnd=end_sec,
            ),
        )

        nodes, edges = self._build_linear_protocol(meta, transform_node, output_path)
        self._execute_protocol(nodes, edges, desc=f"Slicing {os.path.basename(str(output_path))}")

        return self.link_data_file(output_path)

    def downsample(
        self,
        dataset: DatasetRef,
        factor: int,
        output_path: str | Path,
    ) -> DatasetMetadata:
        """Downsample a dataset.

        Parameters
        ----------
        dataset
            Dataset to downsample (metadata or ID).
        factor
            Downsampling factor (e.g., 2, 4, 10).
        output_path
            Where to save the result.

        Returns
        -------
        DatasetMetadata
            Metadata for the new (downsampled) dataset.
        """
        meta = self._resolve_dataset(dataset)

        transform_id = str(uuid.uuid4())
        transform_node = common_pb2.TransformNode(
            id=transform_id,
            type=common_pb2.DOWNSAMPLE,
            invalid=False,
            downsampleParams=common_pb2.DownsampleTransformParams(
                sampleFactor=factor,
            ),
        )

        nodes, edges = self._build_linear_protocol(meta, transform_node, output_path)
        self._execute_protocol(nodes, edges, desc=f"Downsampling {os.path.basename(str(output_path))}")

        return self.link_data_file(output_path)
