//! The swarm (oxygenflow#47, docs/jetstream.md §8.1): a runner with a

//! `[runner.swarm]` section serves the image chunks it holds to its peers,

//! and fetches an image it lacks from its peers and the hub together.

//!

//! The server is a small HTTP/1.1 endpoint on the swarm address: a peer

//! shows the ticket its hub signed and gets a session, then asks which

//! chunks of an image this runner has and for chunks by hash. A chunk is

//! served from the chunk cache, or read out of a whole image by its

//! manifest. Every chunk is checked by the side that asked for it, so a

//! peer can only fail to help.

use crate::home::{Home, RunnerConfig};

use crate::images::{cache_path, chunk_cache, manifest_path, save_manifest, write_image, Client, Moved};

use anyhow::{bail, Context, Result};

use flow_core::api::{SwarmAddr, SwarmPeer};

use jet::images::{is_hash, Manifest};

use jet::swarm::{Bits, Pick, Picker, Source, Ticket};

use oxy_core::identity::Identity;

use std::collections::HashMap;

use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, Write};

use std::net::{TcpListener, TcpStream};

use std::path::{Path, PathBuf};

use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};

use std::sync::{Arc, Mutex};

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

/// Connections served at once; more are told 503 and try elsewhere.

const MAX_CONNS: usize = 16;

/// Requests in flight to one peer, and to the hub.

const PEER_WORKERS: usize = 4;

const HUB_WORKERS: usize = 8;

/// How often a fetcher asks its peers what they have.

const REFRESH: Duration = Duration::from_millis(500);

/// Hub failures a fetch survives before it gives up.

const HUB_FAILURES: usize = 5;

/// The largest request body the server reads: a ticket.

const MAX_BODY: usize = 64 * 1024;

/// Names the part files of chunks two sources fetch at once apart.

static PART: AtomicU64 = AtomicU64::new(0);

fn now_ms() -> u64 {

    SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis() as u64).unwrap_or(0)

}

/// An image this runner offers: which of its needed chunks are here

/// (`None`: all of them, the image is whole) and which it is fetching now.

struct Offer {

    bits: Option<Bits>,

    fetching: Bits,

}

struct State {

    root: PathBuf,

    site: String,

    client: Client,

    /// The hub's identity, which signs tickets; fetched when first needed.

    hub: Mutex<Option<Identity>>,

    offers: Mutex<HashMap<String, Offer>>,

    /// Where a chunk lies in a whole image: the image, offset and length.

    located: Mutex<HashMap<String, (PathBuf, u64, u32)>>,

    /// Sessions: token to (peer name, expiry).

    sessions: Mutex<HashMap<String, (String, u64)>>,

    conns: AtomicUsize,

    /// The upload allowance: bytes in hand (negative: owed) and when it

    /// was last topped up.

    bucket: Mutex<(f64, Instant)>,

    rate: f64,

}

/// This runner's place in the swarm.

pub struct Swarm {

    state: Arc<State>,

    pub addr: SwarmAddr,

    peer: ureq::Agent,

}

impl Swarm {

    /// Start serving peers when the runner is configured for it.

