#![no_std]
#![no_main]
#![feature(type_alias_impl_trait)]

use core::fmt::Debug;
use core::{cell::RefCell, str::from_utf8};

use bme280::{i2c::AsyncBME280, Measurements};
use const_format::formatcp;
use cyw43::PowerManagementMode;
use cyw43_pio::PioSpi;
use defmt::{debug, error, info, warn, Display2Format, Format};
use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, FirmwareUpdaterConfig};
use embassy_executor::Spawner;
use embassy_futures::select::{select3, select4, Either3};
use embassy_net::{dns::DnsQueryType, tcp::TcpSocket, Config, Stack, StackResources};
use embassy_rp::{
    bind_interrupts,
    clocks::RoscRng,
    flash::{Flash, WRITE_SIZE},
    gpio::{Flex, Input, Level, Output, Pin, Pull},
    i2c,
    peripherals::{DMA_CH1, I2C0, PIN_23, PIN_25, PIO0},
    pio::{self, Pio},
    Peripheral,
};
use embassy_sync::blocking_mutex::Mutex;
use embassy_time::{Delay, Duration, Ticker, Timer};
use heapless::Vec;
use nourl::Url;
use rand::{
    rngs::{SmallRng, StdRng},
    RngCore, SeedableRng,
};
use rust_mqtt::{
    client::{
        client::MqttClient,
        client_config::{ClientConfig, MqttVersion},
    },
    packet::v5::publish_packet::QualityOfService,
};
use serde::{Deserialize, Serialize};
use static_cell::make_static;

use {defmt_rtt as _, panic_probe as _};

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => pio::InterruptHandler<PIO0>;
    I2C0_IRQ => i2c::InterruptHandler<I2C0>;
});

const ID: &str = env!("ID");
const TOPIC_BASE: &str = formatcp!("pico/{}", ID);
const TOPIC_STATUS: &str = formatcp!("{}/status", TOPIC_BASE);
const TOPIC_UPDATE: &str = formatcp!("{}/update", TOPIC_BASE);
const TOPIC_SET: &str = formatcp!("{}/set", TOPIC_BASE);
const VERSION: &str = git_version::git_version!();
const PUBLIC_SIGNING_KEY: &[u8; 32] = include_bytes!("../key.pub");
const FLASH_SIZE: usize = 2 * 1024 * 1024;

#[derive(Deserialize)]
struct SetMessage {
    state: State,
}

#[derive(Serialize)]
struct StateMessage {
    state: State,
    manual: bool,
    temperature: f32,
    humidity: f32,
    pressure: f32,
}

impl StateMessage {
    pub fn new<E: Debug>((state, manual): (State, bool), measurements: Measurements<E>) -> Self {
        Self {
            state,
            manual,
            temperature: measurements.temperature,
            humidity: measurements.humidity,
            pressure: measurements.pressure,
        }
    }

    pub fn vec(&self) -> Vec<u8, 128> {
        serde_json_core::to_vec(self)
            .expect("The buffer should be large enough to contain all the data")
    }
}

#[derive(Format, PartialEq, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
enum State {
    Off,
    Low,
    Medium,
    High,
}

impl SetMessage {
    fn get_state(&self) -> State {
        self.state
    }
}

struct Controller<'a, O: Pin, L: Pin, M: Pin, H: Pin> {
    off: Input<'a, O>,
    low: Flex<'a, L>,
    medium: Flex<'a, M>,
    high: Flex<'a, H>,
}

impl<'a, O: Pin, L: Pin, M: Pin, H: Pin> Controller<'a, O, L, M, H> {
    pub fn new(
        off: impl Peripheral<P = O> + 'a,
        low: impl Peripheral<P = L> + 'a,
        medium: impl Peripheral<P = M> + 'a,
        high: impl Peripheral<P = H> + 'a,
    ) -> Self {
        let off = Input::new(off, Pull::None);

        let mut low = Flex::new(low);
        low.set_low();
        low.set_as_input();

        let mut medium = Flex::new(medium);
        medium.set_low();
        medium.set_as_input();

        let mut high = Flex::new(high);
        high.set_low();
        high.set_as_input();

        Self {
            off,
            low,
            medium,
            high,
        }
    }

    pub fn get_state(&mut self) -> (State, bool) {
        let manual = self.off.is_high();
        let state = match (self.low.is_low(), self.medium.is_low(), self.high.is_low()) {
            (false, false, false) => State::Off,
            (true, false, false) => State::Low,
            (false, true, false) => State::Medium,
            (false, false, true) => State::High,
            (a, b, c) => {
                // This happens if the user turns the knob, in this case we should turn off remote
                // control
                debug!("Unknown state: ({}, {}, {})", a, b, c);
                self.set_state(State::Off);
                State::Off
            }
        };

        (state, manual)
    }

