#![no_std]
#![no_main]

use assign_resources::assign_resources;
use bincode::config::Configuration;
use bincode::Decode;
use cyw43::JoinOptions;
use cyw43_pio::PioSpi;
use defmt::*;
use embassy_executor::{Executor, InterruptExecutor, Spawner};
use embassy_net::tcp::client::{TcpClient, TcpClientState};
use embassy_net::{Config, StackResources};
use embassy_rp::bind_interrupts;
use embassy_rp::clocks::RoscRng;
use embassy_rp::gpio::{Level, Output};
use embassy_rp::interrupt;
use embassy_rp::interrupt::{InterruptExt, Priority};
use embassy_rp::multicore::{spawn_core1, Stack};
use embassy_rp::peripherals;
use embassy_rp::peripherals::{DMA_CH0, PIO0};
use embassy_rp::pio::{InterruptHandler, Pio};
use embassy_rp::spi::{Phase, Polarity, Spi};
use embassy_sync::blocking_mutex::raw::ThreadModeRawMutex;
use embassy_time::Timer;
//use embedded_nal_async::TcpConnect;
//use embedded_nal_async::stack::tcp::TcpConnect;
use embedded_nal_async::TcpConnect;
use rand::RngCore;
use smart_leds::RGB8;
use static_cell::StaticCell;
use ws2812_async::{ColorOrder, Ws2812};
use {defmt_rtt as _, panic_probe as _};

use rust_mqtt::{
    client::{client::MqttClient, client_config::ClientConfig},
    utils::rng_generator::CountingRng,
};
//pub type RGB8 = RGB<u8>;

use {defmt_rtt as _, panic_probe as _};

// after observing somewhat jumpy behavior of the neopixel task, I decided to set the scheduler and orhestrator to high priority
// hight priority runs on interrupt
static EXECUTOR_HIGH: InterruptExecutor = InterruptExecutor::new();
// low priority runs in thread-mode
static EXECUTOR_LOW: StaticCell<Executor> = StaticCell::new();

#[global_allocator]
static ALLOCATOR: emballoc::Allocator<4096> = emballoc::Allocator::new();
extern crate alloc;

#[interrupt]
unsafe fn SWI_IRQ_1() {
    EXECUTOR_HIGH.on_interrupt()
}

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => InterruptHandler<PIO0>;
});

const WIFI_NETWORK: &str = "cat";
const WIFI_PASSWORD: &str = "catcat123cat";
const CLIENT_ID: &str = "pico-495f6297-b962-4266-9c00-74138ae5a1f1";
const TOPIC: &str = "hello";
const NUM_LEDS: usize = 100;

static CHANNEL: embassy_sync::channel::Channel<ThreadModeRawMutex, LedStrip, 64> =
    embassy_sync::channel::Channel::new();

#[embassy_executor::task]
async fn cyw43_task(
    runner: cyw43::Runner<'static, Output<'static>, PioSpi<'static, PIO0, 0, DMA_CH0>>,
) -> ! {
    runner.run().await
}

#[embassy_executor::task]
async fn net_task(mut runner: embassy_net::Runner<'static, cyw43::NetDriver<'static>>) -> ! {
    runner.run().await
}

assign_resources! {
    led_peripheral: LedPeripheral {
        inner_spi: SPI1,
        clk_pin: PIN_14, // this is just a dummy pin, the neopixel uses only the mosi pin
        mosi_pin: PIN_15,
        tx_dma_ch: DMA_CH1,
    },
    wifi: WifiResources {
        pwr_pin: PIN_23,
        cs_pin: PIN_25,
        pio_sm: PIO0,
        dio_pin: PIN_24,
        clk_pin: PIN_29,
        dma_ch: DMA_CH0,
    },
}

#[embassy_executor::main]
async fn main(_spawner: Spawner) {
    info!("Main thread!");

    let p = embassy_rp::init(Default::default());
    let r = split_resources!(p);

    interrupt::SWI_IRQ_1.set_priority(Priority::P2);

    let executor = EXECUTOR_LOW.init(Executor::new());

    executor.run(|spawner| {
        // update the RTC
        spawner.spawn(run_mqtt(spawner, r.wifi)).unwrap();

        static mut CORE1_STACK: Stack<4096> = Stack::new();
        static EXECUTOR1: StaticCell<Executor> = StaticCell::new();

        spawn_core1(
            p.CORE1,
            unsafe { &mut *core::ptr::addr_of_mut!(CORE1_STACK) },
            move || {
                let executor1 = EXECUTOR1.init(Executor::new());
                executor1
                    .run(|spawner| unwrap!(spawner.spawn(led_ctrl_core1_task(r.led_peripheral))));
            },
        );
    });
}

