#![no_main]
#![no_std]

extern crate panic_probe;
#[defmt::panic_handler]
fn panic() -> ! {
    panic_probe::hard_fault()
}

use bbqueue::nicknames::Churrasco;
use bbqueue::prod_cons::stream::{StreamConsumer, StreamProducer};
use bbqueue::traits::bbqhdl::BbqHandle;
use core::sync::atomic::{AtomicBool, AtomicI32, AtomicU32, AtomicUsize, Ordering};
use cortex_m_rt::entry;
use defmt;
use defmt::debug;
use defmt_rtt as _;
use hal::Pin;
use hal::Syscon;
use hal::drivers::{Timer, UsbBus, pins, pins::direction::Output};
use hal::prelude::*;
use hal::raw as pac;
use hal::time::Hertz;
use hal::typestates::pin::state::Gpio;
use lpc55_hal as hal;
use pac::interrupt;
use usb_device::{
    bus::{self},
    device::{StringDescriptors, UsbDeviceBuilder, UsbVidPid},
    endpoint::IsochronousSynchronizationType,
};
use usbd_uac2::UsbIsochronousFeedback;
use usbd_uac2::{
    self, AudioClassConfig, RangeEntry, TerminalConfig, UsbAudioClass, UsbAudioClockImpl, UsbSpeed,
    constants::{FunctionCode, TerminalType},
    descriptors::{ChannelConfig, ClockType, FormatType1, LockDelay},
};

use crate::dac::DacImpl;
use crate::traits::Dac;

#[cfg(feature = "ak4490")]
pub mod dac {
    mod ak4490;
    pub use self::ak4490::Ak4490Dac as DacImpl;
}

mod traits;

// Fo = M/(N*2*P) * Fin
// Fo = 3072/(125*2*8) * 16MHz = 24.576MHz
//
const FIFO_LENGTH: usize = 256; // frames
const MCLK_FREQ: u32 = 24576000;
const SAMPLE_RATE: u32 = 88200;
type SampleType = (i32, i32);

struct CodecPins {
    reset: Pin<pins::Pio0_3, Gpio<Output>>,
}

struct ClockSelPins {
    sel_24m: Pin<pins::Pio0_27, Gpio<Output>>,
    sel_22m: Pin<pins::Pio0_31, Gpio<Output>>,
}

struct Clock {
    pins: ClockSelPins,
    cur_rate: u32,
}
impl Clock {
    const RATES: [RangeEntry<u32>; 1] = [RangeEntry::new_fixed(SAMPLE_RATE)];
}
impl UsbAudioClockImpl for Clock {
    const CLOCK_TYPE: usbd_uac2::descriptors::ClockType = ClockType::InternalFixed;
    const SOF_SYNC: bool = false;
    fn get_sample_rate(&self) -> u32 {
        self.cur_rate
    }
    fn set_sample_rate(
        &mut self,
        sample_rate: u32,
    ) -> core::result::Result<(), usbd_uac2::UsbAudioClassError> {
        if 24_576_000u32.is_multiple_of(sample_rate) {
            defmt::info!("[clock] 24M clock selected");
            self.pins.sel_22m.set_low().ok();
            // hal::wait_at_least(1);
            self.pins.sel_24m.set_high().ok();
        } else {
            defmt::info!("[clock] 22M clock selected");
            self.pins.sel_24m.set_low().ok();
            // hal::wait_at_least(1);
            self.pins.sel_22m.set_high().ok();
        };
        self.cur_rate = sample_rate;
        Ok(())
    }
    fn get_rates(
        &self,
    ) -> core::result::Result<&[usbd_uac2::RangeEntry<u32>], usbd_uac2::UsbAudioClassError> {
        Ok(&Clock::RATES)
    }
    fn get_clock_validity(&self) -> core::result::Result<bool, usbd_uac2::UsbAudioClassError> {
        Ok(true)
    }
    fn alt_setting(
        &mut self,
        alt_setting: u8,
    ) -> core::result::Result<(), usbd_uac2::UsbAudioClassError> {
        match alt_setting {
            0 => {
                self.pins.sel_22m.set_low().ok();
                self.pins.sel_24m.set_low().ok();
            }
            1 => {
                self.set_sample_rate(self.cur_rate).ok();
            }
            _ => {}
        };
        Ok(())
    }
}