    pub fn set_state(&mut self, state: State) {
        let manual = self.off.is_high();

        if manual && state != State::Off {
            warn!("Filter is manual controlled, cannot control remotely");
            return;
        }

        debug!("Setting state: {}", state);

        match state {
            State::Off => {
                self.low.set_as_input();
                self.medium.set_as_input();
                self.high.set_as_input();
            }
            State::Low => {
                self.low.set_as_output();
                self.low.set_drive_strength(embassy_rp::gpio::Drive::_12mA);
                self.medium.set_as_input();
                self.high.set_as_input();
            }
            State::Medium => {
                self.low.set_as_input();
                self.medium.set_as_output();
                self.medium
                    .set_drive_strength(embassy_rp::gpio::Drive::_12mA);
                self.high.set_as_input();
            }
            State::High => {
                self.low.set_as_input();
                self.medium.set_as_input();
                self.high.set_as_output();
                self.high.set_drive_strength(embassy_rp::gpio::Drive::_12mA);
            }
        }
    }

    pub async fn watch(&mut self) -> (State, bool) {
        // Wait for change on any of the pins
        select4(
            self.off.wait_for_any_edge(),
            self.low.wait_for_any_edge(),
            self.medium.wait_for_any_edge(),
            self.high.wait_for_any_edge(),
        )
        .await;

        // Give it some time to stabilze
        Timer::after(Duration::from_millis(500)).await;

        if self.off.is_high() {
            // If the filter is in manual mode, set the pico outputs to off
            self.set_state(State::Off);
        }

        // Get the current state
        self.get_state()
    }
}

/// Get the cyw43 firmware blobs
///
/// # Safety
/// When building without `include_firmwares` make sure to flash the firmwares using the following
/// commands:
/// ```bash
/// probe-rs download firmware/43439A0.bin --format bin --chip RP2040 --base-address 0x101BE000
/// probe-rs download firmware/43439A0_clm.bin --format bin --chip RP2040 --base-address 0x101FE000
/// ```
unsafe fn get_firmware() -> (&'static [u8], &'static [u8]) {
    cfg_if::cfg_if! {
        if #[cfg(feature = "include_firmwares")] {
            let fw = include_bytes!("../firmware/43439A0.bin");
            let clm = include_bytes!("../firmware/43439A0_clm.bin");

            (fw, clm)
        } else {
            // TODO: It would be nice if it could automatically get the correct size
            extern "C" {
                #[link_name = "__fw_start"]
                static fw: [u8; 230321];
                #[link_name = "__clm_start"]
                static clm: [u8; 4752];
            }

            (&fw, &clm)
        }
    }
}

async fn wait_for_config(
    stack: &'static Stack<cyw43::NetDriver<'static>>,
) -> embassy_net::StaticConfigV4 {
    for _ in 0..120 {
        // We are essentially busy looping here since there is no Async API for this
        if let Some(config) = stack.config_v4() {
            return config;
        }

        Timer::after_secs(1).await;
    }

    info!("Restarting...");
    cortex_m::peripheral::SCB::sys_reset();
}

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