    pub fn start(home: &Home, cfg: &RunnerConfig) -> Result<Option<Swarm>> {

        let Some(sc) = &cfg.swarm else { return Ok(None) };

        let listener = TcpListener::bind(&sc.listen).with_context(|| format!("swarm: listen on {}", sc.listen))?;

        let client = Client::new(&cfg.hub, &cfg.token, cfg.ca.as_deref())?;

        let hub = client.hub_identity().map_err(|e| eprintln!("swarm: the hub's key: {e:#}; trying again when a peer comes")).ok();

        let rate = sc.upload_mib_s.unwrap_or(50).max(1) as f64 * (1 << 20) as f64;

        let state = Arc::new(State {

            root: home.root.clone(),

            site: sc.site.clone(),

            client,

            hub: Mutex::new(hub),

            offers: Mutex::new(HashMap::new()),

            located: Mutex::new(HashMap::new()),

            sessions: Mutex::new(HashMap::new()),

            conns: AtomicUsize::new(0),

            bucket: Mutex::new((rate, Instant::now())),

            rate,

        });

        // Whole images: offered from their manifests; one without a

        // manifest (imported here, not fetched) is scanned for one, in the

        // background and one at a time.

        let mut unscanned = Vec::new();

        for h in home.images() {

            match std::fs::read(manifest_path(&home.root, &h)).ok().and_then(|b| serde_json::from_slice::<Manifest>(&b).ok()) {

                Some(m) if m.hash == h => offer_whole(&state, &m, &home.rootfs(&h)),

                _ => unscanned.push((h.clone(), home.rootfs(&h), home.image_arch(&h))),

            }

        }

        let scanning = state.clone();

        std::thread::spawn(move || {

            for (h, path, arch) in unscanned {

                match jet::images::scan_file(&path, arch, |_, _| Ok(())) {

                    Ok(m) if m.hash == h => {

                        save_manifest(&scanning.root, &m).ok();

                        offer_whole(&scanning, &m, &path);

                    }

                    Ok(_) => eprintln!("swarm: {} does not hold the image its name says; not offering it", path.display()),

                    Err(e) => eprintln!("swarm: scan {}: {e}", path.display()),

                }

            }

        });

        let serving = state.clone();

        std::thread::spawn(move || accept(listener, serving));

        let addr = SwarmAddr { addr: sc.advertise.clone().unwrap_or_else(|| sc.listen.clone()), site: sc.site.clone() };

        eprintln!("swarm: serving image chunks on {} (site {:?}, up to {} MiB/s)", sc.listen, sc.site, sc.upload_mib_s.unwrap_or(50));

        let peer = ureq::AgentBuilder::new().timeout_connect(Duration::from_secs(2)).timeout(Duration::from_secs(20)).build();

        Ok(Some(Swarm { state, addr, peer }))

    }

    /// Fetch image `hash` into `home` from the peers and the hub (§8.1):

    /// the rarest chunks from peers first, the hub for what no peer has.

    /// `Ok(None)` when the hub does not hold it.