#[derive(Default)]
struct PerfCounters {
    received_frames: AtomicUsize,
    played_frames: AtomicUsize,
    min_fill: AtomicUsize,
    avg_fill: AtomicUsize,
    queue_underflows: AtomicUsize,
    queue_overflows: AtomicUsize,
    audio_underflows: AtomicUsize,
}

impl PerfCounters {
    fn reset(&self) {
        self.received_frames.store(0, Ordering::Relaxed);
        self.played_frames.store(0, Ordering::Relaxed);
        self.min_fill.store(FIFO_LENGTH, Ordering::Relaxed);
        self.avg_fill.store(FIFO_LENGTH / 2, Ordering::Relaxed);
        self.queue_underflows.store(0, Ordering::Relaxed);
        self.queue_overflows.store(0, Ordering::Relaxed);
        self.audio_underflows.store(0, Ordering::Relaxed);
    }
}

impl defmt::Format for PerfCounters {
    fn format(&self, fmt: defmt::Formatter) {
        defmt::write!(
            fmt,
            "frames: {}/{} min_fill: {} avg fill: {} a_underflows: {} q_underflows: {} q_overflows: {}",
            self.played_frames.load(Ordering::Relaxed),
            self.received_frames.load(Ordering::Relaxed),
            self.min_fill.load(Ordering::Relaxed),
            self.avg_fill.load(Ordering::Relaxed),
            self.audio_underflows.load(Ordering::Relaxed),
            self.queue_underflows.load(Ordering::Relaxed),
            self.queue_overflows.load(Ordering::Relaxed)
        )
    }
}

const BYTES_PER_FRAME: usize = 8;
const QUEUE_BYTES: usize = FIFO_LENGTH * BYTES_PER_FRAME;
// We use bbqueue here for performance in the USB driver that runs almost entirely in interrupt free critical section.
static QUEUE: Churrasco<QUEUE_BYTES> = Churrasco::new();
// Used for feedback calculation of current fifo state
static PRODUCED: AtomicU32 = AtomicU32::new(0);
static CONSUMED: AtomicU32 = AtomicU32::new(0);

static PERF: PerfCounters = PerfCounters {
    received_frames: AtomicUsize::new(0),
    played_frames: AtomicUsize::new(0),
    min_fill: AtomicUsize::new(FIFO_LENGTH), // not recording this for now, need to figure out how to make it meaningful, since the queue starts empty
    avg_fill: AtomicUsize::new(FIFO_LENGTH / 2),
    queue_underflows: AtomicUsize::new(0), // ditto here, since we underflow at startup, but we record this one as it can be trended
    queue_overflows: AtomicUsize::new(0),
    audio_underflows: AtomicUsize::new(0),
};

fn cur_fill() -> u32 {
    PRODUCED
        .load(Ordering::Acquire)
        .wrapping_sub(CONSUMED.load(Ordering::Acquire))
        / BYTES_PER_FRAME as u32
}

#[inline]
fn try_write_one_frame<T: BbqHandle>(
    cons: &mut StreamConsumer<T>,
    i2s: &pac::i2s7::RegisterBlock,
) -> bool {
    if let Ok(rgr) = cons.read() {
        if rgr.len() >= BYTES_PER_FRAME {
            let l = u32::from_le_bytes(rgr[0..4].try_into().unwrap());
            let r = u32::from_le_bytes(rgr[4..8].try_into().unwrap());

            i2s.fifowr.write(|w| unsafe { w.bits(l) });
            i2s.fifowr.write(|w| unsafe { w.bits(r) });

            // consume exactly one frame (8 bytes)
            rgr.release(BYTES_PER_FRAME);
            PERF.played_frames.fetch_add(1, Ordering::Relaxed);
            CONSUMED.fetch_add(BYTES_PER_FRAME as u32, Ordering::Relaxed);
            return true;
        } else {
            // Not enough bytes for a full frame: leave it in the queue.
            return false;
        }
    }
    false
}

