//! This example uses the RP Pico W board Wifi chip (cyw43).
//! Creates an Access point Wifi network and creates a TCP endpoint on port 1234.

#![no_std]
#![no_main]
#![allow(async_fn_in_trait)]

mod sensor_manager;

use core::array;
use core::borrow::BorrowMut;
use core::cell::RefCell;
use core::fmt::Formatter;
use core::mem::transmute;
use core::panic::PanicInfo;
use core::str::from_utf8;

use cyw43::JoinOptions;
use cyw43_pio::{PioSpi, DEFAULT_CLOCK_DIVIDER};
use embassy_rp::i2c::{Async, I2c};
use embassy_rp::multicore::{spawn_core1, Stack};
use embassy_rp::pwm::{self, Pwm};
use embassy_sync::blocking_mutex::raw::{CriticalSectionRawMutex, ThreadModeRawMutex};
use embassy_sync::blocking_mutex::{Mutex, ThreadModeMutex};
use embassy_sync::channel::Channel;
use embedded_hal_bus::i2c::RefCellDevice;
use fixed::FixedU16;
use log::*;
//use embassy_rp::i2c::InterruptHandler;
use embassy_executor::Spawner;
use embassy_net::tcp::TcpSocket;
use embassy_net::{Config, StackResources};
use embassy_rp::{bind_interrupts, watchdog};
use embassy_rp::clocks::RoscRng;
use embassy_rp::gpio::{Input, Level, Output, Pull};
use embassy_rp::peripherals::{DMA_CH0, PIO0, USB};
use embassy_rp::pio::{InterruptHandler, Pio};
use embassy_rp::usb::Driver;
use embassy_time::{Delay, Duration, Instant, Timer};
use embedded_io_async::{Read, Write};
use mpu6050::Mpu6050;
use nalgebra::clamp;
use rand::RngCore;
use reqwless::response;
use sensor_manager::{sensor_manager, SensorHardware};
use common::{ControlPacket, SensorData, TelemetryPacket};
use static_cell::StaticCell;
use defmt_rtt as _;

bind_interrupts!(struct Irqs {
    PIO0_IRQ_0 => InterruptHandler<PIO0>;
    I2C0_IRQ => embassy_rp::i2c::InterruptHandler<embassy_rp::peripherals::I2C0>;
    I2C1_IRQ => embassy_rp::i2c::InterruptHandler<embassy_rp::peripherals::I2C1>;
    USBCTRL_IRQ => embassy_rp::usb::InterruptHandler<USB>;
});

static mut CORE1_STACK: Stack<4096> = Stack::new();
pub static CHANNEL: Channel<CriticalSectionRawMutex, (SensorData, Instant), 1> = Channel::new();

#[embassy_executor::task]
async fn logger_task(driver: Driver<'static, USB>) {
    embassy_usb_logger::run!(1024, log::LevelFilter::Debug, driver);
}

#[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
}

