//! The swarm (oxygenflow#47, docs/jetstream.md §8.1): runners fetch image

//! chunks from each other as well as from the hub.

//!

//! This module holds what both ends agree on and what can be tested

//! without a network: the have-map a peer offers ([`Bits`]), the ticket the

//! hub signs to admit a runner to its peers ([`Ticket`]), and the order in

//! which a fetcher asks its sources for chunks ([`Picker`]): rarest first,

//! the hub only for what no peer has, and an endgame at the close.

use oxy_core::hash::KeyId;

use oxy_core::identity::{verify_challenge, Identity, SecretKey, Signature};

use serde::{Deserialize, Serialize};

use std::collections::HashMap;

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

/// How long a ticket admits a runner to its peers.

pub const TICKET_LIFE_MS: u64 = 15 * 60 * 1000;

/// Peers make no progress for this long and the hub is asked for anything.

pub const STALL: Duration = Duration::from_secs(2);

/// Most sources asked for the same chunk in the endgame.

const ENDGAME_COPIES: u8 = 2;

/// One bit per chunk of a manifest's needed list, in its order.

#[derive(Clone, Debug, Default, PartialEq, Eq)]

pub struct Bits {

    words: Vec<u64>,

    len: usize,

}

impl Bits {

    pub fn new(len: usize) -> Bits {

        Bits { words: vec![0; len.div_ceil(64)], len }

    }

    /// Every bit set: a peer that holds the whole image.

    pub fn full(len: usize) -> Bits {

        let mut b = Bits::new(len);

        (0..len).for_each(|i| b.set(i));

        b

    }

    pub fn len(&self) -> usize {

        self.len

    }

    pub fn is_empty(&self) -> bool {

        self.len == 0

    }

    pub fn get(&self, i: usize) -> bool {

        i < self.len && self.words[i / 64] & (1 << (i % 64)) != 0

    }

    pub fn set(&mut self, i: usize) {

        if i < self.len {

            self.words[i / 64] |= 1 << (i % 64);

        }

    }

    pub fn clear(&mut self, i: usize) {

        if i < self.len {

            self.words[i / 64] &= !(1 << (i % 64));

        }

    }

    pub fn count(&self) -> usize {

        self.words.iter().map(|w| w.count_ones() as usize).sum()

    }

    /// Little-endian bytes, as `GET …/have` sends them.

    pub fn to_bytes(&self) -> Vec<u8> {

        let mut out: Vec<u8> = self.words.iter().flat_map(|w| w.to_le_bytes()).collect();

        out.truncate(self.len.div_ceil(8));

        out

    }

    /// The bits of a `have` answer for a list of `len` chunks; bits past

    /// `len` are ignored, and a short answer means the rest are unset.

    pub fn from_bytes(len: usize, bytes: &[u8]) -> Bits {

        let mut b = Bits::new(len);

        for (i, byte) in bytes.iter().enumerate().take(len.div_ceil(8)) {

            b.words[i / 8] |= (*byte as u64) << (8 * (i % 8));

        }

        if let Some(last) = b.words.last_mut() {

            if !len.is_multiple_of(64) {

                *last &= (1u64 << (len % 64)) - 1;

            }

        }

        b

    }

}

/// What the hub signs to admit a runner to its peers.

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]

pub struct Ticket {

    /// The hub's key, hex: a peer accepts only its own hub's tickets.

    pub hub: String,

    pub runner: i64,

    pub name: String,

    pub site: String,

    pub expires_ms: u64,

}

impl Ticket {

    /// The ticket as text: its JSON and the hub's signature over it, hex,

    /// joined by a dot.

    pub fn seal(&self, key: &KeyId, secret: &SecretKey) -> String {

        let body = serde_json::to_vec(self).expect("a ticket serialises");

        let sig = secret.sign_challenge(key, &body).expect("a ticket is longer than 16 bytes");

        format!("{}.{}", hex::encode(&body), hex::encode(Signature::encode_list(&[sig])))

    }

    /// Read and check a ticket: signed by `hub`, not expired at `now_ms`.

