# SPDX-FileCopyrightText: © 2024-2026 Jimmy Fitzpatrick <jcfitzpatrick12@gmail.com>
# This file is part of SPECTRE
# SPDX-License-Identifier: GPL-3.0-or-later

import dataclasses

from gnuradio import sdrplay3
from gnuradio import spectre

import spectre_core.fields

from ._base import Base, BaseModel
from ._constants import FlowgraphConstant


@dataclasses.dataclass(frozen=True)
class RSPduoPort:
    """Specifies one of the antenna ports on the RSPduo.

    These are easier to type in the CLI tool, than the values we must pass to `gr-sdrplay3`.
    Namely 'Tuner 1 50' ohm and 'Tuner 2 50 ohm'.
    """

    TUNER_1 = "tuner_1"
    TUNER_2 = "tuner_2"


def _map_port(antenna_port: str | None) -> str:
    """Maps the CLI-typeable port, to the one used in the constructor for the RSPduo block in the gr-sdrplay3 OOT module."""
    if antenna_port == RSPduoPort.TUNER_1:
        return "Tuner 1 50 ohm"
    elif antenna_port == RSPduoPort.TUNER_2:
        return "Tuner 2 50 ohm"
    else:
        raise ValueError(f"{antenna_port} is not a valid antenna port.")


class RSPduoFixedCenterFrequencyModel(BaseModel):
    sample_rate: spectre_core.fields.Field.sample_rate = 2e6
    batch_size: spectre_core.fields.Field.batch_size = 3
    center_frequency: spectre_core.fields.Field.center_frequency = 95.8e6
    bandwidth: spectre_core.fields.Field.bandwidth = 1.536e6
    if_gain: spectre_core.fields.Field.if_gain = -30
    rf_gain: spectre_core.fields.Field.rf_gain = 0
    output_type: spectre_core.fields.Field.output_type = (
        spectre_core.fields.OutputType.FC32
    )
    antenna_port: spectre_core.fields.Field.antenna_port = RSPduoPort.TUNER_1


class RSPduoFixedCenterFrequency(Base[RSPduoFixedCenterFrequencyModel]):
    def configure(self, tag: str, model: RSPduoFixedCenterFrequencyModel) -> None:
        self.spectre_batched_file_sink = spectre.batched_file_sink(
            self._batches_dir_path,
            tag,
            model.output_type,
            model.batch_size,
            model.sample_rate,
            FlowgraphConstant.GROUP_BY_DATE,
        )
        self.sdrplay3_rspduo = sdrplay3.rspduo(
            "",
            rspduo_mode="Single Tuner",
            stream_args=sdrplay3.stream_args(
                output_type=model.output_type, channels_size=1
            ),
        )
        self.sdrplay3_rspduo.set_sample_rate(model.sample_rate)
        self.sdrplay3_rspduo.set_center_freq(model.center_frequency)
        self.sdrplay3_rspduo.set_bandwidth(model.bandwidth)
        self.sdrplay3_rspduo.set_antenna(_map_port(model.antenna_port))
        self.sdrplay3_rspduo.set_gain_mode(False)
        self.sdrplay3_rspduo.set_gain(model.if_gain, "IF")
        self.sdrplay3_rspduo.set_gain(model.rf_gain, "RF")
        self.sdrplay3_rspduo.set_freq_corr(0)
        self.sdrplay3_rspduo.set_dc_offset_mode(False)
        self.sdrplay3_rspduo.set_iq_balance_mode(False)
        self.sdrplay3_rspduo.set_agc_setpoint(-30)
        self.sdrplay3_rspduo.set_rf_notch_filter(False)
        self.sdrplay3_rspduo.set_dab_notch_filter(False)
        self.sdrplay3_rspduo.set_am_notch_filter(False)
        self.sdrplay3_rspduo.set_biasT(False)
        self.sdrplay3_rspduo.set_debug_mode(False)
        self.sdrplay3_rspduo.set_sample_sequence_gaps_check(False)
        self.sdrplay3_rspduo.set_show_gain_changes(False)

        self.connect((self.sdrplay3_rspduo, 0), (self.spectre_batched_file_sink, 0))


