//! PM_K-1 firmware — Stage 3 bring-up.
//! Milestone 1 (done): blink GP25 → proved boot/flash.
//! Milestone 2 (this): init the ST7796 320×480 display over SPI0 and draw to it.
//! Pins (from the CircuitPython firmware): SCK=GP2, MOSI=GP3, CS=GP5, DC=GP6, RST=GP7;
//! BGR panel, colours inverted. LED on GP25 keeps blinking as a heartbeat.

#![no_std]
#![no_main]

extern crate alloc;
use alloc::vec::Vec;
use embedded_alloc::LlffHeap as Heap;

#[global_allocator]
static HEAP: Heap = Heap::empty();

use embedded_hal::delay::DelayNs;
use embedded_hal::digital::{InputPin, OutputPin};
use embedded_hal::pwm::SetDutyCycle;
use embedded_hal_0_2::adc::OneShot;
use embedded_hal::spi::SpiBus;
use defmt::info;
use defmt_rtt as _; // global defmt logger over RTT (read by `probe-rs run`)
use panic_probe as _; // panic handler that prints the message + location over defmt, then halts
use rp235x_hal as hal;
use hal::fugit::RateExtU32;
use hal::Clock;
use embedded_graphics::pixelcolor::Rgb565;
use embedded_graphics::prelude::*;
use embedded_graphics::primitives::Rectangle;

/// Image definition block — the RP2350 bootrom looks for this to boot the image.
#[link_section = ".start_block"]
#[used]
pub static IMAGE_DEF: hal::block::ImageDef = hal::block::ImageDef::secure_exe();

const XTAL_FREQ_HZ: u32 = 12_000_000;
const WIDTH: u16 = 320;
const HEIGHT: u16 = 480;

/// Dirty-tile grid for incremental updates. TILE wide is sub-screen-width so each tile blit writes
/// faster than the panel scan crosses it → minimal tearing. 320/40=8 cols, 480/40=12 rows = 96 tiles.
const TILE: u16 = 40;
const NX: usize = WIDTH as usize / TILE as usize;
const NY: usize = HEIGHT as usize / TILE as usize;