    pub fn open(text: &str, hub: &Identity, now_ms: u64) -> Result<Ticket, String> {

        let (body, sig) = text.trim().split_once('.').ok_or("not a ticket")?;

        let body = hex::decode(body).map_err(|_| "not a ticket")?;

        let sigs = hex::decode(sig).ok().and_then(|b| Signature::decode_list(&b).ok()).ok_or("not a ticket's signature")?;

        let sig = sigs.first().ok_or("an unsigned ticket")?;

        verify_challenge(hub, sig, &body).map_err(|_| "the ticket is not signed by this runner's hub")?;

        let t: Ticket = serde_json::from_slice(&body).map_err(|_| "not a ticket")?;

        if t.hub != hub.hash().to_hex() {

            return Err("the ticket names another hub".into());

        }

        if t.expires_ms <= now_ms {

            return Err("the ticket has expired".into());

        }

        Ok(t)

    }

}

/// Where a chunk comes from: the hub, or a peer by its index.

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]

pub enum Source {

    Hub,

    Peer(usize),

}

/// What a source should do next.

#[derive(Debug, PartialEq, Eq)]

pub enum Pick {

    /// Ask for this chunk (an index into the needed list).

    Chunk(usize),

    /// Nothing for this source now; ask again shortly.

    Wait,

    /// Every chunk is here.

    Done,

}

/// The order a fetcher asks its sources for chunks (docs/jetstream.md

/// §8.1): the rarest chunk a peer holds first; the hub for chunks no peer

/// holds, or for anything once the peers stall; and in the endgame, a

/// second source for a chunk already asked for.

pub struct Picker {

    have: Bits,

    asked: Vec<u8>,

    /// The source first asked for each chunk in flight: the endgame asks

    /// a different one.

    first: Vec<Option<Source>>,

    /// Each peer's map: what it has, and what it is fetching now.

    peers: HashMap<usize, (Bits, Bits)>,

    /// Missing chunks, rarest first; rebuilt when the peers' maps change.

    order: Vec<usize>,

    /// How many peers hold each chunk, as of the last rebuild.

    holders: Vec<u16>,

    /// Whether some peer is fetching each chunk now: the hub is not asked

    /// for it, the peer will have it shortly.

    coming: Vec<bool>,

    progress: Instant,

    seed: u64,

}

impl Picker {

    /// A picker for `len` chunks, of which `have` are already here.

    pub fn new(len: usize, have: Bits, seed: u64) -> Picker {

        let mut p = Picker {

            have,

            asked: vec![0; len],

            first: vec![None; len],

            peers: HashMap::new(),

            order: Vec::new(),

            holders: vec![0; len],

            coming: vec![false; len],

            progress: Instant::now(),

            seed: seed | 1,

        };

        p.reorder();

        p

    }

    pub fn len(&self) -> usize {

        self.asked.len()

    }

    pub fn is_empty(&self) -> bool {

        self.asked.is_empty()

    }

    pub fn missing(&self) -> usize {

        self.len() - self.have.count()

    }

    pub fn have(&self) -> &Bits {

        &self.have

    }

    /// A peer's have-map, new or refreshed.

    pub fn peer(&mut self, peer: usize, bits: Bits) {

        let none = Bits::new(bits.len());

        self.peer_fetching(peer, bits, none);

    }

    /// A peer's have-map and the chunks it is fetching now.

    pub fn peer_fetching(&mut self, peer: usize, have: Bits, fetching: Bits) {

        self.peers.insert(peer, (have, fetching));

        self.reorder();

    }

    /// A peer that failed: it is asked for nothing more.

    pub fn drop_peer(&mut self, peer: usize) {

        self.peers.remove(&peer);

        self.reorder();

    }

    fn next_rand(&mut self) -> u64 {

        // xorshift64: enough to spread runners that start together.

        self.seed ^= self.seed << 13;

        self.seed ^= self.seed >> 7;

        self.seed ^= self.seed << 17;

        self.seed

    }

    fn reorder(&mut self) {

        let mut keyed: Vec<(u16, u64, usize)> = Vec::new();

        for i in 0..self.len() {

            self.holders[i] = self.peers.values().filter(|(b, _)| b.get(i)).count().min(u16::MAX as usize) as u16;

            self.coming[i] = self.peers.values().any(|(_, f)| f.get(i));

            if !self.have.get(i) {

                let r = self.next_rand();

                keyed.push((self.holders[i], r, i));

            }

        }

        keyed.sort_unstable();

        self.order = keyed.into_iter().map(|(_, _, i)| i).collect();

    }

