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

use core::{cell::RefCell, net::Ipv4Addr, str::FromStr};

use approuter::{AppRouter, AppState};
use bme280::i2c::AsyncBME280;
use cyw43::{JoinOptions, PowerManagementMode};
use cyw43_pio::{PioSpi, DEFAULT_CLOCK_DIVIDER};
use defmt::{debug, info, warn};
use embassy_boot::{AlignedBuffer, BlockingFirmwareUpdater, FirmwareUpdaterConfig};
use embassy_embedded_hal::flash::partition::BlockingPartition;
use embassy_executor::Spawner;
use embassy_net::{
    tcp::{
        client::{TcpClient, TcpClientState},
        TcpSocket,
    },
    Config, DhcpConfig, Stack, StackResources,
};
use embassy_rp::{
    bind_interrupts,
    clocks::RoscRng,
    flash::{Flash, WRITE_SIZE},
    gpio::{Flex, Input, Level, Output, Pin, Pull},
    i2c::{self},
    peripherals::{DMA_CH1, FLASH, I2C0, PIO0},
    pio::{self, Pio},
    Peripheral,
};
use embassy_sync::{
    blocking_mutex::{
        self,
        raw::{CriticalSectionRawMutex, NoopRawMutex},
    },
    mutex::Mutex,
};
use embassy_time::{Delay, Duration, Timer};
use heapless::String;
use picoserve::{make_static, Router};
use rand::{rngs::StdRng, RngCore, SeedableRng};
use static_cell::StaticCell;

use {defmt_rtt as _, panic_probe as _};

use air_filter_types::{FanSpeed, FanState};

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

const FLASH_SIZE: usize = 2 * 1024 * 1024;

struct Controller<'a> {
    off: Input<'a>,
    low: Flex<'a>,
    medium: Flex<'a>,
    high: Flex<'a>,
}

impl<'a> Controller<'a> {
    pub fn new<O: Pin, L: Pin, M: Pin, H: Pin>(
        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) -> FanState {
        let manual = self.off.is_high();
        let speed = match (self.low.is_low(), self.medium.is_low(), self.high.is_low()) {
            (false, false, false) => FanSpeed::Off,
            (true, false, false) => FanSpeed::Low,
            (false, true, false) => FanSpeed::Medium,
            (false, false, true) => FanSpeed::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_speed(FanSpeed::Off);
                FanSpeed::Off
            }
        };

        FanState { speed, manual }
    }

    pub fn set_speed(&mut self, speed: FanSpeed) -> bool {
        let manual = self.off.is_high();

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

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

        match speed {
            FanSpeed::Off => {
                self.low.set_as_input();
                self.medium.set_as_input();
                self.high.set_as_input();
            }
            FanSpeed::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();
            }
            FanSpeed::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();
            }
            FanSpeed::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);
            }
        }

        true
    }

    pub fn check_for_manual(&mut self) {
        // If off is high, that means that the knob is used to override the state
        if self.off.is_high() {
            self.set_speed(FanSpeed::Off);
        }
    }
}

// We need this because otherwise the compiler yells at us
mod approuter {
    use air_filter_types::{SensorData, SetFanSpeed};
    use bme280::i2c::AsyncBME280;
    use defmt::debug;
    use embassy_boot::BlockingFirmwareUpdater;
    use embassy_rp::{
        i2c::{Async, I2c},
        peripherals::I2C0,
    };
    use embassy_sync::{blocking_mutex::raw::CriticalSectionRawMutex, mutex::Mutex};
    use embassy_time::Delay;
    use picoserve::{
        extract, make_static,
        routing::{get, PathRouter},
        Router,
    };
    use updater::firmware_router;

    use crate::{Controller, Partition};

    const VERSION: &str = git_version::git_version!();
    const PUBLIC_SIGNING_KEY: &[u8; 32] = include_bytes!("../key.pub");

    #[derive(Clone, Copy)]
    pub struct SharedController(pub &'static Mutex<CriticalSectionRawMutex, Controller<'static>>);

    impl SharedController {
        fn new(controller: Controller<'static>) -> Self {
            Self(
                make_static!(Mutex<CriticalSectionRawMutex, Controller<'static>>, Mutex::new(controller)),
            )
        }
    }

    #[derive(Clone, Copy)]
    pub struct SharedBME280(
        pub &'static Mutex<CriticalSectionRawMutex, AsyncBME280<I2c<'static, I2C0, Async>>>,
    );

    impl SharedBME280 {
        fn new(bme280: AsyncBME280<I2c<'static, I2C0, Async>>) -> Self {
            Self(
                make_static!(Mutex<CriticalSectionRawMutex, AsyncBME280<I2c<'static, I2C0, Async>>>, Mutex::new(bme280)),
            )
        }
    }

    #[derive(Clone)]
    pub struct AppState {
        pub shared_controller: SharedController,
        pub shared_bme280: SharedBME280,
    }

    impl AppState {
        pub fn new(
            controller: Controller<'static>,
            bme280: AsyncBME280<I2c<'static, I2C0, Async>>,
        ) -> Self {
            Self {
                shared_controller: SharedController::new(controller),
                shared_bme280: SharedBME280::new(bme280),
            }
        }
    }

    impl picoserve::extract::FromRef<AppState> for SharedController {
        fn from_ref(state: &AppState) -> Self {
            state.shared_controller
        }
    }

    impl picoserve::extract::FromRef<AppState> for SharedBME280 {
        fn from_ref(state: &AppState) -> Self {
            state.shared_bme280
        }
    }

