//! Platform-agnostic driver for the **AS7341** 11-channel spectral sensor.
//!
//! This crate uses [`embedded-hal`](https://docs.rs/embedded-hal/1) I2C and
//! delay traits, so it works on any platform that implements them (ESP32,
//! STM32, nRF, RP2040, Linux, …).
//!
//! # Quick start
//!
//! ```rust,ignore
//! let mut sensor = As7341::new(i2c, delay);
//! sensor.init()?;
//!
//! let data = sensor.read_all_channels()?;
//! // data.f1 … data.f8, data.clear, data.nir
//! ```

#![no_std]

mod register;

use embedded_hal::delay::DelayNs;
use embedded_hal::i2c::I2c;

/// Default 7-bit I2C address of the AS7341.
const ADDRESS: u8 = 0x39;

/// Maximum polling iterations (1 ms each) when waiting for hardware.
const POLL_TIMEOUT: u32 = 2000;

// ── Error ────────────────────────────────────────────────────────────

/// Driver errors.
#[derive(Debug)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum Error<E> {
    /// I2C bus error.
    I2c(E),
    /// The chip ID register did not match the expected AS7341 value.
    InvalidDeviceId(u8),
    /// A polling loop exceeded its timeout.
    Timeout,
}

impl<E> From<E> for Error<E> {
    fn from(e: E) -> Self {
        Error::I2c(e)
    }
}

// ── Gain ─────────────────────────────────────────────────────────────

/// ADC gain setting.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[repr(u8)]
pub enum Gain {
    X0_5 = 0,
    X1 = 1,
    X2 = 2,
    X4 = 3,
    X8 = 4,
    X16 = 5,
    X32 = 6,
    X64 = 7,
    X128 = 8,
    X256 = 9,
    X512 = 10,
}

impl Gain {
    /// Multiplier as a floating-point value (e.g. `Gain::X16` → `16.0`).
    pub fn multiplier(self) -> f32 {
        match self {
            Gain::X0_5 => 0.5,
            Gain::X1 => 1.0,
            Gain::X2 => 2.0,
            Gain::X4 => 4.0,
            Gain::X8 => 8.0,
            Gain::X16 => 16.0,
            Gain::X32 => 32.0,
            Gain::X64 => 64.0,
            Gain::X128 => 128.0,
            Gain::X256 => 256.0,
            Gain::X512 => 512.0,
        }
    }
}

// ── Channel data ─────────────────────────────────────────────────────

/// Raw ADC readings for all 10 unique spectral channels.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub struct ChannelData {
    /// F1 – 415 nm (violet)
    pub f1: u16,
    /// F2 – 445 nm (indigo)
    pub f2: u16,
    /// F3 – 480 nm (blue)
    pub f3: u16,
    /// F4 – 515 nm (cyan)
    pub f4: u16,
    /// F5 – 555 nm (green)
    pub f5: u16,
    /// F6 – 590 nm (yellow)
    pub f6: u16,
    /// F7 – 630 nm (orange)
    pub f7: u16,
    /// F8 – 680 nm (red)
    pub f8: u16,
    /// Clear (unfiltered)
    pub clear: u16,
    /// Near-infrared
    pub nir: u16,
}

impl ChannelData {
    /// Channel values as an array: F1–F8, Clear, NIR.
    pub fn as_array(&self) -> [u16; 10] {
        [
            self.f1, self.f2, self.f3, self.f4, self.f5, self.f6, self.f7, self.f8,
            self.clear, self.nir,
        ]
    }
}

/// Channel names for display: `["F1", "F2", … "Clear", "NIR"]`.
pub const CHANNEL_NAMES: [&str; 10] = [
    "F1", "F2", "F3", "F4", "F5", "F6", "F7", "F8", "Clear", "NIR",
];

/// Peak wavelengths in nm (0 = broadband).
pub const CHANNEL_NM: [u16; 10] = [415, 445, 480, 515, 555, 590, 630, 680, 0, 0];

/// All gain variants in ascending order.
pub const ALL_GAINS: [Gain; 11] = [
    Gain::X0_5, Gain::X1, Gain::X2, Gain::X4, Gain::X8, Gain::X16,
    Gain::X32, Gain::X64, Gain::X128, Gain::X256, Gain::X512,
];

