mod mdns;
mod noise;
mod rpc;
mod sntp;

use embassy_futures::select::{select, select3};
use embassy_net::{
    tcp::TcpSocket,
    udp::{PacketMetadata, UdpSocket},
};
use embassy_time::Duration;
use sachy_fmt::unwrap;
use sachy_mdns::{
    GROUP_ADDR_V4,
    service::{MdnsService, Service},
};

use crate::{
    constants::{HOST_NAME, HOST_PORT},
    net::noise::NOISE_PROTO,
    rtc::GlobalRtc,
    updates::UpdateConnection,
    utils::try_static_buffer_with,
};

#[embassy_executor::task]
pub async fn udp_stack(stack: embassy_net::Stack<'static>, rtc: GlobalRtc<'static>) {
    let rx_meta = unwrap!(try_static_buffer_with(16, || PacketMetadata::EMPTY));
    let tx_meta = unwrap!(try_static_buffer_with(16, || PacketMetadata::EMPTY));
    let rx_buffer = unwrap!(try_static_buffer_with(4096, Default::default));
    let tx_buffer = unwrap!(try_static_buffer_with(4096, Default::default));

    let mut udp = UdpSocket::new(stack, rx_meta, rx_buffer, tx_meta, tx_buffer);

    stack.wait_config_up().await;

    unwrap!(stack.join_multicast_group(GROUP_ADDR_V4));

    let mut service = MdnsService::new(Service::new(
        "_picostrike._tcp.local",
        HOST_NAME,
        HOST_NAME,
        stack
            .config_v4()
            .map(|config| config.address.address().into()),
        HOST_PORT,
        [NOISE_PROTO],
    ));

    loop {
        stack.wait_config_up().await;

        if !rtc.is_ready().await {
            sntp::sntp_loop(&mut udp, stack, rtc).await;
        }

        select3(
            stack.wait_link_down(),
            mdns::mdns_loop(&mut service, &mut udp),
            rtc.wait_for_reset(),
        )
        .await;

        udp.close();
    }
}

#[embassy_executor::task]
pub async fn tcp_stack(stack: embassy_net::Stack<'static>) {
    let rx_buffer = unwrap!(try_static_buffer_with(4096, Default::default));
    let tx_buffer = unwrap!(try_static_buffer_with(4096, Default::default));
    let mut tcp = TcpSocket::new(stack, rx_buffer, tx_buffer);
    let net_data = UpdateConnection::get_receiver();

    tcp.set_keep_alive(Some(Duration::from_secs(1)));
    tcp.set_timeout(Some(Duration::from_secs(10)));

    loop {
        stack.wait_config_up().await;

        select(stack.wait_link_down(), rpc::rpc_loop(&mut tcp, &net_data)).await;

        net_data.clear();

        UpdateConnection::disconnect();
    }
}