//! Camera firmware for an ESP32-S3-N16R8 with an OV2640 / OV3660 / OV5640 sensor,
//! written in Rust on top of ESP-IDF.
//!
//! Captures JPEG frames and streams them as UDP datagrams to any number of
//! clients. The wire format lives in the `protocol` crate, shared with the desktop
//! client in `client/`; see that crate's documentation for why the transport is
//! UDP rather than TCP.
//!
//! WiFi credentials come from the `WIFI_SSID` / `WIFI_PASS` environment variables
//! at build time, so no secrets end up in the source.

use std::time::{Duration, Instant};

use anyhow::{anyhow, Result};
use embedded_svc::wifi::{
    AccessPointInfo, AuthMethod, ClientConfiguration, Configuration, PmfConfiguration,
};
use esp_idf_svc::hal::peripherals::Peripherals;
use esp_idf_svc::wifi::{BlockingWifi, EspWifi};
use esp_idf_svc::{eventloop::EspSystemEventLoop, nvs::EspDefaultNvsPartition};
use log::{error, info, warn};

mod camera;
mod udp;

use camera::{Camera, Resolution, SensorHint};
use udp::UdpSender;

const SSID: &str = env!("WIFI_SSID");
const PASSWORD: &str = env!("WIFI_PASS");

/// Authentication mode to assume, overriding whatever the scan reports.
///
/// Usually left unset: the scan tells us the AP's real mode. Pin it with
/// `WIFI_AUTH=wpa2` (or `wpa3`, `wpa2wpa3`, `wpa`, `none`) when the scan cannot
/// see the AP - a hidden SSID, for instance.
fn auth_override() -> Option<AuthMethod> {
    let name = option_env!("WIFI_AUTH")?;

    Some(match name {
        "none" => AuthMethod::None,
        "wep" => AuthMethod::WEP,
        "wpa" => AuthMethod::WPA,
        "wpa2" => AuthMethod::WPA2Personal,
        "wpawpa2" => AuthMethod::WPAWPA2Personal,
        "wpa3" => AuthMethod::WPA3Personal,
        "wpa2wpa3" => AuthMethod::WPA2WPA3Personal,
        "wapi" => AuthMethod::WAPIPersonal,
        other => panic!(
            "WIFI_AUTH={other} is not one of: \
             none, wep, wpa, wpa2, wpawpa2, wpa3, wpa2wpa3, wapi"
        ),
    })
}

/// Capture resolution, overridable at build time: `CAM_RES=vga ./scripts/build.sh`
const RESOLUTION: &str = match option_env!("CAM_RES") {
    Some(value) => value,
    None => "vga",
};

/// JPEG quality 0-63 (lower is better), overridable: `CAM_QUALITY=15 ./scripts/build.sh`
const QUALITY: &str = match option_env!("CAM_QUALITY") {
    Some(value) => value,
    None => "12",
};

/// Whether the driver inverts VSYNC, overridable so polarity can be A/B tested
/// without editing code: `CAM_VSYNC_INVERT=0 ./scripts/build.sh`.
///
/// A function rather than a `const` because matching on a `str` is not yet
/// allowed in const context on this toolchain.
///
/// `true` is esp32-camera's default and correct for most DVP sensors. `false` is
/// what a board that inverts VSYNC in hardware needs.
fn vsync_invert() -> bool {
    !matches!(option_env!("CAM_VSYNC_INVERT"), Some("0"))
}