/// Direct ST7796 driver — a faithful port of the proven MicroPython driver in `pico/main.py`,
/// which runs this exact panel on this exact board. Per-command CS framing (CS low → command →
/// data → CS high), MADCTL 0x48, INVON, NO row/column offset, big-endian RGB565. This replaces
/// mipidsi, whose split-transaction CS and orientation/offset math were mangling the geometry.
/// We render into the RAM `FrameBuf` (embedded-graphics), then `blit` the whole frame in one
/// windowed RAMWR stream — exactly like main.py's `_window` + `spi.write`.
struct St7796<SPI, DC, CS, RST> {
    spi: SPI,
    dc: DC,
    cs: CS,
    rst: RST,
}
impl<SPI, DC, CS, RST> St7796<SPI, DC, CS, RST>
where
    SPI: SpiBus,
    DC: OutputPin,
    CS: OutputPin,
    RST: OutputPin,
{
    fn cmd(&mut self, c: u8, data: &[u8]) {
        let _ = self.cs.set_low();
        let _ = self.dc.set_low();
        let _ = self.spi.write(&[c]);
        if !data.is_empty() {
            let _ = self.dc.set_high();
            let _ = self.spi.write(data);
        }
        let _ = self.cs.set_high();
    }

    /// main.py reset() + init() verbatim.
    fn init(&mut self, delay: &mut impl DelayNs) {
        let _ = self.cs.set_high();
        let _ = self.dc.set_low();
        let _ = self.rst.set_high();
        delay.delay_ms(20);
        let _ = self.rst.set_low();
        delay.delay_ms(40);
        let _ = self.rst.set_high();
        delay.delay_ms(150);

        self.cmd(0x01, &[]); // SWRESET
        delay.delay_ms(120);
        self.cmd(0x11, &[]); // SLPOUT
        delay.delay_ms(120);
        self.cmd(0xF0, &[0xC3]); // unlock extension command set
        self.cmd(0xF0, &[0x96]);
        self.cmd(0x36, &[0x48]); // MADCTL = MX | BGR
        self.cmd(0x3A, &[0x55]); // COLMOD 16bpp
        self.cmd(0xB4, &[0x01]);
        self.cmd(0xB6, &[0x80, 0x02, 0x3B]);
        self.cmd(0xE8, &[0x40, 0x8A, 0x00, 0x00, 0x29, 0x19, 0xA5, 0x33]);
        self.cmd(0xC1, &[0x06]);
        self.cmd(0xC2, &[0xA7]);
        self.cmd(0xC5, &[0x18]);
        delay.delay_ms(120);
        self.cmd(0xE0, &[0xF0, 0x09, 0x0B, 0x06, 0x04, 0x15, 0x2F, 0x54, 0x42, 0x3C, 0x17, 0x14, 0x18, 0x1B]);
        self.cmd(0xE1, &[0xE0, 0x09, 0x0B, 0x06, 0x04, 0x03, 0x2B, 0x43, 0x42, 0x3B, 0x16, 0x14, 0x17, 0x1B]);
        self.cmd(0xF0, &[0x3C]); // lock
        self.cmd(0xF0, &[0x69]);
        delay.delay_ms(120);
        self.cmd(0x21, &[]); // INVON (INVERT_COLORS = true)
        self.cmd(0x29, &[]); // DISPON
        delay.delay_ms(50);
    }

    /// Blit a full-width band of rows `ylo..ylo+h`. CASET is the full width (0..=319, exactly like
    /// the full-screen blit), only RASET is partial — and because a full-width band is *contiguous*
    /// in the row-major framebuffer, `px` is just that contiguous slice, streamed with the identical
    /// loop the full blit uses. The full-screen blit is `blit_band(&px[..], 0, HEIGHT)`.
    fn blit_band(&mut self, px: &[Rgb565], ylo: u16, h: u16) {
        let y1 = ylo + h - 1;
        self.cmd(0x2A, &[0x00, 0x00, 0x01, 0x3F]); // CASET 0..=319 (full width)
        self.cmd(0x2B, &[(ylo >> 8) as u8, (ylo & 0xFF) as u8, (y1 >> 8) as u8, (y1 & 0xFF) as u8]);
        let _ = self.cs.set_low();
        let _ = self.dc.set_low();
        let _ = self.spi.write(&[0x2C]); // RAMWR
        let _ = self.dc.set_high();
        let mut buf = [0u8; 1024];
        let mut n = 0;
        for &c in px {
            let raw = c.into_storage(); // u16 RGB565, big-endian like main.py
            buf[n] = (raw >> 8) as u8;
            buf[n + 1] = (raw & 0xFF) as u8;
            n += 2;
            if n == buf.len() {
                let _ = self.spi.write(&buf);
                n = 0;
            }
        }
        if n > 0 {
            let _ = self.spi.write(&buf[..n]);
        }
        let _ = self.cs.set_high();
    }

    /// Full-screen blit — the known-good path. Just a full-height band.
    fn blit(&mut self, px: &[Rgb565]) {
        self.blit_band(px, 0, HEIGHT);
    }

    /// Blit a sub-rectangle `(x,y,w,h)` gathered from the full framebuffer. Used for small tile
    /// updates: sub-width writes finish a region faster than the panel scan crosses it (datasheet
    /// §10.8 Example 1), so the visible tear is tiny vs a full-width band. Same streaming loop and
    /// 1024-byte buffer as `blit_band` (the 512-byte buffer was what broke the earlier version).
    fn blit_rect(&mut self, fb: &[Rgb565], x: u16, y: u16, w: u16, h: u16) {
        let x1 = x + w - 1;
        let y1 = y + h - 1;
        self.cmd(0x2A, &[(x >> 8) as u8, (x & 0xFF) as u8, (x1 >> 8) as u8, (x1 & 0xFF) as u8]);
        self.cmd(0x2B, &[(y >> 8) as u8, (y & 0xFF) as u8, (y1 >> 8) as u8, (y1 & 0xFF) as u8]);
        let _ = self.cs.set_low();
        let _ = self.dc.set_low();
        let _ = self.spi.write(&[0x2C]); // RAMWR
        let _ = self.dc.set_high();
        let mut buf = [0u8; 1024];
        let mut n = 0;
        for r in 0..h as usize {
            let row = (y as usize + r) * WIDTH as usize + x as usize;
            for c in 0..w as usize {
                let raw = fb[row + c].into_storage();
                buf[n] = (raw >> 8) as u8;
                buf[n + 1] = (raw & 0xFF) as u8;
                n += 2;
                if n == buf.len() {
                    let _ = self.spi.write(&buf);
                    n = 0;
                }
            }
        }
        if n > 0 {
            let _ = self.spi.write(&buf[..n]);
        }
        let _ = self.cs.set_high();
    }
}