#[interrupt]
fn FLEXCOMM7() {
    let i2s = unsafe { &*pac::I2S7::ptr() };

    if i2s.fifostat.read().txlvl().bits() == 0 {
        PERF.audio_underflows.fetch_add(1, Ordering::Relaxed);
    }

    // refil the buffer to 4 frames / 8 samples
    let mut cons = QUEUE.stream_consumer();
    while i2s.fifostat.read().txlvl().bits() <= 6 {
        if !try_write_one_frame(&mut cons, i2s) {
            // No complete frame available: write silence to keep FIFO above threshold
            PERF.queue_underflows.fetch_add(1, Ordering::Relaxed);
            i2s.fifowr.write(|w| unsafe { w.bits(0) });
            i2s.fifowr.write(|w| unsafe { w.bits(0) });
            break;
        }
    }
}

struct Audio<T: BbqHandle, D: Dac<I>, I> {
    running: AtomicBool,
    i2s: I2sTx,
    dac: D,
    producer: StreamProducer<T>,
    integrator: AtomicI32,
    filtered_fill: AtomicI32,
    _marker: core::marker::PhantomData<I>,
}
impl<T: BbqHandle, D: Dac<I>, I> Audio<T, D, I> {
    fn init(&mut self) {
        self.dac.init();
        self.dac.change_rate(SAMPLE_RATE);
    }
    fn start(&self) {
        self.running.store(true, Ordering::Relaxed);
        defmt::info!("playback starting, enabling interrupts");
        self.i2s.i2s.fifointenclr.write(|w| w.txlvl().set_bit());
        // FIFO threshold trigger enable
        self.i2s
            .i2s
            .fifotrig
            .modify(|_, w| unsafe { w.txlvl().bits(6).txlvlena().enabled() });
        // FIFO level interrupt enable
        self.i2s.i2s.fifointenset.modify(|_, w| w.txlvl().enabled());
        unsafe { pac::NVIC::unmask(pac::Interrupt::FLEXCOMM7) };
    }
    fn stop(&self) {
        self.running.store(true, Ordering::Relaxed);
        defmt::info!("playback stopped: {}", PERF);
        PERF.reset();
        pac::NVIC::mask(pac::Interrupt::FLEXCOMM7);
    }
}
impl<T: BbqHandle, D: Dac<I>, I, B: bus::UsbBus> UsbAudioClass<'_, B> for Audio<T, D, I> {
    fn alternate_setting_changed(&mut self, _terminal: usb_device::UsbDirection, alt_setting: u8) {
        match alt_setting {
            0 => self.stop(),
            1 => self.start(),
            _ => defmt::error!("unexpected alt setting {}", alt_setting),
        }
    }
    fn audio_data_rx(
        &mut self,
        ep: &usb_device::endpoint::Endpoint<'_, B, usb_device::endpoint::Out>,
    ) {
        let mut buf = [0; SAMPLE_RATE as usize / 1000 * 64];
        let len = match ep.read(&mut buf) {
            Ok(len) => len,
            Err(e) => {
                defmt::error!("usb error in rx callback");
                return;
            }
        };
        let buf = &buf[..len];
        if let Ok(mut wg) = self.producer.grant_exact(buf.len()) {
            wg.copy_from_slice(buf);
            wg.commit(buf.len());
            PRODUCED.fetch_add(buf.len() as u32, Ordering::Relaxed);
            PERF.received_frames
                .fetch_add(buf.len() / BYTES_PER_FRAME, Ordering::Relaxed);
        } else {
            PERF.queue_overflows.fetch_add(1, Ordering::Relaxed);
            defmt::error!("overflowed bbq, asked {}", buf.len());
        }
    }
    fn feedback(&mut self, nominal_rate: UsbIsochronousFeedback) -> Option<UsbIsochronousFeedback> {
        let target = FIFO_LENGTH as i32 / 2 - nominal_rate.int as i32;

        let fill = cur_fill() as i32;
        let prev = self.filtered_fill.load(Ordering::Relaxed);
        let filtered = prev + ((fill - prev) >> 4); // ~1/16 smoothing
        self.filtered_fill.store(filtered, Ordering::Relaxed);

        let error = filtered - target;

        // Clamp startup excursions.
        let error = error.clamp(-(nominal_rate.int as i32 * 4), nominal_rate.int as i32 * 4);

        // Reset integrator when the error is small
        if error.abs() < 2 {
            self.integrator.store(0, Ordering::Relaxed);
        }
        let mut integrator = self.integrator.load(Ordering::Relaxed);
        integrator = integrator - (integrator >> 6); // ~1/64 leak, reduce windup

        integrator = integrator.clamp(-256, 256);
        self.integrator.store(integrator, Ordering::Relaxed);

        // gains
        let p = error << 3;
        let i = integrator * 0; // disabled for now

        let correction = -((p + i) >> 2);
        let nominal_v = nominal_rate.to_u32_12_13() as i32;

        let mut v = nominal_v + correction;

        // Tight clamp around nominal.
        v = v.clamp(nominal_v - (1 << 12), nominal_v + (1 << 12));

        defmt::debug!(
            "fill:{} err:{} int:{} fb:{=u32:x}",
            fill,
            error,
            integrator,
            v
        );

        Some(UsbIsochronousFeedback::new(v as u32))
    }
}