// ── Driver ───────────────────────────────────────────────────────────

/// AS7341 driver.
pub struct As7341<I2C, D> {
    i2c: I2C,
    delay: D,
}

impl<I2C, D> As7341<I2C, D> {
    /// Create a new driver instance. Does **not** touch the bus yet – call
    /// [`init`](Self::init) to power on the sensor and verify the chip ID.
    pub fn new(i2c: I2C, delay: D) -> Self {
        Self { i2c, delay }
    }

    /// Consume the driver and return the I2C bus and delay.
    pub fn destroy(self) -> (I2C, D) {
        (self.i2c, self.delay)
    }
}

impl<I2C: I2c, D: DelayNs> As7341<I2C, D> {
    // ── Lifecycle ────────────────────────────────────────────────────

    /// Power on the sensor, verify the chip ID, and apply sensible
    /// defaults (ATIME = 29, ASTEP = 599 → ≈ 50 ms integration, gain 8×).
    pub fn init(&mut self) -> Result<(), Error<I2C::Error>> {
        // Power on, everything else off (clears any in-flight measurement).
        self.write_reg(register::ENABLE, register::ENABLE_PON)?;
        self.delay.delay_ms(10);

        // Verify chip identity.
        let id = self.read_reg(register::ID)?;
        if id & 0xFC != register::CHIP_ID {
            return Err(Error::InvalidDeviceId(id));
        }

        // Reasonable defaults.
        self.set_atime(29)?;
        self.set_astep(599)?;
        self.set_gain(Gain::X8)?;
        Ok(())
    }

    /// Read the raw device ID register (0x92).
    pub fn device_id(&mut self) -> Result<u8, Error<I2C::Error>> {
        Ok(self.read_reg(register::ID)?)
    }

    // ── Configuration ────────────────────────────────────────────────

    /// Set the integration step count (0–255).
    ///
    /// Integration time = (ATIME + 1) × (ASTEP + 1) × 2.78 µs.
    pub fn set_atime(&mut self, atime: u8) -> Result<(), Error<I2C::Error>> {
        Ok(self.write_reg(register::ATIME, atime)?)
    }

    /// Set the integration step size (0–65534).
    ///
    /// Integration time = (ATIME + 1) × (ASTEP + 1) × 2.78 µs.
    pub fn set_astep(&mut self, astep: u16) -> Result<(), Error<I2C::Error>> {
        self.write_reg(register::ASTEP_L, astep as u8)?;
        self.write_reg(register::ASTEP_H, (astep >> 8) as u8)?;
        Ok(())
    }

    /// Set the ADC gain.
    pub fn set_gain(&mut self, gain: Gain) -> Result<(), Error<I2C::Error>> {
        Ok(self.write_reg(register::CFG1, gain as u8)?)
    }

    /// Read the current gain setting.
    pub fn gain(&mut self) -> Result<Gain, Error<I2C::Error>> {
        let raw = self.read_reg(register::CFG1)? & 0x1F;
        Ok(match raw {
            0 => Gain::X0_5,
            1 => Gain::X1,
            2 => Gain::X2,
            3 => Gain::X4,
            4 => Gain::X8,
            5 => Gain::X16,
            6 => Gain::X32,
            7 => Gain::X64,
            8 => Gain::X128,
            9 => Gain::X256,
            _ => Gain::X512,
        })
    }

    // ── Measurement ──────────────────────────────────────────────────