#[embassy_executor::task]
pub async fn run_mqtt(spawner: Spawner, r: WifiResources) {
    let mut rng = RoscRng;

    let fw = include_bytes!("../embassy/cyw43-firmware/43439A0.bin");
    let clm = include_bytes!("../embassy/cyw43-firmware/43439A0_clm.bin");

    info!("init wifi");
    let pwr = Output::new(r.pwr_pin, Level::Low);
    let cs = Output::new(r.cs_pin, Level::High);
    let mut pio = Pio::new(r.pio_sm, Irqs);
    let spi = PioSpi::new(
        &mut pio.common,
        pio.sm0,
        pio.irq0,
        cs,
        r.dio_pin,
        r.clk_pin,
        r.dma_ch,
    );

    static STATE: StaticCell<cyw43::State> = StaticCell::new();
    let state = STATE.init(cyw43::State::new());
    let (net_device, mut control, runner) = cyw43::new(state, pwr, spi, fw).await;
    unwrap!(spawner.spawn(cyw43_task(runner)));

    control.init(clm).await;
    control
        .set_power_management(cyw43::PowerManagementMode::PowerSave)
        .await;

    let config = Config::dhcpv4(Default::default());
    // Use static IP configuration instead of DHCP
    //let config = embassy_net::Config::ipv4_static(embassy_net::StaticConfigV4 {
    //    address: Ipv4Cidr::new(Ipv4Address::new(192, 168, 69, 2), 24),
    //    dns_servers: Vec::new(),
    //    gateway: Some(Ipv4Address::new(192, 168, 69, 1)),
    //});

    // Generate random seed
    let seed = rng.next_u64();

    // Init network stack
    static RESOURCES: StaticCell<StackResources<5>> = StaticCell::new();
    let (stack, runner) = embassy_net::new(
        net_device,
        config,
        RESOURCES.init(StackResources::new()),
        seed,
    );

    unwrap!(spawner.spawn(net_task(runner)));

    loop {
        match control
            .join(WIFI_NETWORK, JoinOptions::new(WIFI_PASSWORD.as_bytes()))
            .await
        {
            Ok(_) => break,
            Err(err) => {
                info!("join failed with status={}", err.status);
            }
        }
    }

    // Wait for DHCP, not necessary when using static IP
    info!("waiting for DHCP...");
    while !stack.is_config_up() {
        Timer::after_millis(100).await;
    }
    info!("DHCP is now up!");

    info!("waiting for link up...");
    while !stack.is_link_up() {
        Timer::after_millis(500).await;
    }
    info!("Link is up!");

    info!("waiting for stack to be up...");
    stack.wait_config_up().await;
    info!("Stack is up!");

    let tcp_state: TcpClientState<1, 1024, 1024> = TcpClientState::new();
    let tcp_client = TcpClient::new(stack, &tcp_state);

    let addr = embedded_nal_async::SocketAddr::new(
        embedded_nal_async::Ipv4Addr::new(192, 168, 0, 100).into(),
        1883,
    );

    // let connection = embedded_nal::TcpStream::connect(addr)
    //     .await
    //     .map_err(|_| ReasonCode::NetworkError)
    //     .unwrap();
    //let connection = FromTokio::<TcpStream>::new(connection);
    let mut config = ClientConfig::new(
        rust_mqtt::client::client_config::MqttVersion::MQTTv5,
        CountingRng(20000),
    );

    config.add_max_subscribe_qos(rust_mqtt::packet::v5::publish_packet::QualityOfService::QoS0);
    config.add_client_id(CLIENT_ID);
    config.keep_alive = u16::MAX;

    // config.add_username(USERNAME);
    // config.add_password(PASSWORD);
    config.max_packet_size = 400;
    let mut recv_buffer = [0; 400];
    let mut write_buffer = [0; 400];

    //let r = tcp_client.connect(addr).await;

    let r;

    loop {
        match tcp_client.connect(addr).await {
            Ok(rp) => {
                info!("Attempting to connect to mqtt broker");
                r = rp;
                break;
            }
            Err(err) => {
                info!("Failed to join broker: {}. Trying again!", err);
                Timer::after_millis(100).await;
            }
        };
    }

    let mut client: MqttClient<
        '_,
        embassy_net::tcp::client::TcpConnection<'_, 1, 1024, 1024>,
        5,
        CountingRng,
    > = MqttClient::<_, 5, _>::new(r, &mut write_buffer, 512, &mut recv_buffer, 512, config);

    loop {
        match client.connect_to_broker().await {
            Ok(_) => {
                info!("CAT Connected to broker");
                break;
            }
            Err(err) => {
                info!("DOG connect_to_broker failed with status={}", err);
            }
        }
    }

    loop {
        match client.subscribe_to_topic(TOPIC).await {
            Ok(_) => {
                info!("CAT subscribed to topic");
                break;
            }
            Err(err) => {
                info!("DOG subscribe_to_topic failed with status={}", err);
            }
        }
    }

    let mut old_msg = LedStrip::new(0);

    loop {
        info!("[RECEIVER] Waiting for new message");
        let msg = client.receive_message().await;

        info!("[RECEIVER] Received msg: {}", msg);

        if msg.is_ok() {
            let message = msg.unwrap();

            info!("[RECEIVER] message has size: {}", message.1.len());

            match bincode::decode_from_slice::<LedStrip, Configuration>(
                message.1,
                bincode::config::standard(),
            ) {
                Ok(decoded) => {
                    if old_msg != decoded.0 {
                        CHANNEL.send(decoded.0.clone()).await;
                        old_msg = decoded.0;
                        info!("[RECEIVER] Sent message");
                    }
                }
                Err(_e) => {
                    info!("DOG");
                }
            }
        } else {
            warn!("[RECEIVER] Could not get message with error: {}", msg.err());
        }
        Timer::after(Duration::from_millis(2)).await;
    }
}