/// Off-screen framebuffer: a full 320×480 Rgb565 image in RAM (300 KB). The pm-ui draw functions
/// render into THIS (instant, no SPI), then we push the whole thing to the panel in one
/// `fill_contiguous` sweep — so the panel never shows a clear-then-redraw transition → no flicker.
const FB_LEN: usize = WIDTH as usize * HEIGHT as usize;
static mut FRAMEBUF: [Rgb565; FB_LEN] = [Rgb565::new(0, 0, 0); FB_LEN];

struct FrameBuf {
    px: &'static mut [Rgb565],
}
impl DrawTarget for FrameBuf {
    type Color = Rgb565;
    type Error = core::convert::Infallible;
    fn draw_iter<I>(&mut self, pixels: I) -> Result<(), Self::Error>
    where
        I: IntoIterator<Item = Pixel<Self::Color>>,
    {
        for Pixel(p, c) in pixels {
            if p.x >= 0 && p.y >= 0 && (p.x as u16) < WIDTH && (p.y as u16) < HEIGHT {
                self.px[p.y as usize * WIDTH as usize + p.x as usize] = c;
            }
        }
        Ok(())
    }
    fn fill_contiguous<I>(&mut self, area: &Rectangle, colors: I) -> Result<(), Self::Error>
    where
        I: IntoIterator<Item = Self::Color>,
    {
        // fast path used by clear()/fill — write straight into the backing slice row by row
        let mut it = colors.into_iter();
        let x0 = area.top_left.x.max(0) as u16;
        let y0 = area.top_left.y.max(0) as u16;
        let x1 = (area.top_left.x + area.size.width as i32).min(WIDTH as i32) as u16;
        let y1 = (area.top_left.y + area.size.height as i32).min(HEIGHT as i32) as u16;
        for y in y0..y1 {
            for x in x0..x1 {
                if let Some(c) = it.next() {
                    self.px[y as usize * WIDTH as usize + x as usize] = c;
                }
            }
        }
        Ok(())
    }
}
impl OriginDimensions for FrameBuf {
    fn size(&self) -> Size {
        Size::new(WIDTH as u32, HEIGHT as u32)
    }
}