    /// Perform two SMUX-multiplexed measurements and return all 10
    /// spectral channels.
    ///
    /// This is the main entry point: it configures the SMUX for F1–F4,
    /// reads those four channels plus Clear and NIR, then reconfigures
    /// for F5–F8 and reads again.
    pub fn read_all_channels(&mut self) -> Result<ChannelData, Error<I2C::Error>> {
        // ── First pass: F1 F2 F3 F4 Clear NIR ───────────────────────
        self.disable_spectral()?;
        self.configure_smux(&register::SMUX_F1_F4_CLEAR_NIR)?;
        self.start_spectral_and_wait()?;
        let set1 = self.read_6_adc()?;

        // ── Second pass: F5 F6 F7 F8 Clear NIR ──────────────────────
        self.disable_spectral()?;
        self.configure_smux(&register::SMUX_F5_F8_CLEAR_NIR)?;
        self.start_spectral_and_wait()?;
        let set2 = self.read_6_adc()?;

        Ok(ChannelData {
            f1: set1[0],
            f2: set1[1],
            f3: set1[2],
            f4: set1[3],
            f5: set2[0],
            f6: set2[1],
            f7: set2[2],
            f8: set2[3],
            clear: set1[4],
            nir: set1[5],
        })
    }

    // ── LED control ──────────────────────────────────────────────────

    /// Turn the on-board LED on or off and set its drive current.
    ///
    /// `current_ma` is clamped to the range 4–258 mA (step 2 mA).
    /// Pass `enable = false` to turn the LED off (current value is
    /// still written so that the next `enable = true` uses it).
    pub fn set_led(
        &mut self,
        enable: bool,
        current_ma: u16,
    ) -> Result<(), Error<I2C::Error>> {
        let drive = if current_ma <= 4 {
            0u8
        } else {
            (((current_ma.min(258) - 4) / 2) as u8).min(0x7F)
        };

        // Switch to low register bank.
        self.set_reg_bank(true)?;

        // Enable LED_SEL in CONFIG (0x70).
        let cfg = self.read_reg(register::CONFIG)?;
        self.write_reg(register::CONFIG, cfg | register::CONFIG_LED_SEL)?;

        // Write LED register.
        let led_val = if enable { register::LED_EN | drive } else { drive };
        self.write_reg(register::LED, led_val)?;

        // Switch back to default bank.
        self.set_reg_bank(false)?;
        Ok(())
    }

    // ── Power ────────────────────────────────────────────────────────

    /// Enable internal oscillator (PON).
    pub fn power_on(&mut self) -> Result<(), Error<I2C::Error>> {
        let en = self.read_reg(register::ENABLE)?;
        self.write_reg(register::ENABLE, en | register::ENABLE_PON)?;
        self.delay.delay_ms(10);
        Ok(())
    }

    /// Disable internal oscillator (clears ENABLE to 0).
    pub fn power_off(&mut self) -> Result<(), Error<I2C::Error>> {
        Ok(self.write_reg(register::ENABLE, 0)?)
    }

    // ── Flicker detection ────────────────────────────────────────────

    /// Run a single flicker-detection cycle.
    ///
    /// Returns `Some(100)` or `Some(120)` if a mains frequency is
    /// detected, or `None` if the result is unknown / saturated.
    pub fn detect_flicker(&mut self) -> Result<Option<u16>, Error<I2C::Error>> {
        self.disable_spectral()?;

        // Configure flicker-detection timing.
        // FD_TIME = 0x40 → integration period ≈ 1 ms (good for 50/60 Hz).
        // FD_GAIN = 9 (256×) in bits [7:3] of FD_TIME2.
        // FD_CFG0 = 0 → max periods for averaging.
        self.write_reg(register::FD_CFG0, 0x00)?;
        self.write_reg(register::FD_TIME1, 0x40)?;
        self.write_reg(register::FD_TIME2, 9 << 3)?;

        // Load flicker-detection SMUX.
        self.configure_smux(&register::SMUX_FLICKER)?;

        // Clear FD_STATUS by reading it.
        let _ = self.read_reg(register::FD_STATUS)?;

        // Enable flicker detection.
        let en = self.read_reg(register::ENABLE)?;
        self.write_reg(
            register::ENABLE,
            (en & !register::ENABLE_SP_EN) | register::ENABLE_FDEN,
        )?;

        // Wait for FD_STATUS.FD_VALID (bit 5).
        for _ in 0..POLL_TIMEOUT {
            let st = self.read_reg(register::FD_STATUS)?;
            if st & (1 << 5) != 0 {
                // Disable flicker detection.
                let en = self.read_reg(register::ENABLE)?;
                self.write_reg(register::ENABLE, en & !register::ENABLE_FDEN)?;

                return Ok(match st & 0x07 {
                    0b100 => Some(100),
                    0b101 => Some(120),
                    _ => None,
                });
            }
            self.delay.delay_ms(1);
        }

        // Disable flicker detection before returning timeout.
        let en = self.read_reg(register::ENABLE)?;
        self.write_reg(register::ENABLE, en & !register::ENABLE_FDEN)?;
        Err(Error::Timeout)
    }