#[embassy_executor::task]
async fn led_ctrl_core1_task(led_peripheral: LedPeripheral) {
    info!("Hello from core 1");

    let mut current_ledstrip = LedStrip::new(0);

    // Spi configuration for the neopixel
    let mut spi_config = embassy_rp::spi::Config::default();
    spi_config.frequency = 3_800_000;
    spi_config.phase = Phase::CaptureOnFirstTransition;
    spi_config.polarity = Polarity::IdleLow;
    //    let spi = embassy_rp::spi::Spi::new_txonly(led_peripheral.inner_spi, led_peripheral.clk_pin, led_peripheral.mosi_pin, led_peripheral.tx_dma_ch, spi_config);
    let spi = Spi::new_txonly(
        led_peripheral.inner_spi,
        led_peripheral.clk_pin,
        led_peripheral.mosi_pin,
        led_peripheral.tx_dma_ch,
        spi_config,
    );
    let mut np: Ws2812<_, { 12 * NUM_LEDS }> = Ws2812::new(spi);
    np.set_color_order(ColorOrder::GRB);

    let mut data = [RGB8::default(); NUM_LEDS];

    for led in data.iter_mut().step_by(1) {
        led.r = 250; // blue
        led.g = 150; // red
        led.b = 0; // green
    }

    //np.write(empty.iter().cloned()).await.ok();

    //np.write(data.iter().cloned()).await.ok();
    let mut rgb_arr: [RGB8; NUM_LEDS] = [RGB8::new(0, 0, 0); NUM_LEDS];

    loop {
        let rc = CHANNEL.receive().await;
        info!("OCTOPUS yyyy");

        if current_ledstrip != rc {
            info!("OCTOPUS got a new led strip");

            for i in 0..rc.leds.len() {
                rgb_arr[i].r = rc.leds[i].r;
                rgb_arr[i].g = rc.leds[i].g;
                rgb_arr[i].b = rc.leds[i].b;
            }

            current_ledstrip = rc.clone();

            info!("writing led!");
            np.write(rgb_arr.iter().cloned()).await.ok();

            //            Timer::after(Duration::from_millis(0)).await;
        }
    }
}

#[derive(Debug, Decode, Clone, Copy, PartialEq)]
struct Led {
    r: u8,
    g: u8,
    b: u8,
}

#[derive(Clone, Debug, Decode, PartialEq)]
pub struct LedStrip {
    leds: [Led; NUM_LEDS],
}

impl LedStrip {
    // naive inital start with uniform colour on entire strip
    pub fn new(_r: u8) -> Self {
        let arr: [Led; NUM_LEDS] = [Led { r: 0, g: 0, b: 0 }; NUM_LEDS];

        LedStrip { leds: arr }
    }
}