    pub fn fetch(&self, home: &Home, hash: &str) -> Result<Option<Moved>> {

        let start = Instant::now();

        let client = &self.state.client;

        let Some(m) = client.manifest(hash)? else { return Ok(None) };

        if m.hash != hash {

            bail!("the hub's manifest for {} names {}", &hash[..12], &m.hash[..12]);

        }

        save_manifest(&home.root, &m)?;

        let needed: Vec<String> = m.needed().into_iter().map(String::from).collect();

        let cache = chunk_cache(home);

        let mut have = Bits::new(needed.len());

        for (i, h) in needed.iter().enumerate() {

            if cache_path(&cache, h).is_file() {

                have.set(i);

            }

        }

        self.state.offers.lock().unwrap().insert(hash.into(), Offer { bits: Some(have.clone()), fetching: Bits::new(needed.len()) });

        let announced = client

            .announce(hash, "fetching")

            .map_err(|e| eprintln!("swarm: announce {}: {e:#}; fetching from the hub alone", &hash[..12]))

            .ok();

        let seed = now_ms() ^ (std::process::id() as u64) << 32;

        let picker = Mutex::new(Picker::new(needed.len(), have, seed));

        let peers: Mutex<Vec<Arc<Peer>>> = Mutex::new(Vec::new());

        if let Some(a) = &announced {

            self.join(&a.peers, &a.ticket, hash, needed.len(), &peers, &picker);

        }

        let (bytes, peer_bytes, got) = (AtomicU64::new(0), AtomicU64::new(0), AtomicUsize::new(0));

        let (failed, hub_failures) = (Mutex::new(None::<anyhow::Error>), AtomicUsize::new(0));

        {

            let (picker, peers, needed, cache) = (&picker, &peers, &needed, &cache);

            let (bytes, peer_bytes, got, failed, hub_failures) = (&bytes, &peer_bytes, &got, &failed, &hub_failures);

            let (state, m) = (&self.state, &m);

            let worker = move |src: Source| loop {

                if failed.lock().unwrap().is_some() {

                    break;

                }

                let peer = match src {

                    Source::Peer(p) => Some(peers.lock().unwrap()[p].clone()),

                    Source::Hub => None,

                };

                if peer.as_ref().is_some_and(|p| !p.alive.load(Ordering::Relaxed)) {

                    break;

                }

                // Taken apart from the match, so the lock is not held

                // while waiting.

                let pick = picker.lock().unwrap().pick(src);

                let i = match pick {

                    Pick::Done => break,

                    Pick::Wait => {

                        std::thread::sleep(Duration::from_millis(20));

                        continue;

                    }

                    Pick::Chunk(i) => i,

                };

                if let Some(o) = state.offers.lock().unwrap().get_mut(&m.hash) {

                    o.fetching.set(i);

                }

                let h = &needed[i];

                let fetched = match &peer {

                    Some(p) => self.peer_chunk(p, h),

                    None => client.chunk(h),

                };

                let kept = fetched.and_then(|b| {

                    if blake3::hash(&b).to_hex().as_str() != h {

                        bail!("chunk {} came back with other bytes", &h[..12]);

                    }

                    let p = cache_path(cache, h);

                    std::fs::create_dir_all(p.parent().expect("a chunk directory"))?;

                    let tmp = p.with_extension(format!("part{}", PART.fetch_add(1, Ordering::Relaxed)));

                    std::fs::write(&tmp, &b)?;

                    std::fs::rename(&tmp, &p)?;

                    Ok(b.len() as u64)

                });

                match kept {

                    Ok(n) => {

                        if let Some(Offer { bits: Some(b), fetching }) = state.offers.lock().unwrap().get_mut(&m.hash) {

                            b.set(i);

                            fetching.clear(i);

                        }

                        picker.lock().unwrap().got(i);

                        bytes.fetch_add(n, Ordering::Relaxed);

                        got.fetch_add(1, Ordering::Relaxed);

                        if peer.is_some() {

                            peer_bytes.fetch_add(n, Ordering::Relaxed);

                        }

                    }

                    Err(e) => {

                        picker.lock().unwrap().failed(i);

                        if let Some(o) = state.offers.lock().unwrap().get_mut(&m.hash) {

                            o.fetching.clear(i);

                        }

                        match (&peer, src) {

                            (Some(p), Source::Peer(idx)) => {

                                eprintln!("swarm: peer {} failed ({e:#}); carrying on without it", p.name);

                                p.alive.store(false, Ordering::Relaxed);

                                picker.lock().unwrap().drop_peer(idx);

                                break;

                            }

                            _ => {

                                if hub_failures.fetch_add(1, Ordering::Relaxed) + 1 >= HUB_FAILURES {

                                    *failed.lock().unwrap() = Some(e);

                                    break;

                                }

                                std::thread::sleep(Duration::from_millis(500));

                            }

                        }

                    }

                }

            };

            std::thread::scope(|s| {

                for _ in 0..HUB_WORKERS {

                    s.spawn(move || worker(Source::Hub));

                }

                let mut spawned = 0;

                let mut spawn_new = || {

                    let n = peers.lock().unwrap().len();

                    for p in spawned..n {

                        for _ in 0..PEER_WORKERS {

                            s.spawn(move || worker(Source::Peer(p)));

                        }

                    }

                    spawned = n;

                };

                spawn_new();

                // Keep the peers' maps fresh, and find new peers while

                // fetching, until every chunk is here.

                let mut asked_hub = Instant::now();

                let interval = Duration::from_secs(announced.as_ref().map(|a| a.interval_s).unwrap_or(30).max(5));

                let mut last = Instant::now();

                while picker.lock().unwrap().missing() > 0 && failed.lock().unwrap().is_none() {

                    std::thread::sleep(Duration::from_millis(100));

                    if last.elapsed() < REFRESH {

                        continue;

                    }

                    last = Instant::now();

                    let now: Vec<(usize, Arc<Peer>)> =

                        peers.lock().unwrap().iter().cloned().enumerate().filter(|(_, p)| p.alive.load(Ordering::Relaxed)).collect();

                    for (idx, p) in now {

                        match self.peer_have(&p, hash, needed.len()) {

                            Ok((bits, fetching)) => picker.lock().unwrap().peer_fetching(idx, bits, fetching),

                            Err(_) => {

                                p.alive.store(false, Ordering::Relaxed);

                                picker.lock().unwrap().drop_peer(idx);

                            }

                        }

                    }

                    if announced.is_some() && asked_hub.elapsed() >= interval {

                        asked_hub = Instant::now();

                        if let Ok(a) = client.announce(hash, "fetching") {

                            self.join(&a.peers, &a.ticket, hash, needed.len(), peers, picker);

                            spawn_new();

                        }

                    }

                }

            });

        }

        if let Some(e) = failed.into_inner().unwrap() {

            client.announce(hash, "stopped").ok();

            return Err(e);

        }

        write_image(home, &m)?;

        offer_whole(&self.state, &m, &home.rootfs(hash));

        client.announce(hash, "done").ok();

        Ok(Some(Moved {

            chunks: got.into_inner(),

            bytes: bytes.into_inner(),

            peer_bytes: peer_bytes.into_inner(),

            elapsed: start.elapsed(),

        }))

    }