#[embassy_executor::task]
async fn net_task(stack: &'static Stack<cyw43::NetDriver<'static>>) -> ! {
    stack.run().await
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    info!("Starting...");
    let p = embassy_rp::init(Default::default());

    // TODO: Ideally we use async flash
    // This has issues with alignment right now
    let flash = Flash::<_, _, FLASH_SIZE>::new_blocking(p.FLASH);
    let flash = Mutex::new(RefCell::new(flash));

    let config = FirmwareUpdaterConfig::from_linkerfile_blocking(&flash);
    let mut aligned = AlignedBuffer([0; WRITE_SIZE]);
    let updater = BlockingFirmwareUpdater::new(config, &mut aligned.0);

    let mut updater = updater::Updater::new(updater, TOPIC_STATUS, VERSION, PUBLIC_SIGNING_KEY);

    let mut controller = Controller::new(p.PIN_28, p.PIN_27, p.PIN_26, p.PIN_22);

    let i2c = i2c::I2c::new_async(p.I2C0, p.PIN_9, p.PIN_8, Irqs, i2c::Config::default());
    let mut bme280 = AsyncBME280::new_primary(i2c);
    let mut delay = Delay {};
    bme280.init(&mut delay).await.unwrap();

    let pwr = Output::new(p.PIN_23, Level::Low);
    let cs = Output::new(p.PIN_25, Level::High);

    let mut pio = Pio::new(p.PIO0, Irqs);
    let spi = PioSpi::new(
        &mut pio.common,
        pio.sm0,
        pio.irq0,
        cs,
        p.PIN_24,
        p.PIN_29,
        p.DMA_CH1,
    );

    let (fw, clm) = unsafe { get_firmware() };

    let state = make_static!(cyw43::State::new());
    let (net_device, mut control, runner) = cyw43::new(state, pwr, spi, fw).await;
    spawner.spawn(wifi_task(runner)).unwrap();

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

    // Turn LED on while trying to connect
    control.gpio_set(0, true).await;

    let config = Config::dhcpv4(Default::default());

    // Use the Ring Oscillator of the RP2040 as a source of true randomness to seed the
    // cryptographically secure PRNG
    let mut rng = StdRng::from_rng(&mut RoscRng).unwrap();

    let stack = make_static!(Stack::new(
        net_device,
        config,
        make_static!(StackResources::<6>::new()),
        rng.next_u64(),
    ));

    spawner.spawn(net_task(stack)).unwrap();

    // Connect to wifi
    loop {
        match control
            .join_wpa2(env!("WIFI_NETWORK"), env!("WIFI_PASSWORD"))
            .await
        {
            Ok(_) => break,
            Err(err) => {
                info!("Failed to join with status = {}", err.status)
            }
        }
    }

    info!("Waiting for DHCP...");
    let cfg = wait_for_config(stack).await;
    info!("IP Address: {}", cfg.address.address());

    let mut rx_buffer = [0; 1024];
    let mut tx_buffer = [0; 1024];

    let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
    // socket.set_timeout(Some(Duration::from_secs(10)));
    let url = Url::parse(env!("MQTT_ADDRESS")).unwrap();
    debug!("MQTT URL: {}", url);
    let ip = stack.dns_query(url.host(), DnsQueryType::A).await.unwrap()[0];
    let addr = (ip, url.port_or_default());
    debug!("MQTT ADDR: {}", addr);

    while let Err(e) = socket.connect(addr).await {
        warn!("Connect error: {:?}", e);
        Timer::after(Duration::from_secs(1)).await;
    }
    info!("TCP Connected!");

    let mut config = ClientConfig::new(
        MqttVersion::MQTTv5,
        // Use fast and simple PRNG to generate packet identifiers, there is no need for this to be
        // cryptographically secure
        SmallRng::from_rng(&mut RoscRng).unwrap(),
    );

    config.add_username(env!("MQTT_USERNAME"));
    config.add_password(env!("MQTT_PASSWORD"));
    config.add_max_subscribe_qos(QualityOfService::QoS1);
    config.add_client_id(ID);
    updater.add_will(&mut config);

    let mut recv_buffer = [0; 1024];
    let mut write_buffer = [0; 1024];

    let mut client = MqttClient::<_, 5, _>::new(
        socket,
        &mut write_buffer,
        1024,
        &mut recv_buffer,
        1024,
        config,
    );

    info!("Connecting to MQTT...");
    if client.connect_to_broker().await.is_err() {
        cortex_m::peripheral::SCB::sys_reset();
    };
    info!("MQTT Connected!");

    // We wait with marking as booted until everything is connected
    client.subscribe_to_topic(TOPIC_UPDATE).await.unwrap();
    client.subscribe_to_topic(TOPIC_SET).await.unwrap();
    updater.ready(&mut client).await.unwrap();

    // Turn LED off when connected
    control.gpio_set(0, false).await;

    let mut keep_alive = Ticker::every(Duration::from_secs(30));
    loop {
        let message = match select3(
            keep_alive.next(),
            client.receive_message(),
            controller.watch(),
        )
        .await
        {
            Either3::First(_) => Some(StateMessage::new(
                controller.get_state(),
                bme280.measure(&mut delay).await.unwrap(),
            )),
            Either3::Second(message) => match message {
                Ok((TOPIC_UPDATE, url)) => {
                    let url: Vec<_, 256> = match Vec::from_slice(url) {
                        Ok(url) => url,
                        Err(_) => {
                            error!("URL is longer then buffer size");
                            continue;
                        }
                    };
                    let url = match from_utf8(&url) {
                        Ok(url) => url,
                        Err(err) => {
                            error!("Url is not valid utf-8 string: {}", Display2Format(&err));
                            continue;
                        }
                    };
                    let url = match Url::parse(url) {
                        Ok(url) => url,
                        Err(err) => {
                            error!("Failed to parse url: {}", err);
                            continue;
                        }
                    };

                    if let Err(err) = updater.update(url, stack, &mut rng, &mut client).await {
                        error!("Update failed: {}", err);
                    }

                    None
                }
                Ok((TOPIC_SET, message)) => {
                    let message: SetMessage = match serde_json_core::from_slice(message) {
                        Ok((message, _)) => message,
                        Err(_) => {
                            error!("Unable to parse set message");
                            continue;
                        }
                    };

                    controller.set_state(message.get_state());

                    Some(StateMessage::new(
                        controller.get_state(),
                        bme280.measure(&mut delay).await.unwrap(),
                    ))
                }
                Ok(_) => None,
                Err(err) => {
                    error!("{}", err);
                    info!("Restarting in 5s...");
                    Timer::after(Duration::from_secs(5)).await;
                    info!("Restarting...");
                    cortex_m::peripheral::SCB::sys_reset();
                }
            },
            Either3::Third(state) => Some(StateMessage::new(
                state,
                bme280.measure(&mut delay).await.unwrap(),
            )),
        };

        if let Some(message) = message {
            client
                .send_message(TOPIC_BASE, &message.vec(), QualityOfService::QoS1, true)
                .await
                .unwrap();
        }
    }
}