use core::{
    arch::asm,
    mem,
};

use molybdos::{
    embassy_rp,
    embassy_rp::{
        clocks::clk_sys_freq,
        pio,
        pio::{
            Direction,
            ShiftDirection,
            StateMachine,
        },
        Peripheral,
    },
    embassy_time,
    embassy_time::{
        Duration,
        Instant,
    },
    embedded_io_async,
    fixed::prelude::ToFixed,
    genlog::defmt,
    pio::{
        InSource,
        JmpCondition,
        MovDestination,
        MovOperation,
        MovSource,
        OutDestination,
        SetDestination,
        SideSet,
    },
    pio_proc,
};

pub trait GetSm<'inner, PIO, const IDX: usize>
where
    PIO: pio::Instance,
{
    fn sm<'p>(pio: &'p mut pio::Pio<'inner, PIO>) -> &'p mut StateMachine<'inner, PIO, IDX>;
}

pub struct SmLookup<const IDX: usize>;

macro_rules! impl_getsm {
    ($idx:literal, $name:ident) => {
        impl<'inner, PIO> GetSm<'inner, PIO, $idx> for SmLookup<$idx>
        where
            PIO: embassy_rp::pio::Instance,
        {
            fn sm<'p>(
                pio: &'p mut pio::Pio<'inner, PIO>,
            ) -> &'p mut StateMachine<'inner, PIO, $idx> {
                &mut pio.$name
            }
        }
    };
}

impl_getsm!(0, sm0);
impl_getsm!(1, sm1);
impl_getsm!(2, sm2);
impl_getsm!(3, sm3);

pub type StaticI2S<PIO, const IDX: usize> = I2S<'static, 'static, PIO, IDX>;

pub struct I2S<'b, 'p, PIO, const IDX: usize>
where
    PIO: pio::Instance,
{
    sm: &'b mut StateMachine<'p, PIO, IDX>,
}