class RSPduoSweptCenterFrequencyModel(BaseModel):
    sample_rate: spectre_core.fields.Field.sample_rate = 2e6
    batch_size: spectre_core.fields.Field.batch_size = 3
    bandwidth: spectre_core.fields.Field.bandwidth = 1.536e6
    if_gain: spectre_core.fields.Field.if_gain = -30
    rf_gain: spectre_core.fields.Field.rf_gain = 0
    min_frequency: spectre_core.fields.Field.min_frequency = 95e6
    max_frequency: spectre_core.fields.Field.max_frequency = 100e6
    dwell_time: spectre_core.fields.Field.dwell_time = 0.15
    frequency_hop: spectre_core.fields.Field.frequency_hop = 2e6
    output_type: spectre_core.fields.Field.output_type = (
        spectre_core.fields.OutputType.FC32
    )
    antenna_port: spectre_core.fields.Field.antenna_port = RSPduoPort.TUNER_1


class RSPduoSweptCenterFrequency(Base[RSPduoSweptCenterFrequencyModel]):
    def configure(self, tag: str, model: RSPduoSweptCenterFrequencyModel) -> None:
        retune_cmd_name = "freq"
        self.spectre_frequency_sweeper = spectre.frequency_sweeper(
            model.min_frequency,
            model.max_frequency,
            model.frequency_hop,
            model.dwell_time,
            model.sample_rate,
            retune_cmd_name,
            model.output_type,
        )

        is_tagged = True
        frequency_tag_key = "freq"
        self.spectre_batched_file_sink = spectre.batched_file_sink(
            self._batches_dir_path,
            tag,
            model.output_type,
            model.batch_size,
            model.sample_rate,
            FlowgraphConstant.GROUP_BY_DATE,
            is_tagged,
            frequency_tag_key,
            model.min_frequency,
        )
        self.sdrplay3_rspduo = sdrplay3.rspduo(
            "",
            rspduo_mode="Single Tuner",
            stream_args=sdrplay3.stream_args(
                output_type=model.output_type, channels_size=1
            ),
        )
        self.sdrplay3_rspduo.set_sample_rate(model.sample_rate, True)
        self.sdrplay3_rspduo.set_center_freq(model.min_frequency, True)
        self.sdrplay3_rspduo.set_bandwidth(model.bandwidth)
        self.sdrplay3_rspduo.set_antenna(_map_port(model.antenna_port))
        self.sdrplay3_rspduo.set_gain_mode(False)
        self.sdrplay3_rspduo.set_gain(model.if_gain, "IF", True)
        self.sdrplay3_rspduo.set_gain(model.rf_gain, "RF", True)
        self.sdrplay3_rspduo.set_freq_corr(0)
        self.sdrplay3_rspduo.set_dc_offset_mode(False)
        self.sdrplay3_rspduo.set_iq_balance_mode(False)
        self.sdrplay3_rspduo.set_agc_setpoint(-30)
        self.sdrplay3_rspduo.set_rf_notch_filter(False)
        self.sdrplay3_rspduo.set_dab_notch_filter(False)
        self.sdrplay3_rspduo.set_am_notch_filter(False)
        self.sdrplay3_rspduo.set_biasT(False)
        self.sdrplay3_rspduo.set_stream_tags(True)
        self.sdrplay3_rspduo.set_debug_mode(False)
        self.sdrplay3_rspduo.set_sample_sequence_gaps_check(False)
        self.sdrplay3_rspduo.set_show_gain_changes(False)

        self.msg_connect(
            (self.spectre_frequency_sweeper, "retune_command"),
            (self.sdrplay3_rspduo, "command"),
        )
        self.connect((self.sdrplay3_rspduo, 0), (self.spectre_batched_file_sink, 0))
        self.connect((self.sdrplay3_rspduo, 0), (self.spectre_frequency_sweeper, 0))
