use embassy_stm32::{gpio::Output, mode::Blocking, spi::Spi, spi::Error};
use embassy_time::Timer;

use crate::ad5940_registers::*;

pub struct AD5940 {
    spi: Spi<'static, Blocking>,
    cs: Output<'static>,
    rst: Output<'static>,
}

impl AD5940 {
    pub fn new(spi: Spi<'static, Blocking>, cs: Output<'static>, rst: Output<'static>) -> Self {
        AD5940 { spi, cs, rst }
    }

    pub async fn reset(&mut self) -> Result<(), Error> {
        self.rst.set_low();
        Timer::after_millis(1).await;
        self.rst.set_high();
        Ok(())
    }

    async fn read_reg_16(&mut self, address: Register) -> Result<u16, Error> {
        // Write address command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_SETADDR as u8])?;
        self.spi.blocking_write(&[address as u16])?;
        self.cs.set_high();

        // Wait after cs is high
        Timer::after_nanos(80).await;

        // Read command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_READREG as u8])?;

        let mut data = [0u8; 3]; // First byte dummy, then two data bytes
        self.spi.blocking_read(&mut data)?;
        self.cs.set_high();

        Ok(u16::from_be_bytes([data[1], data[2]]))
    }

    async fn read_reg_32(&mut self, address: Register) -> Result<u32, Error> {
        // Write address command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_SETADDR as u8])?;
        self.spi.blocking_write(&[address as u16])?;
        self.cs.set_high();

        // Wait after cs is high
        Timer::after_nanos(80).await;

        // Read command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_READREG as u8])?;

        let mut data = [0u8; 5]; // First byte dummy, then four data bytes
        self.spi.blocking_read(&mut data)?;
        self.cs.set_high();

        Ok(u32::from_be_bytes([data[1], data[2], data[3], data[4]]))
    }

    async fn write_reg_32(&mut self, address: Register, value: u32) -> Result<(), Error> {
       self.write_reg_32_raw(address as u16, value).await
    }

    async fn write_reg_32_raw(&mut self, address: u16, value: u32) -> Result<(), Error> {
        // Write address command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_SETADDR as u8])?;
        self.spi.blocking_write(&[address as u16])?;
        self.cs.set_high();

        // Wait after cs is high
        Timer::after_nanos(80).await;

        // Write value command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_WRITEREG as u8])?;
        self.spi.blocking_write(&[value])?;
        self.cs.set_high();

        Ok(())
    }

    async fn write_reg_16_raw(&mut self, address: u16, value: u16) -> Result<(), Error> {
        // Write address command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_SETADDR as u8])?;
        self.spi.blocking_write(&[address as u16])?;
        self.cs.set_high();

        // Wait after cs is high
        Timer::after_nanos(80).await;

        // Write value command
        self.cs.set_low();
        self.spi.blocking_write(&[Command::SPICMD_WRITEREG as u8])?;
        self.spi.blocking_write(&[value])?;
        self.cs.set_high();

        Ok(())
    }

    pub async fn system_init(&mut self) -> Result<(), Error> {
        // See table 14
        self.write_reg_16_raw(0x0908, 0x02C9).await?;
        self.write_reg_16_raw(0x0C08, 0x206C).await?;
        self.write_reg_16_raw(0x21F0, 0x0010).await?;
        self.write_reg_16_raw(0x0410, 0x02C9).await?;
        self.write_reg_16_raw(0x0A28, 0x0009).await?;
        self.write_reg_16_raw(0x238C, 0x0104).await?;
        self.write_reg_16_raw(0x0A04, 0x4859).await?;
        self.write_reg_16_raw(0x0A04, 0xF27B).await?;
        self.write_reg_16_raw(0x0A00, 0x8009).await?;
        self.write_reg_16_raw(0x22F0, 0x0000).await?;

        // According to library
        // self.write_reg_32_raw(0x2230, 0xDE87A5AF).await?;
        // self.write_reg_16_raw(0x2250, 0x103F).await?;
        // self.write_reg_16_raw(0x22B0, 0x203C).await?;
        // self.write_reg_32_raw(0x2230, 0xDE87A5A0).await?;

        Ok(())
    }