pub struct I2sTx {
    pub i2s: pac::I2S7,
}

pub fn init_i2s(mut fc7: pac::FLEXCOMM7, i2s7: pac::I2S7, syscon: &mut Syscon) -> I2sTx {
    defmt::debug!("init i2s");
    // Enable BOTH
    syscon.reset(&mut fc7);
    syscon.enable_clock(&mut fc7);

    unsafe {
        pac::IOCON::ptr().as_ref().unwrap().pio0_23.modify(|_, w| {
            w.func()
                .alt1() // MCLK
                .mode()
                .inactive()
                .slew()
                .fast()
                .invert()
                .disabled()
                .digimode()
                .digital()
                .od()
                .normal()
        });
        pac::SYSCON::ptr()
            .as_ref()
            .unwrap()
            .mclkio
            .modify(|_, w| w.mclkio().input());
        pac::SYSCON::ptr()
            .as_ref()
            .unwrap()
            .fcclksel7()
            .modify(|_, w| w.sel().enum_0x5()); // MCLK
    };

    // Select I2S TX function
    fc7.pselid.write(|w| w.persel().i2s_transmit());

    let regs = i2s7;

    // Enable TX FIFO only
    regs.fifocfg.modify(|_, w| {
        w.enabletx()
            .enabled()
            .enablerx()
            .disabled()
            .dmatx()
            .disabled()
            .txi2se0()
            .zero()
    });

    // Flush
    regs.fifocfg.modify(|_, w| w.emptytx().set_bit());

    regs.cfg2
        .modify(|_, w| unsafe { w.position().bits(0).framelen().bits(63) }); // framelen = 64

    let bclk_div = (MCLK_FREQ / SAMPLE_RATE / 64) as u16;
    regs.div
        .modify(|_, w| unsafe { w.div().bits(bclk_div - 1) }); // Clock source is MCLK (12.288MHz) / 4 = 3MHz

    // Config
    regs.cfg1.modify(|_, w| unsafe {
        w.mstslvcfg()
            .normal_master()
            .onechannel()
            .dual_channel()
            .datalen()
            .bits(31)
            .mainenable()
            .enabled()
            .mode()
            .classic_mode()
            .datapause()
            .normal()
    });

    I2sTx { i2s: regs }
}

