#![no_std]
#![no_main]
#![allow(async_fn_in_trait)]
#![feature(impl_trait_in_assoc_type)]
#![feature(slice_split_once)]

use cyw43_pio::{DEFAULT_CLOCK_DIVIDER, PioSpi};
use embassy_executor::Spawner;
use embassy_net::{Config, StackResources};
use embassy_rp::bind_interrupts;
use embassy_rp::clocks::RoscRng;
use embassy_rp::gpio::{Level, Output};
use embassy_rp::peripherals::USB;
use embassy_rp::peripherals::{DMA_CH0, PIO0};
use embassy_rp::pio::{InterruptHandler as PioInterruptHandler, Pio};
use embassy_rp::usb::{Driver, InterruptHandler as UsbInterruptHandler};
use pico_website::unwrap;
use rand_core::RngCore;
use static_cell::StaticCell;
use {defmt_rtt as _, panic_probe as _};

mod dhcp;
mod game;
mod socket;

bind_interrupts!(struct Irqs {
    USBCTRL_IRQ => UsbInterruptHandler<USB>;
    PIO0_IRQ_0 => PioInterruptHandler<PIO0>;
});

#[embassy_executor::task]
async fn logger_task(driver: Driver<'static, USB>) {
    embassy_usb_logger::run!(1024, log::LevelFilter::Info, 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
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    let p = embassy_rp::init(Default::default());
    let driver = Driver::new(p.USB, Irqs);
    spawner.spawn(logger_task(driver)).unwrap();
    let mut rng = RoscRng;

    let fw = include_bytes!("../cyw43-firmware/43439A0.bin");
    let clm = include_bytes!("../cyw43-firmware/43439A0_clm.bin");
    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;
    unwrap(spawner.spawn(cyw43_task(runner))).await;

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

    #[cfg(not(feature = "wifi-connect"))]
    // 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(192, 254, 0, 2), 24),
        dns_servers: heapless::Vec::new(),
        gateway: None,
    });

    #[cfg(feature = "wifi-connect")]
    let (wifi_conf, config) = {
        let wifi_conf = unwrap(serde_json_core::from_slice::<WifiConnectConf>(
            include_bytes!("../wifi.json"),
        ))
        .await
        .0;
        // Use a link-local address for communication without DHCP server
        // let config = Config::dhcpv4(embassy_net::DhcpConfig::default());
        let config = match wifi_conf.ip {
            Some(ip) => Config::ipv4_static(embassy_net::StaticConfigV4 {
                address: embassy_net::Ipv4Cidr::new(ip, 24),
                dns_servers: heapless::Vec::new(),
                gateway: None,
            }),
            None => Config::dhcpv4(DhcpConfig::default()),
        };
        (wifi_conf, config)
    };

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

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

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

    #[cfg(not(feature = "wifi-connect"))]
    //control.start_ap_open("cyw43", 5).await;
    control.start_ap_wpa2("cyw43", "password", 5).await;

    #[cfg(feature = "wifi-connect")]
    {
        loop {
            match control
                .join(
                    wifi_conf.name,
                    cyw43::JoinOptions::new(wifi_conf.password.as_bytes()),
                )
                .await
            {
                Ok(_) => break,
                Err(err) => {
                    info!("join failed with status={}", err.status);
                }
            }
        }
        info!("Network joined!");
        info!("waiting for link...");
        stack.wait_link_up().await;
        // Wait for DHCP, not necessary when using static IP
        info!("waiting for DHCP...");
        stack.wait_config_up().await;
        // while !stack.is_config_up() {
        //     Timer::after_millis(100).await;
        // }
        info!("DHCP is now up!");
        info!(
            "ip : {}",
            unwrap(stack.config_v4().ok_or("no dhcp config"))
                .await
                .address
        )
    }

    unwrap(spawner.spawn(dhcp::dhcp_server(stack))).await;

    unwrap(spawner.spawn(socket::listen_task(stack, game::Team::Zero, 80))).await;
    unwrap(spawner.spawn(socket::listen_task(stack, game::Team::One, 81))).await;
}

#[cfg(feature = "wifi-connect")]
#[derive(serde::Deserialize)]
struct WifiConnectConf<'a> {
    name: &'a str,
    password: &'a str,
    ip: Option<Ipv4Addr>,
}