    fn state_router() -> Router<impl PathRouter<AppState>, AppState> {
        Router::new()
            .route(
                "/fan",
                get(
                    |extract::State(SharedController(controller)): extract::State<
                        SharedController,
                    >| async {
                        debug!("Getting fan state");

                        let state = controller.lock().await.get_state();
                        picoserve::response::Json(state)
                    },
                )
                .put(
                    |extract::State(SharedController(controller)): extract::State<
                        SharedController,
                    >,
                     extract::Json(message): extract::Json<SetFanSpeed>| async move {
                        debug!("Setting fan state");

                        let success = controller.lock().await.set_speed(message.speed());

                        if success {
                            picoserve::response::StatusCode::OK
                        } else {
                            picoserve::response::StatusCode::SERVICE_UNAVAILABLE
                        }
                    },
                ),
            )
            .route(
                "/sensor",
                get(
                    |extract::State(SharedBME280(bme280)): extract::State<SharedBME280>| async {
                        debug!("Getting sensor state");

                        let measurement = bme280
                            .lock()
                            .await
                            .measure(&mut Delay {})
                            .await
                            .expect("Measurement should work");
                        let sensor_data = SensorData::new(measurement);

                        picoserve::response::Json(sensor_data)
                    },
                ),
            )
    }

    pub type AppRouter = impl PathRouter<AppState>;
    pub fn make_app(
        updater: &'static Mutex<
            CriticalSectionRawMutex,
            BlockingFirmwareUpdater<'static, Partition, Partition>,
        >,
    ) -> Router<AppRouter, AppState> {
        Router::new().nest("/state", state_router()).nest(
            "/firmware",
            firmware_router(VERSION, updater, PUBLIC_SIGNING_KEY),
        )
    }
}

/// 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)
        }
    }
}

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

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

type Partition = BlockingPartition<
    'static,
    NoopRawMutex,
    Flash<'static, FLASH, embassy_rp::flash::Blocking, FLASH_SIZE>,
>;

const WEB_TASK_POOL_SIZE: usize = 8;

#[embassy_executor::task(pool_size = WEB_TASK_POOL_SIZE)]
async fn web_task(
    id: usize,
    stack: Stack<'static>,
    app: &'static Router<AppRouter, AppState>,
    config: &'static picoserve::Config<Duration>,
    state: AppState,
) -> ! {
    let port = 80;
    let mut tcp_rx_buffer = [0; 1024];
    let mut tcp_tx_buffer = [0; 1024];
    let mut http_buffer = [0; 2048];

    picoserve::listen_and_serve_with_state(
        id,
        app,
        config,
        stack,
        port,
        &mut tcp_rx_buffer,
        &mut tcp_tx_buffer,
        &mut http_buffer,
        &state,
    )
    .await;
}

#[embassy_executor::task]
async fn controller_task(state: AppState) {
    loop {
        state.shared_controller.0.lock().await.check_for_manual();
        Timer::after(Duration::from_millis(500)).await;
    }
}

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

    let flash = Flash::<_, _, FLASH_SIZE>::new_blocking(p.FLASH);
    let flash = make_static!(blocking_mutex::Mutex<NoopRawMutex, RefCell<Flash<'static, FLASH, embassy_rp::flash::Blocking, FLASH_SIZE>>>, blocking_mutex::Mutex::new(RefCell::new(flash)));

    let config = FirmwareUpdaterConfig::from_linkerfile_blocking(flash, flash);
    let aligned = make_static!(AlignedBuffer<WRITE_SIZE>, AlignedBuffer([0; WRITE_SIZE]));
    let updater = BlockingFirmwareUpdater::new(config, &mut aligned.0);
    let updater = make_static!(Mutex<CriticalSectionRawMutex, BlockingFirmwareUpdater<'static, Partition, Partition>>, Mutex::new(updater));

    let 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,
        DEFAULT_CLOCK_DIVIDER,
        pio.irq0,
        cs,
        p.PIN_24,
        p.PIN_29,
        p.DMA_CH1,
    );

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

    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;
    spawner.spawn(cyw43_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 mut config: DhcpConfig = Default::default();
    config.hostname = Some(String::from_str("airfilter").expect("Is valid"));
    let config = Config::dhcpv4(config);

    // 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();

    static RESOURCES: StaticCell<StackResources<16>> = StaticCell::new();
    let (stack, runner) = embassy_net::new(
        net_device,
        config,
        RESOURCES.init(StackResources::new()),
        rng.next_u64(),
    );
    spawner.spawn(net_task(runner)).unwrap();

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

    // Mark the update as successful after connecting to the wifi
    updater.lock().await.mark_booted().unwrap();

    info!("Waiting for link up...");
    stack.wait_link_up().await;
    info!("Link is up!");

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

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

    let cfg = stack.config_v4().expect("Config is up");
    info!("IP Address: {}", cfg.address.address());

    let config = make_static!(
        picoserve::Config<Duration>,
        picoserve::Config::new(picoserve::Timeouts {
            start_read_request: Some(Duration::from_secs(5)),
            read_request: Some(Duration::from_secs(1)),
            write: Some(Duration::from_secs(1)),
        })
        .keep_connection_alive()
    );

    let app = make_static!(Router<AppRouter, AppState>, approuter::make_app(updater));

    let state = AppState::new(controller, bme280);

    for id in 0..WEB_TASK_POOL_SIZE {
        spawner.must_spawn(web_task(id, stack, app, config, state.clone()));
    }

    spawner.must_spawn(controller_task(state));

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

    loop {
        // TODO: In the future, use reqless to push the current state to automation_rs
        control.gpio_set(0, true).await;
        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
        socket.set_timeout(Some(Duration::from_secs(1)));
        socket
            .connect(("10.0.0.2".parse::<Ipv4Addr>().unwrap(), 80))
            .await
            .ok();
        control.gpio_set(0, false).await;
        Timer::after(Duration::from_secs(60)).await;
    }
}