impl<'b, 'p, PIO, const IDX: usize> I2S<'b, 'p, PIO, IDX>
where
    PIO: pio::Instance,
{
    pub const DEFAULT_BUFFER_SIZE: usize = 960;
    pub const SAMPLE_RATE_HZ: usize = 48_000 * 32;

    pub fn new(
        pio: &'b mut pio::Pio<'p, PIO>,
        bclk: impl Peripheral<P = impl pio::PioPin> + 'static,
        data: impl Peripheral<P = impl pio::PioPin> + 'static,
        ws: impl Peripheral<P = impl pio::PioPin> + 'static,
    ) -> Self
    where
        SmLookup<IDX>: GetSm<'p, PIO, IDX>,
    {
        defmt::info!("start");

        let bclk = pio.common.make_pio_pin(bclk);
        let data = pio.common.make_pio_pin(data);
        let ws = pio.common.make_pio_pin(ws);

        defmt::info!("pins done");

        let program = program::<1024>(None);
        defmt::info!("made program: {}", defmt::Debug2Format(&program));

        let program = pio.common.try_load_program(&program);

        let program = match program {
            Err(e) => {
                defmt::panic!("loading pio: {}", e);
            },
            Ok(program) => program,
        };

        defmt::info!("program loaded");

        let config = Self::config(&program, &bclk, &data, &ws);

        defmt::info!("config run");

        let sm = <SmLookup<IDX> as GetSm<PIO, IDX>>::sm(pio);
        sm.set_enable(false);
        sm.set_config(&config);
        sm.clear_fifos();

        sm.set_pin_dirs(Direction::Out, &[&bclk, &ws]);
        sm.set_pin_dirs(Direction::In, &[&data]);

        Self {
            sm,
        }
    }

    pub async fn run<const BUFFER_SIZE: usize>(
        &mut self,
        dma: impl Peripheral<P = impl embassy_rp::dma::Channel>,
        mut pipe: impl embedded_io_async::Write,
    ) {
        self.sm.set_enable(true);
        defmt::info!("i2s state machine started");

        let rx = self.sm.rx();
        let mut dma = dma.into_ref();

        let mut front_buffer = [0u8; BUFFER_SIZE];
        let mut back_buffer = [0u8; BUFFER_SIZE];

        let per_sample_dur =
            core::time::Duration::from_secs(1) * BUFFER_SIZE as _ / (Self::SAMPLE_RATE_HZ as _);

        defmt::info!("ready to loop, sample_dur = {}us", per_sample_dur.as_micros());

        let per_sample_dur = embassy_time::Duration::from_nanos(per_sample_dur.as_nanos() as _);

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

            let transfer = rx.dma_pull(dma.reborrow(), &mut front_buffer);
            let now = Instant::now();
            let deadline = now + per_sample_dur;

            if let Err(e) =
                molybdos::embassy_time::with_deadline(deadline, pipe.write_all(&back_buffer)).await
            {
                defmt::error!("audio timeout");
            }

            transfer.await;
            mem::swap(&mut back_buffer, &mut front_buffer);
        }
    }

    fn config<'pin>(
        prog: &pio::LoadedProgram<'p, PIO>,
        bclk: &pio::Pin<'pin, PIO>,
        data: &pio::Pin<'pin, PIO>,
        ws: &pio::Pin<'pin, PIO>,
    ) -> pio::Config<'p, PIO>
    where
        'pin: 'p,
    {
        let clock_frequency: usize = Self::SAMPLE_RATE_HZ;
        defmt::info!("using clock freq: {}", clock_frequency);

        let mut cfg = pio::Config::default();

        cfg.use_program(&prog, &[&bclk]);
        cfg.set_in_pins(&[&data]);
        cfg.set_set_pins(&[&ws]);

        cfg.fifo_join = pio::FifoJoin::RxOnly;

        cfg.shift_in = pio::ShiftConfig {
            threshold: 32,
            direction: ShiftDirection::Left,
            auto_fill: true,
        };

        cfg.clock_divider = (clk_sys_freq() as f64 / clock_frequency as f64 / 2.).to_fixed();
        defmt::info!(
            "set clock divider: {} based on sys freq: {}Hz",
            defmt::Debug2Format(&cfg.clock_divider),
            clk_sys_freq()
        );

        cfg
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct WarmupConfig {
    clock_freq:    usize,
    warmup_period: embassy_time::Duration,
}

/// I2S pio driver program.
///
/// Pass `warmup` to have the PIO idle for a period before starting the clock.
pub fn program<const PROGRAM_SIZE: usize>(
    warmup: Option<WarmupConfig>,
) -> molybdos::pio::Program<PROGRAM_SIZE> {
    let mut asm =
        molybdos::pio::Assembler::<PROGRAM_SIZE>::new_with_side_set(SideSet::new(true, 1, false));

    asm.set(SetDestination::PINS, 1);

    // on boot, SPH0644LM4H wants 30ms warmup time with the clock running at full frequency. it's
    // nice to do this in the PIO itself because it makes the host agnostic -- it can just wait
    // until there is data available in the fifo.
    if let Some(WarmupConfig {
        clock_freq,
        warmup_period,
    }) = warmup
    {
        const SET_MAX_BITS: usize = 5;

        let warmup_period_clock_cycles = clock_freq / 1000 * warmup_period.as_millis() as usize;
        debug_assert!(warmup_period_clock_cycles <= 1 << 31);

        asm.r#in(InSource::NULL, 32);

        let sig_bits = 32 - warmup_period_clock_cycles.leading_zeros();

        for i in 0..(sig_bits / 5) {
            const MASK: usize = 0b11111;
            let shift = i * SET_MAX_BITS as u32;

            asm.set(SetDestination::X, ((warmup_period_clock_cycles >> shift) & MASK) as u8);
            asm.r#in(InSource::X, SET_MAX_BITS as _);
        }

        let rem_bits = sig_bits % SET_MAX_BITS as u32;
        if rem_bits != 0 {
            let mask = (1 << (rem_bits + 1)) - 1;

            asm.set(
                SetDestination::X,
                ((warmup_period_clock_cycles >> (sig_bits - rem_bits)) & mask) as u8,
            );
            asm.r#in(InSource::X, SET_MAX_BITS as _);
        }

        asm.mov(MovDestination::X, MovOperation::None, MovSource::ISR);

        let mut bootloop = asm.label();
        {
            asm.nop_with_side_set(1);
            asm.jmp_with_side_set(JmpCondition::XDecNonZero, &mut bootloop, 0);
        }
    }

    let mut wrap_target = asm.label();
    asm.bind(&mut wrap_target);

    asm.in_with_side_set(InSource::PINS, 1, 0);
    asm.nop_with_side_set(1);

    let mut wrap = asm.label();
    asm.bind(&mut wrap);

    defmt::info!("trying assemble...");
    asm.assemble_with_wrap(wrap, wrap_target)
}