use embassy_time::{Duration, Instant, Timer};
use embedded_graphics::{
    Drawable, draw_target::DrawTarget, image::{Image, ImageRaw}, mono_font::{MonoTextStyleBuilder, iso_8859_1::FONT_10X20}, pixelcolor::BinaryColor, prelude::Point, text::{Baseline, Text}
};
use ssd1306::{prelude::*, I2CDisplayInterface, Ssd1306, mode::BufferedGraphicsMode};
use embassy_rp::i2c::I2c;
use core::fmt::Write;

use crate::wpm::WPM_CHANNEL;

// Helper enum for all possible animation states
enum PuppyAnim {
    Sleep,
    Idle,
    Happy,
}

/// Display type encapsulating the I2C interface, its size and graphics mode
pub type Display = Ssd1306<
    I2CInterface<I2c<'static, embassy_rp::peripherals::I2C0, embassy_rp::i2c::Async>>,
    DisplaySize128x32,
    BufferedGraphicsMode<DisplaySize128x32>
>;

/// Function that initializes the display with given settings
pub fn init_display(bus: I2c<'static, embassy_rp::peripherals::I2C0, embassy_rp::i2c::Async>) -> Option<Display> {
    let interface = I2CDisplayInterface::new(bus);
    let mut display = Ssd1306::new(
        interface,
        DisplaySize128x32,
        DisplayRotation::Rotate180,
    )
    .into_buffered_graphics_mode();

    match display.init() {
        Ok(_) => Some(display),
        Err(_) => None,
    }
}

pub async fn run_display(display: &mut Option<Display>) -> Option<()> {
    if let Some(display) = display {
        let mut wpm = 0;
        let mut frame = 0;
        let mut blink_frame = 0;
        let mut last_frame_ms = Instant::now();

        loop {
            // Choose state based on wpm
            let (state, frame_delay_ms) = if wpm == 0 {
                (PuppyAnim::Sleep, 250)
            } else if wpm < 60 {
                // faster tail wagging, connected to wpm
                let delay = (3000 / wpm.max(10)).clamp(100, 300);
                (PuppyAnim::Idle, delay)
            } else {
                // even faster tail wagging!!
                let delay = (4000 / wpm).clamp(16, 100);
                (PuppyAnim::Happy, delay)
            };

            if last_frame_ms.elapsed().as_millis() >= frame_delay_ms {
                // used for determining when to blink
                blink_frame += 1;
                // switch between tail left and right
                frame = (frame + 1) % 2;
                last_frame_ms = Instant::now();
            }

            // Choose image based on current state and tail position
            let image_data = match (state, frame) {
                (PuppyAnim::Sleep, 0) => include_bytes!("../assets/sleep-1.raw"),
                (PuppyAnim::Sleep, 1) => include_bytes!("../assets/sleep-2.raw"),
                (PuppyAnim::Idle, 0) => {
                    // blink every once in a while
                    if blink_frame % 20 < 5 {
                        include_bytes!("../assets/idle-blink-1.raw")
                    } else {
                        include_bytes!("../assets/idle-normal-1.raw")
                    }
                },
                (PuppyAnim::Idle, 1) => {
                    if blink_frame % 20 < 5 {
                        include_bytes!("../assets/idle-blink-2.raw")
                    } else {
                        include_bytes!("../assets/idle-normal-2.raw")
                    }
                },
                (PuppyAnim::Happy, 0) => include_bytes!("../assets/happy-1.raw"),
                (PuppyAnim::Happy, 1) => include_bytes!("../assets/happy-2.raw"),
                _ => include_bytes!("../assets/idle-normal-1.raw"),
            };


            display.clear(BinaryColor::Off).ok()?;
        
            // Load and draw the chosen image
            let raw: ImageRaw<BinaryColor> = ImageRaw::new(image_data, 32);
            let im = Image::new(&raw, Point::new(0, 0));
            im.draw(display).ok()?;

            let text_style = MonoTextStyleBuilder::new()
                .font(&FONT_10X20)
                .text_color(BinaryColor::On)
                .build();

            match WPM_CHANNEL.try_receive() {
                Ok(current_wpm) => {
                    wpm = current_wpm;
                },
                Err(_) => {},
            }

            let mut wpm_str = heapless::String::<32>::new();
            write!(&mut wpm_str, "WPM: {}", wpm).ok();

            // Print current WPM
            Text::with_baseline(&wpm_str, Point::new(32 + 4, 6), text_style, Baseline::Top)
                .draw(display)
                .ok()?;
    
            display.flush().ok()?;

            Timer::after(Duration::from_millis(16)).await;
        }
    } else {
        None
    }
}

/// Helper function to clear the display
pub fn clear_display(
    display: &mut Option<Display>,
) -> Option<()> {
    if let Some(display) = display {
        display.clear(BinaryColor::Off).ok()?;
        display.flush().ok()?;
        Some(())
    } else {
        None
    }
}