#[entry]
fn main() -> ! {
    let hal = hal::new();

    let mut anactrl = hal.anactrl;
    let mut pmc = hal.pmc;
    let mut syscon = hal.syscon;
    let mut gpio = hal.gpio.enabled(&mut syscon);
    let mut iocon = hal.iocon.enabled(&mut syscon);

    debug!("start");

    debug!("iocon");
    let usb0_vbus_pin = pins::Pio0_22::take()
        .unwrap()
        .into_usb0_vbus_pin(&mut iocon);
    let codec_i2c_pins = (
        pins::Pio0_16::take().unwrap().into_i2c4_scl_pin(&mut iocon),
        pins::Pio0_5::take().unwrap().into_i2c4_sda_pin(&mut iocon),
    );
    let codec_i2s_pins = (
        pins::Pio0_21::take().unwrap().into_spi7_sck_pin(&mut iocon),
        pins::Pio0_20::take().unwrap().into_i2s7_sda_pin(&mut iocon),
        pins::Pio0_19::take().unwrap().into_i2s7_ws_pin(&mut iocon),
        pins::Pio0_23::take().unwrap(),
    );
    let codec_gpio_pins = CodecPins {
        reset: pins::Pio0_3::take()
            .unwrap()
            .into_gpio_pin(&mut iocon, &mut gpio)
            .into_output_low(),
    };
    let clock_sel_pins = ClockSelPins {
        sel_24m: pins::Pio0_27::take()
            .unwrap()
            .into_gpio_pin(&mut iocon, &mut gpio)
            .into_output_low(),
        sel_22m: pins::Pio0_31::take()
            .unwrap()
            .into_gpio_pin(&mut iocon, &mut gpio)
            .into_output_low(),
    };
    let leds = (
        pins::Pio0_13::take()
            .unwrap()
            .into_gpio_pin(&mut iocon, &mut gpio)
            .into_output_low(),
        pins::Pio0_14::take()
            .unwrap()
            .into_gpio_pin(&mut iocon, &mut gpio)
            .into_output_low(),
    );

    // iocon.disabled(&mut syscon).release(); // save the environment :)

    debug!("clocks");
    let clocks = hal::ClockRequirements::default()
        .system_frequency(96.MHz())
        .configure(&mut anactrl, &mut pmc, &mut syscon)
        .unwrap();
    let mut delay_timer = Timer::new(
        hal.ctimer
            .0
            .enabled(&mut syscon, clocks.support_1mhz_fro_token().unwrap()),
    );

    debug!("peripherals");

    let i2c_peripheral = hal
        .flexcomm
        .4
        .enabled_as_i2c(&mut syscon, &clocks.support_flexcomm_token().unwrap());

    let i2c_bus = I2cMaster::new(
        i2c_peripheral,
        codec_i2c_pins,
        Hertz::try_from(400.kHz()).unwrap(),
    );
    let dac_impl = DacImpl::new(i2c_bus, codec_gpio_pins);

    let i2s_peripheral = {
        let fc7 = hal.flexcomm.7.release();
        init_i2s(fc7.0, fc7.2, &mut syscon)
    };

    let usb_peripheral = hal.usbhs.enabled_as_device(
        &mut anactrl,
        &mut pmc,
        &mut syscon,
        &mut delay_timer,
        clocks.support_usbhs_token().unwrap(),
    );

    defmt::info!("audio init");
    let mut audio = Audio {
        i2s: i2s_peripheral,
        dac: dac_impl,
        producer: QUEUE.stream_producer(),
        running: AtomicBool::new(false),
        integrator: AtomicI32::new(0),
        filtered_fill: AtomicI32::new(0),
        _marker: core::marker::PhantomData,
    };
    audio.init();

    let usb_bus = UsbBus::new(usb_peripheral, usb0_vbus_pin);
    let mut clock = Clock {
        pins: clock_sel_pins,
        cur_rate: SAMPLE_RATE,
    };

    let mut usb_dev = UsbDeviceBuilder::new(&usb_bus, UsbVidPid(0x1209, 0xcc1d))
        .composite_with_iads()
        .strings(&[StringDescriptors::default()
            .manufacturer("VE7XEN")
            .product("Guac Tortilla")
            .serial_number("123456789")])
        .unwrap()
        .max_packet_size_0(64)
        .unwrap()
        .device_class(0xef)
        .device_sub_class(0x02)
        .device_protocol(0x01)
        .build();

    let config = AudioClassConfig::new(UsbSpeed::High, FunctionCode::Other, &mut clock, &mut audio)
        .with_output_config(TerminalConfig::new(
            2,
            1,
            2,
            FormatType1 {
                bit_resolution: 32,
                bytes_per_sample: 4,
            },
            TerminalType::OutHeadphones,
            ChannelConfig::default_chans(2),
            IsochronousSynchronizationType::Asynchronous,
            LockDelay::Undefined(0),
            None,
        ));

    let mut uac2 = config.build(&usb_bus).unwrap();


    defmt::info!("main loop");

    loop {
        usb_dev.poll(&mut [&mut uac2]);
    }
}