    /// The next chunk for `src` to fetch.

    pub fn pick(&mut self, src: Source) -> Pick {

        if self.missing() == 0 {

            return Pick::Done;

        }

        let holds = |i: usize| match src {

            Source::Hub => true,

            Source::Peer(p) => self.peers.get(&p).is_some_and(|(b, _)| b.get(i)),

        };

        let stalled = self.progress.elapsed() >= STALL;

        // Not yet asked for: a peer takes the rarest it holds; the hub

        // takes what no peer holds, or anything once the peers stall.

        let fresh = self.order.iter().copied().find(|&i| {

            !self.have.get(i)

                && self.asked[i] == 0

                && holds(i)

                && (src != Source::Hub || stalled || self.peers.is_empty() || (self.holders[i] == 0 && !self.coming[i]))

        });

        if let Some(i) = fresh {

            self.asked[i] += 1;

            self.first[i] = Some(src);

            return Pick::Chunk(i);

        }

        // The endgame: everything missing is asked for already. The hub

        // joins it only once the peers stall.

        if src == Source::Hub && !stalled && !self.peers.is_empty() {

            return Pick::Wait;

        }

        let mut left = self.order.iter().copied().filter(|&i| !self.have.get(i));

        if left.clone().all(|i| self.asked[i] > 0) {

            if let Some(i) = left.find(|&i| self.asked[i] < ENDGAME_COPIES && self.first[i] != Some(src) && holds(i)) {

                self.asked[i] += 1;

                return Pick::Chunk(i);

            }

        }

        Pick::Wait

    }

    /// Chunk `i` arrived and was checked.

    pub fn got(&mut self, i: usize) {

        self.have.set(i);

        self.asked[i] = self.asked[i].saturating_sub(1);

        self.first[i] = None;

        self.progress = Instant::now();

    }

    /// Chunk `i` did not come: it can be asked for again.

    pub fn failed(&mut self, i: usize) {

        self.asked[i] = self.asked[i].saturating_sub(1);

        if self.asked[i] == 0 {

            self.first[i] = None;

        }

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn bits_travel_as_bytes() {

        let mut b = Bits::new(70);

        [0, 9, 63, 64, 69].iter().for_each(|&i| b.set(i));

        let back = Bits::from_bytes(70, &b.to_bytes());

        assert_eq!(back, b);

        assert_eq!(b.to_bytes().len(), 9);

        assert_eq!(Bits::from_bytes(10, &[0xff, 0xff, 0xff]).count(), 10, "bits past the list are ignored");

        assert_eq!(Bits::from_bytes(10, &[]).count(), 0, "a short answer leaves the rest unset");

        assert_eq!(Bits::full(130).count(), 130);

    }

    #[test]

    fn a_ticket_opens_only_with_its_hub_and_in_time() {

        let sk = SecretKey::generate();

        let hub = sk.identity("hub", "hub@example", 1);

        let t = Ticket { hub: hub.hash().to_hex(), runner: 7, name: "earth-1".into(), site: "home".into(), expires_ms: 10_000 };

        let text = t.seal(&hub.hash(), &sk);

        assert_eq!(Ticket::open(&text, &hub, 9_999).unwrap(), t);

        assert!(Ticket::open(&text, &hub, 10_000).unwrap_err().contains("expired"));

        let other_sk = SecretKey::generate();

        let other = other_sk.identity("other", "o@example", 1);

        assert!(Ticket::open(&text, &other, 1).is_err(), "another hub's key");

        let mut forged = text.clone();

        forged.replace_range(2..4, "ff");

        assert!(Ticket::open(&forged, &hub, 1).is_err(), "a changed body");

        assert!(Ticket::open("nonsense", &hub, 1).is_err());

    }

    #[test]

    fn peers_are_asked_for_the_rarest_chunk_they_hold() {

        let mut p = Picker::new(4, Bits::new(4), 42);

        let mut a = Bits::new(4);

        [0, 1, 2].iter().for_each(|&i| a.set(i));

        let mut b = Bits::new(4);

        [0, 1].iter().for_each(|&i| b.set(i));

        p.peer(0, a);

        p.peer(1, b);

        // Chunk 2 only peer 0 holds; chunk 3 nobody does.

        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(2));

        assert_eq!(p.pick(Source::Hub), Pick::Chunk(3), "the hub seeds what no peer holds");

        assert_eq!(p.pick(Source::Hub), Pick::Wait, "and leaves the rest to the peers");

        let first = p.pick(Source::Peer(1));

        assert!(matches!(first, Pick::Chunk(0) | Pick::Chunk(1)));

    }

