#![no_std]
#![no_main]

extern crate alloc;

use core::{net::SocketAddr, pin::pin};
use embassy_executor::Spawner;
use embassy_futures::select::{Either, select};
use embassy_net::{
  Runner, Stack, StackResources,
  dns::DnsQueryType,
  tcp::{TcpSocket, TcpWriter},
  udp::{PacketMetadata, UdpSocket},
};
use embassy_time::{Duration, Timer};
use esp_backtrace as _;
use esp_hal::{
  clock::CpuClock, interrupt::software::SoftwareInterruptControl, timer::timg::TimerGroup,
};
use esp_radio::wifi::{self, Interface, WifiController, sta::StationConfig};
use static_cell::StaticCell;
use wtx::{
  calendar::{Instant, fetch_and_set_epoch_offset},
  collections::Vector,
  http::{HttpRecvParams, MsgBufferString, Protocol, Request, Response, StatusCode},
  http2::{Http2, Http2Buffer, Http2RecvStatus, ServerStream},
  misc::SecretContext,
  rng::{ChaCha20, Rng},
  stream::Stream,
  tls::{CipherSuite, MaxFragmentLength, NamedGroup, TlsAcceptor, TlsConfig, TlsModeVerified},
};
static PUBLIC_KEY: &[u8] = include_bytes!("../.certs/cert.pem");
static SECRET_KEY: &[u8; 119] = include_bytes!("../.certs/key.pem");
static STACK_RESOURCES: StaticCell<StackResources<3>> = StaticCell::new();
const WIFI_PW: &str = env!("WIFI_PW");
const WIFI_SSID: &str = env!("WIFI_SSID");

esp_bootloader_esp_idf::esp_app_desc!();

#[esp_rtos::main]
async fn main(spawner: Spawner) {
  esp_println::logger::init_logger(log::LevelFilter::Debug);
  let (mut rng, stack) = init_peripherals(&spawner).await;
  sync_epoch(stack).await;
  let (mut rx_buffer, mut tx_buffer) = ([0; 4096], [0; 4096]);
  let mut socket = TcpSocket::new(stack, &mut rx_buffer, &mut tx_buffer);
  let tls_config = create_tls_config(&mut rng);
  loop {
    log::debug!("Awaiting for a connection");
    socket.accept(9000).await.unwrap();
    manage_connection(&mut rng, &mut socket, &tls_config).await;
    socket.abort();
  }
}

fn create_tls_config(rng: &mut ChaCha20) -> TlsConfig<TlsModeVerified> {
  let secret_context = SecretContext::new(rng).unwrap();
  let mut tls_config = TlsConfig::from_keys_pem(
    TlsModeVerified::new(),
    PUBLIC_KEY,
    rng,
    (secret_context, &mut SECRET_KEY.clone()),
  )
  .unwrap();
  tls_config
    .alpn_mut()
    .get_or_insert_default()
    .protocol_name_list
    .push(b"h2".try_into().unwrap())
    .unwrap();
  tls_config.cipher_suites_mut().clear();
  tls_config.cipher_suites_mut().push(CipherSuite::Chacha20Poly1305Sha256).unwrap();
  *tls_config.max_fragment_length_mut() = Some(MaxFragmentLength::_2048);
  *tls_config.max_fragment_length_send_mut() = Some(MaxFragmentLength::_2048);
  tls_config.supported_groups_mut().named_group_list.clear();
  tls_config.supported_groups_mut().named_group_list.push(NamedGroup::X25519).unwrap();
  tls_config
}

