//! The swarm's tracker (oxygenhub#144, oxygenflow docs/jetstream.md §8.1):
//! runners that share images tell the hub where peers reach them with each
//! claim; a runner fetching an image announces it and is given the peers
//! in its site that hold or are fetching that image, each with the
//! fingerprint of its TLS certificate, and a ticket, signed with the hub
//! key, that those peers accept.
//!
//! What the tracker knows lives in memory: runners repeat their address
//! with every claim and fetchers announce every `INTERVAL_S`, so a
//! restarted hub knows the swarm again within seconds.

use crate::api::ApiError;
use crate::db;
use crate::Shared;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::Json;
use flow_core::api::{Announce, Announced, SwarmAddr, SwarmPeer};
use jet::swarm::{Ticket, TICKET_LIFE_MS};
use std::collections::HashMap;
use std::sync::Mutex;

/// How often a fetcher announces again.
pub const INTERVAL_S: u64 = 15;
/// A runner whose address is older than this is not offered as a peer.
const ADDR_FRESH_MS: i64 = 2 * 60 * 1000;
/// A fetcher that has not announced for this long is no longer fetching.
const FETCHING_FRESH_MS: i64 = 3 * INTERVAL_S as i64 * 1000;
/// A ticket is made again once it has less than this to run.
const TICKET_RENEW_MS: u64 = 5 * 60 * 1000;

#[derive(Default)]
pub struct Tracker {
    /// Runner id to its swarm address and when it last said so.
    addrs: Mutex<HashMap<i64, (SwarmAddr, i64)>>,
    /// Image to the runners fetching it, with when each last announced.
    fetching: Mutex<HashMap<String, HashMap<i64, i64>>>,
    /// Runner id to its current ticket and when that expires.
    tickets: Mutex<HashMap<i64, (String, u64)>>,
}

impl Tracker {
    /// What a runner's claim said: where peers reach it, or that it is not
    /// in a swarm.
    pub fn claimed(&self, runner: i64, swarm: Option<&SwarmAddr>) {
        let mut addrs = self.addrs.lock().unwrap_or_else(|p| p.into_inner());
        // A peer is trusted by its certificate, so one without is not in
        // the swarm (an older flow).
        match swarm.filter(|s| plausible(&s.addr) && jet::images::is_hash(&s.cert)) {
            Some(s) => {
                addrs.insert(runner, (s.clone(), db::now_ms()));
            }
            None => {
                addrs.remove(&runner);
            }
        }
    }

    /// Where peers reach `runner`, if it is in a swarm now.
    pub fn addr(&self, runner: i64) -> Option<SwarmAddr> {
        let addrs = self.addrs.lock().unwrap_or_else(|p| p.into_inner());
        addrs.get(&runner).filter(|(_, seen)| db::now_ms() - seen < ADDR_FRESH_MS).map(|(a, _)| a.clone())
    }

    fn fetch(&self, image: &str, runner: i64, state: &str) {
        let mut fetching = self.fetching.lock().unwrap_or_else(|p| p.into_inner());
        let now = db::now_ms();
        for who in fetching.values_mut() {
            who.retain(|_, seen| now - *seen < FETCHING_FRESH_MS);
        }
        fetching.retain(|_, who| !who.is_empty());
        if state == "fetching" {
            fetching.entry(image.to_string()).or_default().insert(runner, now);
        } else if let Some(who) = fetching.get_mut(image) {
            who.remove(&runner);
        }
    }

    /// Who is fetching which image now, by image.
    pub fn fetching_now(&self) -> HashMap<String, Vec<i64>> {
        let fetching = self.fetching.lock().unwrap_or_else(|p| p.into_inner());
        let now = db::now_ms();
        fetching
            .iter()
            .map(|(image, who)| {
                (image.clone(), who.iter().filter(|(_, seen)| now - **seen < FETCHING_FRESH_MS).map(|(r, _)| *r).collect::<Vec<_>>())
            })
            .filter(|(_, who)| !who.is_empty())
            .collect()
    }

