"""RTL-SDR device integration for the RIA Toolkit."""

import time
import warnings
from typing import Optional

import numpy as np

try:
    from rtlsdr import RtlSdr
except ImportError as exc:  # pragma: no cover - dependency provided by end user
    raise ImportError("pyrtlsdr is required to use the RTLSDR class") from exc

from ria_toolkit_oss.datatypes.recording import Recording
from ria_toolkit_oss.sdr.sdr import SDR


class RTLSDR(SDR):
    """SDR interface for RTL-SDR dongles using pyrtlsdr."""

    def __init__(self, identifier: Optional[str] = None):
        """
        Initialize a Pluto SDR device object and connect to the SDR hardware.

        This software supports the ADALM Pluto SDR created by Analog Devices.

        :param identifier: The value of the parameter that identifies the device.
        :type identifier: str = "192.168.3.1", "pluto.local", etc

        If no identifier is provided, it will select the first device found, with a warning.
        If more than one device is found with the identifier, it will select the first of those devices.
        """
        print(f"Initializing Pluto radio with identifier [{identifier}].")
        try:
            super().__init__()

            if identifier is None:
                self.radio = RtlSdr()
            else:
                self.radio = RtlSdr(identifier)

            self.rx_buffer_size = 256_000
            self.rx_channel = 0

            print(f"Initialized RTL-SDR with identifier [{identifier}].")

        except Exception as e:
            print(f"Failed to find RTL-SDR with identifier [{identifier}].")
            raise e

    def init_rx(
        self,
        sample_rate: int | float,
        center_frequency: int | float,
        gain: Optional[int],
        channel: int,
        gain_mode: Optional[str] = "absolute",
        buffer_size: Optional[int] = 256_000,
        bias_t: bool = False,
    ):
        if channel not in (0, None):
            raise ValueError("RTL-SDR supports only channel 0 for RX.")

        self.set_rx_sample_rate(sample_rate=sample_rate)
        self.set_rx_center_frequency(center_frequency=center_frequency)
        self.set_rx_gain(gain=gain, gain_mode=gain_mode)

        self.rx_buffer_size = int(buffer_size or self.rx_buffer_size)
        self.rx_channel = 0

        if bias_t:
            self.set_bias_tee(True)
            time.sleep(1)

        self._rx_initialized = True
        self._tx_initialized = False

        return {"sample_rate": self.rx_sample_rate, "center_frequency": self.rx_center_frequency, "gain": self.rx_gain}

    def set_rx_sample_rate(self, sample_rate):
        self.radio.sample_rate = float(sample_rate)
        self.rx_sample_rate = self.radio.sample_rate
        print(f"RTL RX Sample Rate = {self.radio.get_sample_rate()}")

    def set_rx_center_frequency(self, center_frequency):
        self.radio.center_freq = float(center_frequency)
        self.rx_center_frequency = self.radio.center_freq
        print(f"RTL RX Center Frequency = {self.radio.get_center_freq()}")

    def set_rx_gain(self, gain, gain_mode="absolute"):
        available_gains = self.radio.get_gains()

        if gain is None:
            self.radio.gain = "auto"
            self.rx_gain = "auto"
        else:
            if not available_gains:
                warnings.warn(
                    "No gain table reported by RTL-SDR; applying requested gain directly.",
                    RuntimeWarning,
                )
                target_gain = gain
            else:
                min_gain = min(available_gains)
                max_gain = max(available_gains)

                if gain_mode == "relative":
                    if gain > 0:
                        raise ValueError(
                            "When gain_mode = 'relative', gain must be < 0. This sets\
                                the gain relative to the maximum possible gain."
                        )
                    target_gain = max_gain + gain
                else:
                    target_gain = gain

                if target_gain < min_gain or target_gain > max_gain:
                    print(
                        f"Requested gain {target_gain} dB out of range;\
                              clamping to valid span {min_gain}-{max_gain} dB."
                    )
                    target_gain = min(max(target_gain, min_gain), max_gain)

                target_gain = min(available_gains, key=lambda g: abs(g - target_gain))

            self.radio.set_gain(target_gain)
            self.rx_gain = self.radio.get_gain()

        print(f"RTL RX Gain = {self.radio.get_gain()}")
        print(f"Available RTL RX Gains: {available_gains}")

    def record(self, num_samples: Optional[int] = None, rx_time: Optional[int | float] = None):
        """
        Create a radio recording (iq samples and metadata) of a given length from the RTL-SDR.
        Either num_samples or rx_time must be provided.
        init_rx() must be called before record()

        :param num_samples: The number of samples to record.
        :type num_samples: int, optional
        :param rx_time: The time to record.
        :type rx_time: int or float, optional

        returns: Recording object (iq samples and metadata)
        """

        if not self._rx_initialized:
            raise RuntimeError("RX was not initialized. init_rx() must be called before record().")

        if num_samples is not None and rx_time is not None:
            raise ValueError("Only input one of num_samples or rx_time")
        elif num_samples is not None:
            pass
        elif rx_time is not None:
            num_samples = int(rx_time * self.rx_sample_rate)
        else:
            raise ValueError("Must provide input of one of num_samples or rx_time")

        # RTL-SDR has USB buffer limitations - use consistent 256k chunks
        # Always read full chunks to avoid USB overflow issues with partial reads
        max_samples_per_read = 262144  # 256k samples = stable chunk size
        num_full_reads = num_samples // max_samples_per_read
        remainder = num_samples % max_samples_per_read
        signal = np.array([], dtype=np.complex64)

        print("RTL-SDR Starting RX...")

        # Read full chunks
        for _ in range(num_full_reads):
            try:
                chunk = self.radio.read_samples(max_samples_per_read)
                signal = np.append(signal, chunk)
            except Exception as e:
                print(f"Error while reading samples: {e}")
                break

        # Read remainder if needed (round up to power of 2 for USB compatibility)
        if remainder > 0 and len(signal) == num_full_reads * max_samples_per_read:
            # Round up to next 16k boundary for USB stability
            padded_remainder = ((remainder + 16383) // 16384) * 16384
            try:
                chunk = self.radio.read_samples(padded_remainder)
                signal = np.append(signal, chunk[:remainder])  # Only keep what we need
            except Exception as e:
                print(f"Error while reading final chunk: {e}")

        print("RTL-SDR RX Completed.")

        metadata = {
            "source": self.__class__.__name__,
            "sample_rate": self.rx_sample_rate,
            "center_frequency": self.rx_center_frequency,
            "gain": self.rx_gain,
        }

        return Recording(data=signal, metadata=metadata)

    def _stream_rx(self, callback):
        if not self._rx_initialized:
            raise RuntimeError("RX was not initialized. init_rx() must be called before _stream_rx() or record().")

        print("RTL-SDR Starting RX...")
        self._enable_rx = True
        try:
            while self._enable_rx:
                samples = self.radio.read_samples(self.rx_buffer_size)
                callback(buffer=np.asarray(samples, dtype=np.complex64), metadata=None)
        finally:
            print("RTL-SDR RX Completed.")

    def _stream_tx(self, callback):  # pragma: no cover - RTL-SDR is RX only
        raise NotImplementedError("RTL-SDR does not support transmit operations")

    def init_tx(self, *args, **kwargs):  # pragma: no cover - RTL-SDR is RX only
        raise NotImplementedError("RTL-SDR does not support transmit operations")

    def tx_recording(
        self, recording: Recording | np.ndarray | list, num_samples=None, tx_time=None
    ):  # pragma: no cover - RTL-SDR is RX only
        raise NotImplementedError("RTL-SDR does not support transmit operations")

    def supports_bias_tee(self) -> bool:
        return True

    def set_bias_tee(self, enable: bool):
        self.radio.set_bias_tee(bool(enable))
        state = "enabled" if enable else "disabled"
        print(f"RTL-SDR bias tee {state}.")

    def set_clock_source(self, source):  # pragma: no cover - not applicable to RTL-SDR
        raise NotImplementedError("RTL-SDR does not support external clock configuration")

    def close(self):
        try:
            self.radio.close()
        finally:
            self._enable_rx = False
            self._enable_tx = False