    // ── Private helpers ──────────────────────────────────────────────

    fn read_reg(&mut self, reg: u8) -> Result<u8, I2C::Error> {
        let mut buf = [0u8];
        self.i2c.write_read(ADDRESS, &[reg], &mut buf)?;
        Ok(buf[0])
    }

    fn write_reg(&mut self, reg: u8, val: u8) -> Result<(), I2C::Error> {
        self.i2c.write(ADDRESS, &[reg, val])
    }

    fn read_regs(&mut self, start: u8, buf: &mut [u8]) -> Result<(), I2C::Error> {
        self.i2c.write_read(ADDRESS, &[start], buf)
    }

    /// Read the six 16-bit ADC channels (CH0–CH5) starting at 0x95.
    fn read_6_adc(&mut self) -> Result<[u16; 6], Error<I2C::Error>> {
        let mut buf = [0u8; 12];
        self.read_regs(register::CH0_DATA_L, &mut buf)?;
        Ok([
            u16::from_le_bytes([buf[0], buf[1]]),
            u16::from_le_bytes([buf[2], buf[3]]),
            u16::from_le_bytes([buf[4], buf[5]]),
            u16::from_le_bytes([buf[6], buf[7]]),
            u16::from_le_bytes([buf[8], buf[9]]),
            u16::from_le_bytes([buf[10], buf[11]]),
        ])
    }

    /// Write a 20-byte SMUX configuration and execute it.
    fn configure_smux(&mut self, config: &[u8; 20]) -> Result<(), Error<I2C::Error>> {
        // Write SMUX RAM (addresses 0x00–0x13).
        for (i, &val) in config.iter().enumerate() {
            self.write_reg(i as u8, val)?;
        }

        // Set SMUX_CMD = WRITE in CFG6.
        self.write_reg(register::CFG6, register::CFG6_SMUX_CMD_WRITE)?;

        // Trigger SMUX execution.
        let en = self.read_reg(register::ENABLE)?;
        self.write_reg(register::ENABLE, en | register::ENABLE_SMUXEN)?;

        // Wait for SMUXEN to auto-clear.
        for _ in 0..POLL_TIMEOUT {
            if self.read_reg(register::ENABLE)? & register::ENABLE_SMUXEN == 0 {
                return Ok(());
            }
            self.delay.delay_ms(1);
        }
        Err(Error::Timeout)
    }

    /// Set SP_EN to start a spectral measurement and wait for AVALID.
    fn start_spectral_and_wait(&mut self) -> Result<(), Error<I2C::Error>> {
        let en = self.read_reg(register::ENABLE)?;
        self.write_reg(register::ENABLE, en | register::ENABLE_SP_EN)?;

        for _ in 0..POLL_TIMEOUT {
            if self.read_reg(register::STATUS2)? & register::STATUS2_AVALID != 0 {
                return Ok(());
            }
            self.delay.delay_ms(1);
        }
        Err(Error::Timeout)
    }

    /// Clear SP_EN.
    fn disable_spectral(&mut self) -> Result<(), Error<I2C::Error>> {
        let en = self.read_reg(register::ENABLE)?;
        self.write_reg(register::ENABLE, en & !register::ENABLE_SP_EN)?;
        Ok(())
    }

    /// Switch the register bank. `low = true` → access 0x60–0x74
    /// (CONFIG, LED, …). `low = false` → default bank.
    fn set_reg_bank(&mut self, low: bool) -> Result<(), Error<I2C::Error>> {
        let cfg0 = self.read_reg(register::CFG0)?;
        let val = if low {
            cfg0 | register::CFG0_REG_BANK
        } else {
            cfg0 & !register::CFG0_REG_BANK
        };
        self.write_reg(register::CFG0, val)?;
        Ok(())
    }
}