async fn init_peripherals(spawner: &Spawner) -> (ChaCha20, Stack<'static>) {
  let config = esp_hal::Config::default().with_cpu_clock(CpuClock::max());
  let peripherals = esp_hal::init(config);
  esp_alloc::heap_allocator!(#[esp_hal::ram(reclaimed)] size: 96000);

  let timg0 = TimerGroup::new(peripherals.TIMG0);
  let sw_interrupt = SoftwareInterruptControl::new(peripherals.SW_INTERRUPT);
  esp_rtos::start(timg0.timer0, sw_interrupt.software_interrupt0);

  let esp_rng = esp_hal::rng::Rng::new();
  let mut seed = [0u8; 32];
  for chunk in seed.chunks_mut(4) {
    chunk.copy_from_slice(&esp_rng.random().to_ne_bytes());
  }
  let mut rng = ChaCha20::from_key(seed);

  let (wifi_controller, interfaces) = wifi::new(peripherals.WIFI, <_>::default()).unwrap();
  let (stack, runner) = embassy_net::new(
    interfaces.station,
    embassy_net::Config::dhcpv4(Default::default()),
    STACK_RESOURCES.init(StackResources::new()),
    u64::from_be_bytes(rng.u8_8()),
  );

  spawner.spawn(net_task(runner).unwrap());
  spawner.spawn(wifi_task(wifi_controller).unwrap());

  log::debug!("Waiting for network");
  stack.wait_config_up().await;
  log::debug!("Got IP: {}", stack.config_v4().unwrap().address);

  (rng, stack)
}

async fn manage_connection(
  rng: &mut ChaCha20,
  socket: &mut TcpSocket<'_>,
  tls_config: &TlsConfig<TlsModeVerified>,
) {
  log::debug!("Initiating TLS");
  let tcr = TlsAcceptor::new(&tls_config, &mut *rng, socket).accept().await.unwrap();
  let parts = tcr.tls_stream.into_split().unwrap();

  log::debug!("Initiating HTTP/2");
  let hb = Http2Buffer::new(rng);
  let hrp = HttpRecvParams::with_optioned_params();
  let (frame_reader, http2) = Http2::accept(hb, hrp, parts).await.unwrap();
  let mut frame_reader_pin = pin!(frame_reader);
  let mut stream_pin = pin!(http2.stream(stream_cb));
  loop {
    match select(frame_reader_pin.as_mut(), stream_pin.as_mut()).await {
      Either::First(_) => break,
      Either::Second(rslt) => {
        if !manage_stream(rslt.unwrap()).await {
          break;
        }
        stream_pin.set(http2.stream(stream_cb));
      }
    }
  }
}

async fn manage_stream(tuple: Option<(ServerStream<TcpWriter<'_>, TlsModeVerified>, ())>) -> bool {
  let Some((mut stream, _)) = tuple else {
    return false;
  };
  let (hrs, _req) = stream.recv_req().await.unwrap();
  if let Http2RecvStatus::ClosedConnection | Http2RecvStatus::ClosedStream(_) = hrs {
    return false;
  }
  let res = Response::new(b"Hello World\n", StatusCode::Ok);
  let _ = stream.send_res(&mut Vector::new(), res).await.unwrap();
  true
}

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

fn stream_cb(_: Request<&mut MsgBufferString>, _: Option<Protocol>) {}

async fn sync_epoch(stack: Stack<'_>) {
  log::debug!("Synchronizing epoch");
  let (rx_buffer, rx_meta) = (&mut [0; 128], &mut [PacketMetadata::EMPTY; 4]);
  let (tx_buffer, tx_meta) = (&mut [0; 128], &mut [PacketMetadata::EMPTY; 4]);
  let mut socket = UdpSocket::new(stack, rx_meta, rx_buffer, tx_meta, tx_buffer);
  socket.bind(0).unwrap();
  let addrs = stack.dns_query("pool.ntp.org", DnsQueryType::A).await.unwrap();
  let socket_addr = SocketAddr::new(addrs.into_iter().next().unwrap().into(), 123);
  loop {
    if fetch_and_set_epoch_offset(socket_addr, &mut socket).await.unwrap() {
      break;
    }
  }
  log::debug!("Epoch has been successfully synchronized: {:?}", Instant::now_date_time().unwrap());
}

#[embassy_executor::task]
async fn wifi_task(mut controller: WifiController<'static>) {
  let config = wifi::Config::Station(
    StationConfig::default().with_ssid(WIFI_SSID).with_password(WIFI_PW.into()),
  );
  controller.set_config(&config).unwrap();
  log::debug!("Initializing connecting with WiFi");
  loop {
    match controller.connect_async().await {
      Ok(_) => break,
      Err(err) => {
        log::debug!("Attempt to connect to WiFi resulted in an error: {err:?}");
        Timer::after(Duration::from_millis(2000)).await;
        continue;
      }
    }
  }
  log::debug!("WiFi connection has been successfully established");
  let _ = controller.wait_for_disconnect_async().await;
}