#[panic_handler]
fn panic( info: &PanicInfo) -> ! {
    error!("{}", info);
    loop { }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    info!("Hello World!");

    let p = embassy_rp::init(Default::default());
    let mut rng = RoscRng;

    let driver = Driver::new(p.USB, Irqs);
    spawner.spawn(logger_task(driver)).unwrap();

    let limit_switch = Input::new(p.PIN_22, Pull::Up);

    let mut d: pwm::Config = Default::default();
    d.divider = 40.into();
    d.top = 62500; // 20ms
    d.compare_b = 4687; // 1.5ms
    d.compare_a = 4687; // 1.5ms
    let mut drive_pwm = Pwm::new_output_ab(p.PWM_SLICE1, p.PIN_18, p.PIN_19, d.clone());

    let mut f: pwm::Config = Default::default();
    f.divider = 40.into();
    f.top = 62500; // 20ms
    f.compare_b = 4687; // 1.5ms
    let flip_pwm = Pwm::new_output_b(p.PWM_SLICE0, p.PIN_17, f.clone());

    let sda = p.PIN_20;
    let scl = p.PIN_21;
    let config = embassy_rp::i2c::Config::default();
    let bus = embassy_rp::i2c::I2c::new_async(p.I2C0, scl, sda, Irqs, config);

    let sda = p.PIN_26;
    let scl = p.PIN_27;
    let config = embassy_rp::i2c::Config::default();
    let bus1 = embassy_rp::i2c::I2c::new_async(p.I2C1, scl, sda, Irqs, config);
    let tof2enable = Output::new(p.PIN_16, Level::Low);

    let hardware = SensorHardware {
        bus_tof: unsafe {transmute(bus)},
        bus_gyro: unsafe {transmute(bus1)},
        tof2enable,
    };

    spawn_core1(
        p.CORE1, 
        unsafe { &mut *core::ptr::addr_of_mut!(CORE1_STACK) }, 
        move || {
            sensor_manager(hardware)
        }
    );

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

    // To make flashing faster for development, you may want to flash the firmwares independently
    // at hardcoded addresses, instead of baking them into the program with `include_bytes!`:
    //     probe-rs download 43439A0.bin --binary-format bin --chip RP2040 --base-address 0x10100000
    //     probe-rs download 43439A0_clm.bin --binary-format bin --chip RP2040 --base-address 0x10140000
    //let fw = unsafe { core::slice::from_raw_parts(0x10100000 as *const u8, 230321) };
    //let clm = unsafe { core::slice::from_raw_parts(0x10140000 as *const u8, 4752) };

    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_CH0,
    );

    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;
    defmt::unwrap!(spawner.spawn(cyw43_task(runner)));

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

    // Use a link-local address for communication without DHCP server
    let config = Config::ipv4_static(embassy_net::StaticConfigV4 {
        address: embassy_net::Ipv4Cidr::new(embassy_net::Ipv4Address::new(169, 254, 1, 1), 16),
        dns_servers: heapless::Vec::new(),
        gateway: None,
    });

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

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

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

    // password is not terribly private information
    //control.start_ap_wpa2("cruisecontrol", "dxSk2avMFvsY", 5).await;

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

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

    let cam_state: RefCell<CamState> = RefCell::new(CamState::Idle);

    let cam_state_forever = unsafe { transmute(&cam_state) }; // SAFETY: if main exits we shall
                                                              // shortly die.
    spawner.spawn(cam_manager(flip_pwm, limit_switch, cam_state_forever)).unwrap();

    loop {
        let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
        socket.set_timeout(Some(Duration::from_secs(3)));

        let mut c: pwm::Config = Default::default();
        c.divider = 40.into();
        c.top = 62500; // 20ms
        c.compare_b = 4687; // 1.5ms
        c.compare_a = 4687; // 1.5ms
        drive_pwm.set_config(&c);

        cam_state.replace(CamState::Idle);

        control.gpio_set(0, false).await;
        info!("Listening on TCP:1234...");
        if let Err(e) = socket.accept(1234).await {
            warn!("accept error: {:?}", e);
            continue;
        }

        info!("Received connection from {:?}", socket.remote_endpoint());
        control.gpio_set(0, true).await;

        loop {
            let mut len: [u8; 4] = [0;4];
            if let Err(e) = socket.read_exact(&mut len).await {
                warn!("read err: {e:?}");
                break;
            }
            let len = u32::from_be_bytes(len) as usize;
            let data = &mut buf[0..len];
            if let Err(e) = socket.read_exact(data).await {
                warn!("read err: {e:?}");
                break;
            }
            
            if let Ok(data) = postcard::from_bytes::<ControlPacket>(&data) {
                info!("got {data:?}");
                match data {
                    ControlPacket::Twist(forward, right) => {
                        info!("left to {}", clamp(forward+right, -1., 1.));
                        info!("right to {}", clamp(forward-right, -1., 1.));

                        c.compare_a = calc_speed(forward - right);
                        c.compare_b = calc_speed(forward + right);
                        drive_pwm.set_config(&c);
                    },
                    ControlPacket::Fire => {
                        cam_state.replace(CamState::Fire);
                    },
                    ControlPacket::Stop => {
                        c.compare_a = calc_speed(0.);
                        c.compare_b = calc_speed(0.);
                        drive_pwm.set_config(&c);
                        
                        cam_state.replace(CamState::Idle);
                    },
                    ControlPacket::FireOverride(speed) => {
                        cam_state.replace(CamState::Override(speed));
                    },
                    ControlPacket::Arm(enable) => {
                        if enable {
                            cam_state.replace(CamState::Charging);
                        } else {
                            cam_state.replace(CamState::Idle);
                        }
                    }
                    d => {error!("unimplemented: {d:?}")} // TODO
                }
            }

            if let Ok((data, _ts)) = CHANNEL.try_receive() {
                let telem = TelemetryPacket {
                    sensors: data,
                    cam_state: cam_state.borrow().to_telem(),
                };

                if let Ok(telem) = postcard::to_slice(&telem, &mut telem_buf) {
                    if let Err(e) = socket.write_all(&telem).await {
                        warn!("write error: {:?}", e);
                        break;
                    }
                }
            }
        }
    }
}

/// -1 to 1
fn calc_speed(speed: f32) -> u16 {
    const COUNTS_PER_MS: f32 = 3125.;
    let speed = speed.clamp(-1., 1.);

    let ms = (speed/2.0)+1.5;

    (COUNTS_PER_MS * ms) as u16
}

#[derive(Debug)]
enum CamState {
    Override(f32),
    Charging,
    Charged,
    Fire,
    Idle,
}

impl CamState {
    fn to_telem(&self) -> common::CamState {
        match self {
            Self::Override(_) | Self::Fire => common::CamState::Firing,
            Self::Charging => common::CamState::Charging,
            Self::Charged => common::CamState::Charged,
            Self::Idle => common::CamState::Idle,
        }
    }
}

#[embassy_executor::task]
async fn cam_manager(mut cam: Pwm<'static>, limit: Input<'static>, cam_state: &'static RefCell<CamState>) {
    loop {
        let mut cam_state = cam_state.borrow_mut();
        handle_cam(&mut cam, &mut (*cam_state), &limit);
        drop(cam_state);

        Timer::after_micros(500).await;
    }
}

fn handle_cam(cam: &mut Pwm, state: &mut CamState, limit: &Input) {
    let speed = match state {
        CamState::Override(speed) => {
            *speed
        },
        CamState::Charging => {
            if limit.is_low() {
                *state = CamState::Charged;
                0.0
            } else {
                1.0
            }
        },
        CamState::Fire => {
            if limit.is_high() {
                *state = CamState::Charging;
            }
            1.0
        },
        CamState::Charged | CamState::Idle => 0.0,
    };



    let mut f: pwm::Config = Default::default();
    f.divider = 40.into();
    f.top = 62500; // 20ms
    f.compare_b = calc_speed(-speed); // 1.5ms
    cam.set_config(&f);
}