    pub async fn afecon(&mut self, ctr: ConfigurationRegister, state: bool) {
        let reg = self.read_reg_32(Register::AFECON).await.unwrap();
        let mut reg = ConfigurationRegister::from_bits_truncate(reg);

        let flags = [
            ConfigurationRegister::DACBUFEN,
            ConfigurationRegister::DACREFEN,
            ConfigurationRegister::SINC2EN,
            ConfigurationRegister::DFTEN,
            ConfigurationRegister::WAVEGENEN,
            ConfigurationRegister::TEMPCONVEN,
            ConfigurationRegister::TEMPSENSEN,
            ConfigurationRegister::TIAEN,
            ConfigurationRegister::INAMPEN,
            ConfigurationRegister::EXBUFEN,
            ConfigurationRegister::ADCCONVEN,
            ConfigurationRegister::ADCEN,
            ConfigurationRegister::DACEN,
            ConfigurationRegister::HSREFDIS,
        ];

        for &flag in flags.iter() {
            if ctr.contains(flag) {
                reg = if state {
                    reg | flag
                } else {
                    reg & !flag
                };
            }
        }

        self.write_reg_32(Register::AFECON, reg.bits()).await.unwrap();
    }

    pub async fn swcon(&mut self, ctr: SwitchMatrixConfigurationRegister) {
        let reg = self.read_reg_32(Register::SWCON).await.unwrap();
        let mut reg = SwitchMatrixConfigurationRegister::from_bits_truncate(reg);

        // NMUXCON
        if (ctr & SwitchMatrixConfigurationRegister::NMUXCON_MSK) == SwitchMatrixConfigurationRegister::NMUXCON_N2 {
            reg &= !SwitchMatrixConfigurationRegister::NMUXCON_MSK;
            reg |= SwitchMatrixConfigurationRegister::NMUXCON_N2;
        } else if (ctr & SwitchMatrixConfigurationRegister::NMUXCON_MSK) == SwitchMatrixConfigurationRegister::NMUXCON_N5 {
            reg &= !SwitchMatrixConfigurationRegister::NMUXCON_MSK;
            reg |= SwitchMatrixConfigurationRegister::NMUXCON_N5;
        };

        // PMUXCON
        if (ctr & SwitchMatrixConfigurationRegister::PMUXCON_MSK) == SwitchMatrixConfigurationRegister::PMUXCON_P2 {
            reg &= !SwitchMatrixConfigurationRegister::PMUXCON_MSK;
            reg |= SwitchMatrixConfigurationRegister::PMUXCON_P2;
        } else if (ctr & SwitchMatrixConfigurationRegister::PMUXCON_MSK) == SwitchMatrixConfigurationRegister::PMUXCON_P11 {
            reg &= !SwitchMatrixConfigurationRegister::PMUXCON_MSK;
            reg |= SwitchMatrixConfigurationRegister::PMUXCON_P11;
        }

        // DMUXCON
        if ctr.contains(SwitchMatrixConfigurationRegister::DMUXCON_D5) {
            reg &= !SwitchMatrixConfigurationRegister::DMUXCON_MSK;
            reg |= SwitchMatrixConfigurationRegister::DMUXCON_D5;
        }

        self.write_reg_32(Register::SWCON, reg.bits()).await.unwrap();
    }

    pub async fn wgfcw(&mut self, frequency: u32) {
        let sinefcw = (frequency as f64) * (1_073_741_824.0 / 16_000_000.0);
        let sinefcw = 0x00FFFFFF & sinefcw as u32;
        self.write_reg_32(Register::WGFCW, sinefcw).await.unwrap();
    }

    pub async fn get_chipid(&mut self) -> Result<u16, Error> {
        self.read_reg_16(Register::CHIPID).await
    }