    fn fetchers(&self, image: &str) -> Vec<i64> {
        let fetching = self.fetching.lock().unwrap_or_else(|p| p.into_inner());
        fetching.get(image).map(|who| who.keys().copied().collect()).unwrap_or_default()
    }
}

/// An address a peer could dial: `host:port`, nothing else.
fn plausible(addr: &str) -> bool {
    addr.len() <= 255 && addr.rsplit_once(':').is_some_and(|(host, port)| !host.is_empty() && port.parse::<u16>().is_ok_and(|p| p > 0))
}

/// `POST /flow/swarm/announce`: a runner starting, continuing or ending a
/// fetch. The peers it may ask, in its site, and its ticket.
pub async fn announce(State(app): State<Shared>, headers: HeaderMap, Json(a): Json<Announce>) -> Result<Json<Announced>, ApiError> {
    let me = crate::flow::runner_auth(&app, &headers).await?;
    if !jet::images::is_hash(&a.image) || !matches!(a.state.as_str(), "fetching" | "done" | "stopped") {
        return Err(ApiError(StatusCode::BAD_REQUEST, "announce an image's hash with state fetching, done or stopped".into()));
    }
    let image = a.image.to_ascii_lowercase();
    let Some(mine) = app.swarm.addr(me.id) else {
        return Err(ApiError(
            StatusCode::CONFLICT,
            "this runner has not said where peers reach it: give it a [runner.swarm] section and let it claim".into(),
        ));
    };
    app.swarm.fetch(&image, me.id, &a.state);
    let fetchers = app.swarm.fetchers(&image);
    let mut peers = Vec::new();
    if a.state == "fetching" {
        for r in db::runners::list(&app.db).await? {
            if r.id == me.id || r.state == "disabled" || !r.online() {
                continue;
            }
            let Some(addr) = app.swarm.addr(r.id).filter(|s| s.site == mine.site) else { continue };
            let complete = r.image_list().is_some_and(|held| held.iter().any(|h| h.eq_ignore_ascii_case(&image)));
            if complete || fetchers.contains(&r.id) {
                peers.push(SwarmPeer { runner: r.id, name: r.name.clone(), addr: addr.addr, cert: addr.cert, complete });
            }
        }
    }
    let ticket = ticket(&app, &me, &mine.site)?;
    // An image only runners hold (oxygenflow#94): the hub has its manifest
    // and none of its chunks, so the fetcher asks peers alone.
    let seed = Some(
        sqlx::query_scalar::<_, String>("SELECT source FROM jet_images WHERE hash = $1")
            .bind(&image)
            .fetch_optional(&app.db)
            .await?
            .is_none_or(|s| s != "peer"),
    );
    Ok(Json(Announced { ticket, peers, interval_s: INTERVAL_S, seed }))
}

/// The runner's ticket, made again when it nears its end or its site
/// changed.
fn ticket(app: &Shared, me: &db::runners::Runner, site: &str) -> Result<String, ApiError> {
    let now = db::now_ms() as u64;
    let mut tickets = app.swarm.tickets.lock().unwrap_or_else(|p| p.into_inner());
    if let Some((t, until)) = tickets.get(&me.id) {
        if *until > now + TICKET_RENEW_MS && Ticket::open(t, &app.hub_key.identity, now).is_ok_and(|t| t.site == site) {
            return Ok(t.clone());
        }
    }
    let (key, secret) = crate::keys::load_secret(&app.config.keys_dir())
        .map_err(|e| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("the hub key: {e}")))?;
    let expires_ms = now + TICKET_LIFE_MS;
    let t = Ticket { hub: app.hub_key.key.to_hex(), runner: me.id, name: me.name.clone(), site: site.to_string(), expires_ms }
        .seal(&key, &secret);
    tickets.insert(me.id, (t.clone(), expires_ms));
    Ok(t)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn only_host_and_port_addresses_are_kept() {
        assert!(plausible("10.0.0.5:7401"));
        assert!(plausible("[::1]:7401"));
        assert!(plausible("runner.lan:7401"));
        assert!(!plausible("10.0.0.5"));
        assert!(!plausible(":7401"));
        assert!(!plausible("10.0.0.5:0"));
        assert!(!plausible("10.0.0.5:99999"));
    }
}