fn main() -> Result<()> {
    // Must run once so the esp-idf-sys runtime patches link correctly.
    esp_idf_svc::sys::link_patches();
    esp_idf_svc::log::EspLogger::initialize_default();

    let resolution: Resolution = RESOLUTION.parse()?;
    let quality: u8 = QUALITY.parse()?;

    info!("vsync invert: {}", vsync_invert());

    let camera = Camera::init(SensorHint::Auto, resolution, quality, vsync_invert())?;

    let peripherals = Peripherals::take()?;
    let sys_loop = EspSystemEventLoop::take()?;
    let nvs = EspDefaultNvsPartition::take()?;

    let mut wifi = BlockingWifi::wrap(
        EspWifi::new(peripherals.modem, sys_loop.clone(), Some(nvs))?,
        sys_loop,
    )?;

    // `connect` starts the driver, then scans and associates, so that an SSID on
    // a 5 GHz network (which this radio cannot see at all) is reported as such
    // instead of a bare association timeout.
    if let Err(err) = connect(&mut wifi) {
        error!("could not join \"{SSID}\": {err}");
        error!("check that the SSID is a 2.4 GHz network and the password is right");
        return Err(err);
    }

    let ip = wifi.wifi().sta_netif().get_ip_info()?.ip;

    // Modem power save trades latency for current, and for a video stream that
    // trade is backwards: with PS on, an idle beacon interval costs hundreds of
    // milliseconds, which shows up as stutter and a ping of 200-300 ms on an
    // otherwise good link. The board is mains-powered anyway.
    //
    // This one binds as a bare `esp_err_t` rather than a `Result`.
    let ps_rc = unsafe {
        esp_idf_svc::sys::esp_wifi_set_ps(esp_idf_svc::sys::wifi_ps_type_t_WIFI_PS_NONE)
    };

    if ps_rc == 0 {
        info!("WiFi power save disabled");
    } else {
        warn!("could not disable WiFi power save (esp_err 0x{ps_rc:x})");
    }

    // The board identifies itself by an id derived from its MAC, not its IP: DHCP
    // hands out a different address across reboots, and a client that saved its
    // binding to one would have to be reconfigured every time.
    let sender_id = udp::local_sender_id();
    let mut sender = UdpSender::new(sender_id)?;

    info!("ESP32 camera is live on {ip}");
    info!("clients will find it by broadcast; no configuration needed on this side");

    let mut last_report = Instant::now();
    let mut last_stats = sender.stats();

    loop {
        udp::pump(&camera, &mut sender)?;

        let now = Instant::now();

        if now.duration_since(last_report) >= Duration::from_secs(5) {
            // Deltas rather than totals: the totals alone cannot distinguish "the
            // sensor is slow" from "the sender is not keeping up", which is the
            // distinction that matters when the frame rate is wrong.
            let s = sender.stats();
            let elapsed = now.duration_since(last_report).as_secs_f32();

            info!(
                "{:.1} fps, {} datagrams/s, {:.1} kB/s, {} send errors, \
                 {} clients subscribed, {} frames skipped (no client)",
                (s.frames_sent - last_stats.frames_sent) as f32 / elapsed,
                (s.fragments_sent - last_stats.fragments_sent) as f32 / elapsed,
                (s.bytes_sent - last_stats.bytes_sent) as f32 / elapsed / 1024.0,
                s.errors - last_stats.errors,
                sender.peer_count(),
                s.skipped_no_peer - last_stats.skipped_no_peer,
            );

            last_stats = s;
            last_report = now;
        }
    }
}

/// Joins `SSID`, scanning first and retrying the plausible authentication modes.
///
/// The scan is not decoration. This radio is 2.4 GHz only, so a 5 GHz SSID never
/// shows up at all and a bare "association timed out" tells you nothing.
///
/// The retries are because the AP's advertised mode is not the same thing as the
/// mode it will accept. Routers commonly run WPA/WPA2 mixed, WPA2/WPA3 mixed, or
/// WPA3 with PMF mandatory, and a client pinned to the wrong one fails the 4-way
/// handshake *after* it has already associated - the failure looks identical to
/// a wrong password, so it is worth walking the list instead of guessing.
fn connect(wifi: &mut BlockingWifi<EspWifi<'static>>) -> Result<()> {
    let pinned_auth = auth_override();

    // A configuration has to exist before start(), which is what decides
    // station vs AP mode; the real one is picked after the scan.
    wifi.set_configuration(&client_config(pinned_auth.unwrap_or_else(default_auth_method))?)?;
    wifi.start()?;
    info!("wifi started");

    info!("scanning for 2.4 GHz networks...");
    let advertised = match wifi.wifi_mut().scan() {
        Ok(aps) => {
            report_scan(&aps);
            advertised_auth(&aps)
        }
        Err(err) => {
            warn!("scan failed: {err}");
            None
        }
    };

    for auth in candidates(advertised, pinned_auth) {
        info!("trying {auth:?}...");

        wifi.set_configuration(&client_config(auth)?)?;
        // Driver-level (non-blocking) connect: `BlockingWifi::connect()` waits
        // out the driver's own ~20 s association timeout, which would make
        // walking the list of modes take minutes.
        wifi.wifi_mut().connect()?;

        // Bounded, so one bad mode costs a few seconds instead of the timeout.
        let associated = wifi
            .wifi_wait_while(
                || wifi.wifi().is_connected().map(|c| !c),
                Some(Duration::from_secs(5)),
            )
            .is_ok();

        let got_ip = associated
            && wifi
                .ip_wait_while(
                    || wifi.wifi().is_up().map(|up| !up),
                    Some(Duration::from_secs(4)),
                )
                .is_ok();

        if got_ip {
            let ip_info = wifi.wifi().sta_netif().get_ip_info()?;
            if !ip_info.ip.is_unspecified() {
                info!("wifi connected with {auth:?}: {ip_info:?}");
                return Ok(());
            }
        }

        warn!(
            "{auth:?} did not complete the handshake ({})",
            if associated { "associated but no DHCP address" } else { "no association" }
        );
        wifi.wifi_mut().disconnect().ok();
    }

    Err(anyhow!(
        "no authentication mode joined \"{SSID}\" (tried {:?}); \
         association succeeded but the 4-way handshake did not, \
         which almost always means the password is wrong",
        candidates(advertised, pinned_auth)
    ))
}