    /// Admit this runner to the peers it does not know yet: a session from

    /// each, and what each has.

    fn join(&self, found: &[SwarmPeer], ticket: &str, hash: &str, len: usize, peers: &Mutex<Vec<Arc<Peer>>>, picker: &Mutex<Picker>) {

        for sp in found {

            if peers.lock().unwrap().iter().any(|p| p.runner == sp.runner) {

                continue;

            }

            let joined = self.hello(&sp.addr, ticket).and_then(|session| {

                let p = Peer {

                    runner: sp.runner,

                    name: sp.name.clone(),

                    base: format!("http://{}/swarm/v1", sp.addr),

                    session,

                    alive: AtomicBool::new(true),

                };

                let (bits, fetching) = self.peer_have(&p, hash, len)?;

                Ok((p, bits, fetching))

            });

            match joined {

                Ok((p, bits, fetching)) => {

                    let mut list = peers.lock().unwrap();

                    picker.lock().unwrap().peer_fetching(list.len(), bits, fetching);

                    list.push(Arc::new(p));

                }

                Err(e) => eprintln!("swarm: peer {} ({}): {e:#}", sp.name, sp.addr),

            }

        }

    }

    fn hello(&self, addr: &str, ticket: &str) -> Result<String> {

        let url = format!("http://{addr}/swarm/v1/hello");

        let r: serde_json::Value = self.peer.post(&url).send_string(ticket).map_err(|e| peer_error(&url, e))?.into_json()?;

        r["session"].as_str().map(String::from).context("the peer gave no session")

    }

    /// What a peer has of image `hash`, and what it is fetching now:

    /// nothing yet when it has not begun (the tracker may name it a moment

    /// before it starts).

    fn peer_have(&self, p: &Peer, hash: &str, len: usize) -> Result<(Bits, Bits)> {

        let url = format!("{}/images/{hash}/have", p.base);

        let r = match self.peer.get(&url).set("Authorization", &format!("Bearer {}", p.session)).call() {

            Ok(r) => r,

            Err(ureq::Error::Status(404, _)) => return Ok((Bits::new(len), Bits::new(len))),

            Err(e) => return Err(peer_error(&url, e)),

        };

        let mut body = Vec::new();

        r.into_reader().take(len as u64 / 4 + 64).read_to_end(&mut body)?;

        if body == b"complete" {

            return Ok((Bits::full(len), Bits::new(len)));

        }

        let (have, fetching) = body.split_at(len.div_ceil(8).min(body.len()));

        Ok((Bits::from_bytes(len, have), Bits::from_bytes(len, fetching)))

    }

    fn peer_chunk(&self, p: &Peer, hash: &str) -> Result<Vec<u8>> {

        let url = format!("{}/chunks/{hash}", p.base);

        let r = self.peer.get(&url).set("Authorization", &format!("Bearer {}", p.session)).call().map_err(|e| peer_error(&url, e))?;

        let mut out = Vec::new();

        r.into_reader().take(oxy_core::chunk::MAX_SIZE as u64 + 1).read_to_end(&mut out)?;

        Ok(out)

    }

}

/// A peer this fetch asks for chunks.

struct Peer {

    runner: i64,

    name: String,

    base: String,

    session: String,

    alive: AtomicBool,

}

fn peer_error(url: &str, e: ureq::Error) -> anyhow::Error {

    match e {

        ureq::Error::Status(code, r) => anyhow::anyhow!("{url}: HTTP {code}: {}", r.into_string().unwrap_or_default().trim()),

        e => anyhow::anyhow!("{url}: {e}"),

    }

}

/// Offer a whole image: every chunk of it can be read out of `path`.

fn offer_whole(state: &State, m: &Manifest, path: &Path) {

    let mut located = state.located.lock().unwrap();

    let mut at = 0u64;

    for c in &m.chunks {

        if !c.is_zeros() {