    #[test]

    fn the_hub_is_not_asked_for_what_a_peer_is_fetching() {

        let mut p = Picker::new(2, Bits::new(2), 4);

        let mut fetching = Bits::new(2);

        fetching.set(0);

        p.peer_fetching(0, Bits::new(2), fetching);

        assert_eq!(p.pick(Source::Hub), Pick::Chunk(1));

        assert_eq!(p.pick(Source::Hub), Pick::Wait, "chunk 0 is coming from the peer");

        let mut have = Bits::new(2);

        have.set(0);

        p.peer(0, have);

        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(0));

    }

    #[test]

    fn without_peers_the_hub_serves_everything() {

        let mut p = Picker::new(3, Bits::new(3), 1);

        let mut seen = Vec::new();

        while let Pick::Chunk(i) = p.pick(Source::Hub) {

            seen.push(i);

            p.got(i);

        }

        seen.sort();

        assert_eq!(seen, vec![0, 1, 2]);

        assert_eq!(p.pick(Source::Hub), Pick::Done);

    }

    #[test]

    fn the_endgame_never_asks_the_same_source_twice() {

        let mut p = Picker::new(1, Bits::new(1), 9);

        assert_eq!(p.pick(Source::Hub), Pick::Chunk(0));

        assert_eq!(p.pick(Source::Hub), Pick::Wait);

        p.peer(0, Bits::full(1));

        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(0), "a peer may race the hub for it");

        assert_eq!(p.pick(Source::Hub), Pick::Wait, "but not the hub a peer, while the peers make progress");

    }

    #[test]

    fn the_endgame_asks_a_second_source_and_a_failure_frees_the_chunk() {

        let mut p = Picker::new(1, Bits::new(1), 3);

        p.peer(0, Bits::full(1));

        p.peer(1, Bits::full(1));

        assert_eq!(p.pick(Source::Peer(0)), Pick::Chunk(0));

        assert_eq!(p.pick(Source::Peer(1)), Pick::Chunk(0), "the endgame");

        assert_eq!(p.pick(Source::Peer(1)), Pick::Wait, "at most two copies");

        p.failed(0);

        p.failed(0);

        p.drop_peer(0);

        assert_eq!(p.pick(Source::Peer(1)), Pick::Chunk(0));

        p.got(0);

        assert_eq!(p.pick(Source::Peer(1)), Pick::Done);

    }

    #[test]

    fn chunks_already_here_are_never_asked_for() {

        let mut have = Bits::new(3);

        have.set(1);

        let mut p = Picker::new(3, have, 5);

        assert_eq!(p.missing(), 2);

        let mut seen = Vec::new();

        while let Pick::Chunk(i) = p.pick(Source::Hub) {

            seen.push(i);

            p.got(i);

        }

        assert!(!seen.contains(&1));

    }

    #[test]

    fn runners_starting_together_ask_the_hub_for_different_chunks() {

        let a: Vec<_> = {

            let mut p = Picker::new(64, Bits::new(64), 11);

            (0..8).filter_map(|_| if let Pick::Chunk(i) = p.pick(Source::Hub) { Some(i) } else { None }).collect()

        };

        let b: Vec<_> = {

            let mut p = Picker::new(64, Bits::new(64), 12);

            (0..8).filter_map(|_| if let Pick::Chunk(i) = p.pick(Source::Hub) { Some(i) } else { None }).collect()

        };

        assert_ne!(a, b);

    }

}