use defmt::{error, info};
use embassy_executor::Spawner;
use embassy_futures::select::{Either, select};
use embassy_net::{Runner, Stack};
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_time::{Duration, Timer};
use embedded_rpc::{RequestDroppedError, RpcService};
use esp_hal::peripherals::WIFI;
use esp_radio::wifi::sta::StationConfig;
use esp_radio::wifi::{ControllerConfig, Interface, WifiController};

pub enum WifiCommand {
    SetStationConfig(esp_radio::wifi::Config),
    SetDhcpConfig(embassy_net::DhcpConfig),
}

impl WifiCommand {
    pub async fn set_station_config(
        config: esp_radio::wifi::Config,
    ) -> Result<(), RequestDroppedError> {
        WIFI_SERVICE
            .request(WifiCommand::SetStationConfig(config))
            .await
    }
}

pub static WIFI_SERVICE: RpcService<CriticalSectionRawMutex, WifiCommand, ()> = RpcService::new();

pub fn start_wifi(spawner: Spawner, wifi: WIFI<'static>) -> Stack<'static> {
    let wifi_config = esp_radio::wifi::Config::Station(StationConfig::default());

    let (wifi_controller, wifi_interfaces) = esp_radio::wifi::new(
        wifi,
        ControllerConfig::default().with_initial_config(wifi_config),
    )
    .expect("Failed to initialize Wi-Fi controller");

    let mut dhcp_config = embassy_net::DhcpConfig::default();
    dhcp_config.hostname = Some("cover-theif".try_into().unwrap());
    let net_config = embassy_net::Config::dhcpv4(dhcp_config);

    let rng = esp_hal::rng::Rng::new();
    let seed = (rng.random() as u64) << 32 | rng.random() as u64;

    let (net_stack, net_runner) = embassy_net::new(
        wifi_interfaces.station,
        net_config,
        crate::mk_static!(
            embassy_net::StackResources<3>,
            embassy_net::StackResources::<3>::new()
        ),
        seed,
    );

    spawner.spawn(net_task(net_runner).unwrap());
    spawner.spawn(wifi_handler(wifi_controller).unwrap());
    spawner.spawn(ip_logger(net_stack).unwrap());
    spawner.spawn(handler(net_stack).unwrap());

    net_stack
}

#[embassy_executor::task]
async fn handler(net_stack: Stack<'static>) {
    loop {
        let (req, served) = WIFI_SERVICE.serve().await;

        match req {
            WifiCommand::SetStationConfig(config) => {
                if WIFI_HANDLER_RPC.request(config).await.is_err() {
                    continue;
                }
            }
            WifiCommand::SetDhcpConfig(dhcp_config) => {
                net_stack.set_config_v4(embassy_net::ConfigV4::Dhcp(dhcp_config));
            }
        }

        served.respond(());
    }
}

static WIFI_HANDLER_RPC: RpcService<CriticalSectionRawMutex, esp_radio::wifi::Config, ()> =
    RpcService::new();

#[embassy_executor::task]
async fn wifi_handler(mut controller: WifiController<'static>) {
    let connection = async |controller: &mut WifiController<'static>| {
        info!("Connecting to WiFi...");

        match controller.connect_async().await {
            Ok(info) => {
                info!("Wifi connected to {:?}", info);

                // wait until we're no longer connected
                let info = controller.wait_for_disconnect_async().await.ok();
                info!("Disconnected: {:?}", info);
            }
            Err(e) => {
                error!("Failed to connect to wifi: {:?}", e);
            }
        }
    };

    loop {
        if let Either::First((req, served)) =
            select(WIFI_HANDLER_RPC.serve(), connection(&mut controller)).await
        {
            if controller.set_config(&req).is_ok() {
                served.respond(());
            }
        }
    }
}

#[embassy_executor::task]
async fn ip_logger(net_stack: Stack<'static>) {
    loop {
        info!("Getting IP address...");
        net_stack.wait_config_up().await;

        match net_stack.config_v4() {
            Some(info) => {
                info!("Got IP: {:?}", info);

                // wait until we're no longer connected
                net_stack.wait_config_down().await;
                info!("Address no longer configured");
            }
            None => {
                error!("Failed to configure an address");
            }
        }

        Timer::after(Duration::from_millis(5000)).await
    }
}

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