#[hal::entry]
fn main() -> ! {
    // heap for track parsing (track-format uses alloc). RP2350 has 520 KB SRAM.
    const HEAP_SIZE: usize = 16 * 1024;
    {
        use core::mem::MaybeUninit;
        static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
        unsafe { HEAP.init(core::ptr::addr_of_mut!(HEAP_MEM) as usize, HEAP_SIZE) }
    }
    info!("== pm-kit boot == heap {}KB, {} free", HEAP_SIZE / 1024, HEAP.free());

    let mut pac = hal::pac::Peripherals::take().unwrap();
    let mut watchdog = hal::Watchdog::new(pac.WATCHDOG);
    let clocks = hal::clocks::init_clocks_and_plls(
        XTAL_FREQ_HZ,
        pac.XOSC,
        pac.CLOCKS,
        pac.PLL_SYS,
        pac.PLL_USB,
        &mut pac.RESETS,
        &mut watchdog,
    )
    .ok()
    .unwrap();

    let mut timer = hal::Timer::new_timer0(pac.TIMER0, &mut pac.RESETS, &clocks);
    let sio = hal::Sio::new(pac.SIO);
    let pins = hal::gpio::Pins::new(pac.IO_BANK0, pac.PADS_BANK0, sio.gpio_bank0, &mut pac.RESETS);

    let mut led = pins.gpio25.into_push_pull_output();
    led.set_high().unwrap(); // solid ON during init: if it stays solid → hung in init; slow blink → reached the loop

    // --- ST7796 over SPI0 (direct driver, ported from pico/main.py) ---
    let sclk = pins.gpio2.into_function::<hal::gpio::FunctionSpi>();
    let mosi = pins.gpio3.into_function::<hal::gpio::FunctionSpi>();
    let dc = pins.gpio6.into_push_pull_output();
    let rst = pins.gpio7.into_push_pull_output();
    let cs = pins.gpio5.into_push_pull_output();

    let spi = hal::spi::Spi::<_, _, _, 8>::new(pac.SPI0, (mosi, sclk));
    let spi = spi.init(
        &mut pac.RESETS,
        clocks.peripheral_clock.freq(),
        62500.kHz(), // pico-cp SPI_BAUD: 62.5 MHz — "faster SPI = smaller tearing window" (drop to 40 MHz if unstable)
        embedded_hal::spi::MODE_0,
    );

    let mut st = St7796 { spi, dc, cs, rst };
    st.init(&mut timer);
    info!("display init ok (ported st7796 driver), {} free", HEAP.free());

    // ---- inputs + speaker ----
    let mut btn_a = pins.gpio15.into_pull_up_input(); // A = play/stop
    let mut btn_b = pins.gpio14.into_pull_up_input(); // B = grid/notation view
    let mut adc = hal::adc::Adc::new(pac.ADC, &mut pac.RESETS);
    let mut joy_x = hal::adc::AdcPin::new(pins.gpio26).unwrap();
    let mut joy_y = hal::adc::AdcPin::new(pins.gpio27).unwrap();
    let pwm_slices = hal::pwm::Slices::new(pac.PWM, &mut pac.RESETS);
    let mut spk = pwm_slices.pwm6; // GP13 = PWM slice 6, channel B
    spk.set_div_int(125);
    spk.set_top(600); // ~2 kHz click
    spk.enable();
    spk.channel_b.output_to(pins.gpio13);
    info!("scaffolding (buttons/adc/pwm) up, {} free", HEAP.free());

    // ---- built-in grooves ----
    const GROOVES: [&str; 4] = [
        "t120;kick:4=X.x.;snare:4=.X.X;hatClosed:4/2=xxxxxxxx",
        "t92;kick:4/3=X....x...x..;snare:4/3=..gg.gX.gg.g;hatClosed:4/3=X.xX.xX.xX.x",
        "t140;kick:4=X..x;snare:4=.X.X;hatClosed:4/4=xxxxxxxxxxxxxxxx",
        "t100;kick:4=Xxxx;claves:5=Xxxxx~",
    ];
    const NAMES: [&str; 4] = ["Four on the floor", "Half-time shuffle", "Driving 16ths", "5 over 4"];

    let mut idx = 0usize;
    let mut track = track_format::parse(GROOVES[idx]);
    info!("parsed groove 0: bpm={} lanes={}, {} free", track.bpm, track.lanes.len(), HEAP.free());
    let mut tempo: i64 = track.bpm;
    let mut playing = true;
    // View cycle on button-B: Grid → Staff → TUBS → Konnakol → Grid …
    #[derive(Clone, Copy, PartialEq)]
    enum View {
        Grid,
        Staff,
        Tubs,
        Konnakol,
    }
    let mut view = View::Grid;

    // double-buffer: draw into the RAM framebuffer, then push only the full-width ROW BANDS that
    // changed vs the last frame. Full-width windows (CASET 0..319) behave exactly like the proven
    // full-screen blit — unlike sub-width tiles, which the panel's MX bit mirror-maps wrongly.
    let mut fb = FrameBuf { px: unsafe { &mut *core::ptr::addr_of_mut!(FRAMEBUF) } };
    let mut tile_hash = [0u64; NX * NY]; // FNV-1a of each tile at its last blit (0 → force first paint)

    let mut bar_start = timer.get_counter().ticks();
    let mut last_step = usize::MAX; // audio: click on a new step
    let mut drawn_step = usize::MAX; // draw: advance the playhead
    let mut click_off_us: u64 = 0;
    let (mut pa, mut pb) = (false, false);
    let mut joy_zone = 0i8;
    let mut dirty = true; // force the first frame
    let mut force_full = true; // first paint + big changes use the known-good full blit
    let mut hb = false;
    let mut hb_us = 0u64;
    led.set_low().unwrap();
    info!("entering metronome loop (double-buffered)");

    loop {
        let now = timer.get_counter().ticks();

        // ---- inputs ----
        let a = btn_a.is_low().unwrap_or(false);
        let b = btn_b.is_low().unwrap_or(false);
        if a && !pa {
            playing = !playing;
            if playing {
                bar_start = now;
                last_step = usize::MAX;
            }
            dirty = true;
        }
        if b && !pb {
            view = match view {
                View::Grid => View::Staff,
                View::Staff => View::Tubs,
                View::Tubs => View::Konnakol,
                View::Konnakol => View::Grid,
            };
            dirty = true;
            force_full = true; // whole screen changes — use the clean full blit
        }
        pa = a;
        pb = b;

        let jx = adc.read(&mut joy_x).unwrap_or(2048);
        let jy = adc.read(&mut joy_y).unwrap_or(2048);
        // rotate the joystick reading 90° CCW for control
        let rx = jy as i32;
        let ry = 4095 - jx as i32;
        let zone: i8 = if ry > 3200 { 1 } else if ry < 900 { 2 } else if rx > 3200 { 3 } else if rx < 900 { 4 } else { 0 };
        if zone != 0 && joy_zone == 0 {
            match zone {
                1 => tempo = (tempo + 4).min(300),                 // up → tempo+
                2 => tempo = (tempo - 4).max(30),                  // down → tempo-
                3 => {
                    idx = (idx + 1) % GROOVES.len();               // right → next groove
                    track = track_format::parse(GROOVES[idx]);
                    tempo = track.bpm;
                    bar_start = now;
                    last_step = usize::MAX;
                }
                4 => {
                    idx = (idx + GROOVES.len() - 1) % GROOVES.len(); // left → prev
                    track = track_format::parse(GROOVES[idx]);
                    tempo = track.bpm;
                    bar_start = now;
                    last_step = usize::MAX;
                }
                _ => {}
            }
            dirty = true;
            if zone == 3 || zone == 4 {
                force_full = true; // groove change → new lanes, use the clean full blit
            }
        }
        joy_zone = zone;

        // ---- clock ----
        let master = &track.lanes[0];
        let msteps = master.levels.len().max(1) as u64;
        let beats = master.groups.iter().map(|&g| g as u64).sum::<u64>().max(1);
        let bar_us = 60_000_000u64 * beats / tempo.max(1) as u64;
        let step_us = (bar_us / msteps).max(1);
        let elapsed = now.wrapping_sub(bar_start) % bar_us;
        let cur_step = (elapsed / step_us) as usize;
        let phase = elapsed as f32 / bar_us as f32;

        // ---- audio: click on a new step that has a hit ----
        if playing && cur_step != last_step {
            let lvl = master.levels[cur_step.min(msteps as usize - 1)];
            if lvl > 0 {
                let _ = spk.channel_b.set_duty_cycle(if lvl == 2 { 380 } else { 240 });
                click_off_us = now + if lvl == 2 { 28_000 } else { 14_000 };
            }
            last_step = cur_step;
        }
        if now >= click_off_us {
            let _ = spk.channel_b.set_duty_cycle(0);
        }

        // redraw when the playhead steps to a new cell (animation), or on any state change
        if playing && cur_step != drawn_step {
            dirty = true;
        }

        // ---- draw into the framebuffer, then blit ONCE (flicker-free) ----
        if dirty {
            let lanes: Vec<pm_ui::LaneView> = track
                .lanes
                .iter()
                .map(|l| pm_ui::LaneView {
                    name: &l.sound,
                    levels: &l.levels,
                    orns: &l.orns,
                    groups: &l.groups,
                    beats: l.groups.iter().map(|&g| g as u32).sum::<u32>().min(255) as u8,
                    poly: l.poly,
                    muted: l.mute,
                })
                .collect();
            let screen = pm_ui::Screen { name: NAMES[idx], bpm: tempo, playing, phase, lanes: &lanes };
            match view {
                View::Grid => pm_ui::draw_metronome(&mut fb, &screen).ok(),
                View::Staff => pm_ui::notation::draw(&mut fb, &screen, pm_ui::ViewMode::Staff).ok(),
                View::Tubs => pm_ui::notation::draw(&mut fb, &screen, pm_ui::ViewMode::Tubs).ok(),
                View::Konnakol => pm_ui::notation::draw(&mut fb, &screen, pm_ui::ViewMode::Konnakol).ok(),
            };
            // FNV-1a of one TILE×TILE block at grid cell (tx, ty)
            fn tile_fnv(px: &[Rgb565], tx: usize, ty: usize) -> u64 {
                let mut h = 0xcbf29ce484222325u64; // FNV-1a offset basis
                let x0 = tx * TILE as usize;
                let y0 = ty * TILE as usize;
                for r in 0..TILE as usize {
                    let base = (y0 + r) * WIDTH as usize + x0;
                    for c in 0..TILE as usize {
                        h ^= px[base + c].into_storage() as u64;
                        h = h.wrapping_mul(0x100000001b3);
                    }
                }
                h
            }
            if force_full {
                // known-good full-screen blit; seed the tile hashes so increments diff against it
                st.blit(&fb.px[..]);
                for ty in 0..NY {
                    for tx in 0..NX {
                        tile_hash[ty * NX + tx] = tile_fnv(&fb.px[..], tx, ty);
                    }
                }
                force_full = false;
                info!("FULL blit (seeded {} tiles)", (NX * NY) as u32);
            } else {
                // ---- dirty-tile blit: push only the changed TILE×TILE blocks (sub-width = less tearing) ----
                let mut pushed = 0u32;
                for ty in 0..NY {
                    for tx in 0..NX {
                        let h = tile_fnv(&fb.px[..], tx, ty);
                        let idx = ty * NX + tx;
                        if h != tile_hash[idx] {
                            tile_hash[idx] = h;
                            st.blit_rect(&fb.px[..], (tx * TILE as usize) as u16, (ty * TILE as usize) as u16, TILE, TILE);
                            pushed += 1;
                        }
                    }
                }
                if pushed > 0 {
                    info!("blit {} / {} tiles", pushed, (NX * NY) as u32);
                }
            }
            drawn_step = cur_step;
            dirty = false;
        }

        // heartbeat LED + log (~2 Hz)
        if now.wrapping_sub(hb_us) > 500_000 {
            hb = !hb;
            let _ = if hb { led.set_high() } else { led.set_low() };
            info!("alive: idx={} step={} playing={} free={}", idx, cur_step, playing, HEAP.free());
            hb_us = now;
        }
        timer.delay_ms(4);
    }
}

/// picotool metadata (visible via `picotool info`).
#[link_section = ".bi_entries"]
#[used]
pub static PICOTOOL_ENTRIES: [hal::binary_info::EntryAddr; 1] =
    [hal::binary_info::rp_program_name!(c"pm-kit display")];