    pub async fn get_adiid(&mut self) -> Result<u16, Error> {
        self.read_reg_16(Register::ADIID).await
    }

    pub async fn init_temperature(&mut self) -> Result<(), Error> {
        // AFECON:
        self.afecon(ConfigurationRegister::TEMPSENSEN | ConfigurationRegister::ADCEN, true).await;

        // ADCCON
        let config_adccon = ADCConfigurationRegister::GNPGA_1_5.bits() | ADCConfigurationRegister::MUXSELN_TEMP.bits() | ADCConfigurationRegister::MUXSELP_TEMP.bits();
        self.write_reg_32(Register::ADCCON, config_adccon).await?;

        // AFECON - start conversion
        self.afecon(ConfigurationRegister::TEMPCONVEN | ConfigurationRegister::ADCCONVEN, true).await;

        Ok(())
    }

    pub async fn get_temperature(&mut self) -> Result<f32, Error> {
        // See page 57 of the datasheet for temperature calculation
        let pga_gain = 1.5;
        let k = 8.13;

        let mut tempsensdat0: u32 = self.read_reg_32(Register::TEMPSENSDAT).await?;
        tempsensdat0 &= 0x0000FFFF;

        Ok((tempsensdat0 as f32/(pga_gain * k)) - 273.15)
    }

    pub async fn init_waveform(&mut self) -> Result<(), Error> {
        // Set frequency
        let freq: u32 = 1000;
        self.wgfcw(freq).await;

        // Init amplitude
        let wg_amplitude = 2047; // 2047 is the maximum amplitude for a 12-bit DAC --> 1.62V peak-to-peak
        self.write_reg_32(Register::WGAMPLITUDE, wg_amplitude).await?;

        // WGCON, enable waveform generator and set to sinusoidal, NO DACCAL
        let config_wgcon = WaveformGeneratorConfigurationRegister::TYPESEL_SIN.bits();
        self.write_reg_32(Register::WGCON, config_wgcon).await?;

        // SWCON - set up switch matrix
        self.swcon(
            SwitchMatrixConfigurationRegister::NMUXCON_N2
            | SwitchMatrixConfigurationRegister::PMUXCON_P11
            | SwitchMatrixConfigurationRegister::DMUXCON_D5).await;

        // AFECON:
        self.afecon(
            ConfigurationRegister::DACREFEN
            | ConfigurationRegister::EXBUFEN
            | ConfigurationRegister::INAMPEN
            | ConfigurationRegister::DACEN
            | ConfigurationRegister::WAVEGENEN,
            true,
        )
        .await;

        Ok(())
    }
}

#[allow(dead_code)]
#[allow(non_camel_case_types)]
enum Command {
    SPICMD_SETADDR  = 0x20,
    SPICMD_READREG  = 0x6D,
    SPICMD_WRITEREG = 0x2D,
    SPICMD_READFIFO = 0x5F,
}

#[allow(dead_code)]
#[allow(non_camel_case_types)]
#[repr(u16)]
enum Register {
    ADIID       = 0x0000_0400, // Analog Devices Inc., identification register
    CHIPID      = 0x0000_0404, // Chip identification register
    AFECON      = 0x0000_2000, // Configuration Register
    SWCON       = 0x0000_200C, // Switch Matrix Configuration Register
    WGCON       = 0x0000_2014, // Waveform Generator Configuration Register
    WGFCW       = 0x0000_2030, // Waveform Generator, Sinusoid Frequency Control Word Register
    WGAMPLITUDE = 0x0000_203C, // Waveform Generator, Sinusoid Amplitude Register
    TEMPSENSDAT = 0x0000_2084, // Temperature Sensor Result Register
    ADCDAT      = 0x0000_2074, // ADC Raw Result Register
    TEMPSENS    = 0x0000_2174, // Temperature Sensor Configuration Register
    ADCCON      = 0x0000_21A8, // ADC Configuration Register
    PMBW        = 0x0000_22F0, // Power Mode Configuration Register
}