/// Authentication modes to try, most likely first.
///
/// A mode pinned with `WIFI_AUTH` is the only one tried - if it does not work,
/// trying others would just be guessing at a value the builder already stated.
fn candidates(advertised: Option<AuthMethod>, pinned: Option<AuthMethod>) -> Vec<AuthMethod> {
    if let Some(auth) = pinned {
        return vec![auth];
    }

    // "WPA" proper is a pre-2004 mode no current router speaks, and "None" is
    // only meaningful without a password, so neither is worth a 10 s attempt.
    let mut modes = vec![advertised.unwrap_or_else(default_auth_method)];
    modes.extend([
        AuthMethod::WPA2Personal,
        AuthMethod::WPA2WPA3Personal,
        AuthMethod::WPA3Personal,
    ]);

    if PASSWORD.is_empty() {
        modes.push(AuthMethod::None);
    }

    let mut seen = Vec::new();
    modes.retain(|m| {
        let fresh = !seen.contains(m);
        seen.push(*m);
        fresh
    });
    modes
}

fn client_config(auth_method: AuthMethod) -> Result<Configuration> {
    Ok(Configuration::Client(ClientConfiguration {
        ssid: SSID
            .try_into()
            .map_err(|_| anyhow!("WIFI_SSID too long (max 32 bytes)"))?,
        bssid: None,
        auth_method,
        password: PASSWORD
            .try_into()
            .map_err(|_| anyhow!("WIFI_PASS too long (max 64 bytes)"))?,
        channel: None,
        // Advertising PMF keeps routers that offer WPA2+WPA3 mixed mode from
        // dropping the handshake, while `required: false` still lets a WPA2-only
        // AP through. Leaving this at NotCapable fails on the former.
        pmf_cfg: PmfConfiguration::Capable { required: false },
        ..Default::default()
    }))
}

fn default_auth_method() -> AuthMethod {
    if PASSWORD.is_empty() {
        AuthMethod::None
    } else {
        AuthMethod::WPA2Personal
    }
}

/// What `SSID` advertises, if the scan saw it.
///
/// `None` means the scan did not find this network at all, which on its own is
/// the most useful diagnostic there is: on a 2.4 GHz-only radio it almost always
/// means the SSID belongs to a 5 GHz network, or the AP is on a channel this
/// region's regulatory domain forbids.
fn advertised_auth(aps: &[AccessPointInfo]) -> Option<AuthMethod> {
    aps.iter()
        .find(|ap| ap.ssid.as_str() == SSID)
        .and_then(|ap| ap.auth_method)
}

/// Logs the visible networks, marking the one we are trying to join.
fn report_scan(aps: &[AccessPointInfo]) {
    if aps.is_empty() {
        warn!("scan found no networks at all");
        return;
    }

    for ap in aps {
        info!(
            "  {:<32} ch {:>2}  {:>4} dBm  {:?}{}",
            ap.ssid,
            ap.channel,
            ap.signal_strength,
            ap.auth_method,
            if ap.ssid.as_str() == SSID { "  <- target" } else { "" }
        );
    }

    if advertised_auth(aps).is_none() {
        error!("\"{SSID}\" was not in the scan results");
        error!("this radio is 2.4 GHz only, so the AP is likely 5 GHz or on a forbidden channel");
    }
}