//! Blobs: content-addressed bytes fetched in chunks from several peers at once.
//!
//! Large transfers can't go in one Mycelium message (it must arrive within Mycelium's
//! five-minute deadline), and a blob every node holds identically can come from many
//! of them. A blob is an Automerge change by its hash (e.g. an epoch's snapshot change,
//! which every node on that epoch holds byte for byte), or content by its BLAKE3 hash
//! (e.g. the new document of a `replace` or `import` remake).
//!
//! * Serving: a request without chunks gets the manifest (size, chunk size, a BLAKE3
//!   hash per chunk); one with chunks gets each chunk as a message of its own. A node
//!   that doesn't hold the blob answers `Missing`. The bytes of the blob served last
//!   are kept, so serving a download doesn't extract them for every chunk. Serving runs
//!   in tasks of its own, at most [`MAX_SERVING`] chunk sends at a time.
//! * Fetching: the manifest from the first peer that has it (closest first), then the
//!   chunks, a few in flight per peer from up to [`MAX_SOURCES`] peers. Every chunk is
//!   checked against the manifest; a peer that sends a bad one is dropped, a chunk
//!   that doesn't arrive in time is asked of another peer. The caller checks the whole
//!   (for a change: its hash).
//!
//! The chunk size is the serving node's (`--blob-chunk-kib`); the manifest carries it.

use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;

use anyhow::{Result, anyhow, bail};
use automerge::ChangeHash;
use tokio::sync::{Semaphore, mpsc};
use tokio::time::Instant;
use tracing::{debug, info, warn};

use crate::state::AppState;
use crate::transport::{BlobBody, BlobId, BlobMessage, ParsedMessage};

/// Chunk sends in progress at once, over all downloads served.
const MAX_SERVING: usize = 8;
/// Peers a download fetches chunks from at once.
const MAX_SOURCES: usize = 4;
/// Chunks in flight per peer.
const WINDOW: usize = 4;
/// Largest blob a download accepts.
const MAX_SIZE: u64 = 4 << 30;
/// Largest chunk a download accepts.
const MAX_CHUNK: u32 = 4 << 20;
/// How long to wait for a manifest from one peer.
const MANIFEST_TIMEOUT: Duration = Duration::from_secs(15);
/// How long to wait for a chunk before asking another peer (raised by the peer's
/// round-trip time).
const CHUNK_TIMEOUT: Duration = Duration::from_secs(60);
/// A download gives up after this long without any chunk arriving.
const STALL_TIMEOUT: Duration = Duration::from_secs(180);
/// Bad chunks after which a peer is no longer asked.
const MAX_BAD: u32 = 2;

/// Per node: what is being served and fetched.
pub(crate) struct Blobs {
    chunk_size: u32,
    serving: Arc<Semaphore>,
    /// The blob served last.
    cache: Mutex<Option<(BlobId, Arc<Prepared>)>>,
    /// Answers for downloads in progress, by blob.
    fetches: Mutex<HashMap<BlobId, mpsc::UnboundedSender<(String, BlobBody)>>>,
}

struct Prepared {
    bytes: Arc<Vec<u8>>,
    hashes: Vec<[u8; 32]>,
}

impl Blobs {
    pub(crate) fn new(chunk_size: u32) -> Self {
        Self {
            chunk_size: chunk_size.max(1),
            serving: Arc::new(Semaphore::new(MAX_SERVING)),
            cache: Mutex::new(None),
            fetches: Mutex::new(HashMap::new()),
        }
    }

    fn fetches(
        &self,
    ) -> MutexGuard<'_, HashMap<BlobId, mpsc::UnboundedSender<(String, BlobBody)>>> {
        self.fetches
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }

    fn cached(&self, id: BlobId) -> Option<Arc<Prepared>> {
        self.cache
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .as_ref()
            .filter(|(cached, _)| *cached == id)
            .map(|(_, prepared)| Arc::clone(prepared))
    }

    fn remember(&self, id: BlobId, prepared: Arc<Prepared>) {
        *self
            .cache
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some((id, prepared));
    }
}

fn chunk_hashes(bytes: &[u8], chunk_size: u32) -> Vec<[u8; 32]> {
    bytes
        .chunks(chunk_size as usize)
        .map(|chunk| *blake3::hash(chunk).as_bytes())
        .collect()
}

fn chunk_count(size: u64, chunk_size: u32) -> u64 {
    size.div_ceil(u64::from(chunk_size))
}

/// A blob message from `peer`.
pub(crate) async fn handle(state: &AppState, peer: &str, message: BlobMessage) {
    match message.body {
        BlobBody::Request { chunks } => {
            let state = state.clone();
            let peer = peer.to_owned();
            // Off the inbound loop: answering may send many large messages.
            tokio::spawn(async move { serve(&state, &peer, message.id, chunks).await });
        }
        body => {
            let sender = state.blobs.fetches().get(&message.id).cloned();
            match sender {
                Some(sender) => {
                    let _ = sender.send((peer.to_owned(), body));
                }
                None => {
                    debug!(peer = %peer, blob = %message.id, "Blob answer for no download in progress")
                }
            }
        }
    }
}

async fn prepare(state: &AppState, id: BlobId) -> Option<Arc<Prepared>> {
    if let Some(prepared) = state.blobs.cached(id) {
        return Some(prepared);
    }
    let bytes = match id {
        BlobId::Change(hash) => {
            let doc = Arc::clone(&state.doc);
            tokio::task::spawn_blocking(move || {
                // get_change_by_hash needs the document mutably.
                let mut replica = doc.blocking_write();
                replica
                    .doc
                    .get_change_by_hash(&hash)
                    .map(|change| change.raw_bytes().to_vec())
            })
            .await
            .ok()
            .flatten()
            .map(Arc::new)?
        }
        BlobId::Content(hash) => state.remake.content(&hash)?,
    };
    let chunk_size = state.blobs.chunk_size;
    let prepared = tokio::task::spawn_blocking(move || {
        let hashes = chunk_hashes(&bytes, chunk_size);
        Prepared { bytes, hashes }
    })
    .await
    .ok()?;
    let prepared = Arc::new(prepared);
    state.blobs.remember(id, Arc::clone(&prepared));
    Some(prepared)
}

async fn reply(state: &AppState, peer: &str, id: BlobId, body: BlobBody) {
    let message = ParsedMessage::Blob(BlobMessage {
        epoch: state.local_epoch(),
        id,
        body,
    });
    if let Err(error) = state.transport.send_message(peer, &message).await {
        debug!(peer = %peer, blob = %id, error = %format!("{error:#}"), "Blob answer failed");
    }
}

async fn serve(state: &AppState, peer: &str, id: BlobId, chunks: Vec<u32>) {
    let Some(prepared) = prepare(state, id).await else {
        return reply(state, peer, id, BlobBody::Missing).await;
    };
    let chunk_size = state.blobs.chunk_size;
    if chunks.is_empty() {
        let body = BlobBody::Manifest {
            size: prepared.bytes.len() as u64,
            chunk_size,
            hashes: prepared.hashes.clone(),
        };
        return reply(state, peer, id, body).await;
    }
    for index in chunks.into_iter().take(WINDOW * 2) {
        let start = index as usize * chunk_size as usize;
        if start >= prepared.bytes.len() {
            continue;
        }
        let end = (start + chunk_size as usize).min(prepared.bytes.len());
        let Ok(_permit) = state.blobs.serving.acquire().await else {
            return;
        };
        let data = bytes::Bytes::copy_from_slice(&prepared.bytes[start..end]);
        reply(state, peer, id, BlobBody::Chunk { index, data }).await;
    }
}

// --- fetching ----------------------------------------------------------------------

/// Removes the download's answer channel when it ends, however it ends.
struct Registration<'a> {
    state: &'a AppState,
    id: BlobId,
}

impl Drop for Registration<'_> {
    fn drop(&mut self) {
        self.state.blobs.fetches().remove(&self.id);
    }
}

async fn request(state: &AppState, peer: &str, id: BlobId, chunks: Vec<u32>) {
    let message = ParsedMessage::Blob(BlobMessage {
        epoch: state.local_epoch(),
        id,
        body: BlobBody::Request { chunks },
    });
    if let Err(error) = state.transport.send_message(peer, &message).await {
        debug!(peer = %peer, blob = %id, error = %format!("{error:#}"), "Blob request failed");
    }
}

struct Manifest {
    size: u64,
    chunk_size: u32,
    hashes: Vec<[u8; 32]>,
}

/// One peer during a download.
struct Source {
    peer: String,
    in_flight: HashMap<u32, Instant>,
    window: usize,
    bad: u32,
    timeout: Duration,
}

/// Downloads blob `id` from `peers` (all thought to hold it).
pub(crate) async fn fetch(state: &AppState, id: BlobId, peers: Vec<String>) -> Result<Vec<u8>> {
    let (sender, mut answers) = mpsc::unbounded_channel();
    if state.blobs.fetches().insert(id, sender).is_some() {
        bail!("{id} is being downloaded already");
    }
    let _registration = Registration { state, id };

    // Closest first; peers without a round-trip time last.
    let now = Instant::now();
    let mut peers: Vec<(String, Option<Duration>)> = {
        let rtt = state.lock_rtt();
        peers
            .into_iter()
            .map(|peer| {
                let estimate = rtt.effective(&peer, now, state.timing.rtt_fresh_for());
                (peer, estimate)
            })
            .collect()
    };
    peers.sort_by_key(|(_, rtt)| rtt.unwrap_or(Duration::MAX));
    let rtts: HashMap<String, Option<Duration>> = peers.iter().cloned().collect();
    let mut candidates: VecDeque<String> = peers.into_iter().map(|(peer, _)| peer).collect();
    if candidates.is_empty() {
        bail!("no peer to download {id} from");
    }

    // The manifest, from the first peer that has the blob.
    let mut manifest = None;
    let mut missing: HashSet<String> = HashSet::new();
    for peer in &candidates {
        request(state, peer, id, Vec::new()).await;
        let deadline = Instant::now() + MANIFEST_TIMEOUT;
        loop {
            let answer = tokio::time::timeout_at(deadline, answers.recv()).await;
            match answer {
                Ok(Some((
                    from,
                    BlobBody::Manifest {
                        size,
                        chunk_size,
                        hashes,
                    },
                ))) => {
                    if size > MAX_SIZE
                        || chunk_size == 0
                        || chunk_size > MAX_CHUNK
                        || hashes.len() as u64 != chunk_count(size, chunk_size)
                    {
                        warn!(peer = %from, blob = %id, size, chunk_size, "Unusable blob manifest; asking another peer");
                        missing.insert(from);
                        break;
                    }
                    manifest = Some(Manifest {
                        size,
                        chunk_size,
                        hashes,
                    });
                    break;
                }
                Ok(Some((from, BlobBody::Missing))) => {
                    missing.insert(from.clone());
                    if &from == peer {
                        break;
                    }
                }
                Ok(Some(_)) => {} // chunks of an earlier attempt
                Ok(None) | Err(_) => break,
            }
        }
        if manifest.is_some() {
            break;
        }
    }
    let Some(manifest) = manifest else {
        bail!("no peer sent a manifest for {id}");
    };
    candidates.retain(|peer| !missing.contains(peer));

    let total = manifest.hashes.len();
    info!(blob = %id, bytes = manifest.size, chunks = total, peers = candidates.len(), "Downloading blob");
    let started = Instant::now();
    let mut chunks: Vec<Option<bytes::Bytes>> = vec![None; total];
    let mut done = 0usize;
    let mut sources: Vec<Source> = Vec::new();
    let mut last_progress = Instant::now();

    while done < total {
        // Fill up the sources, then their windows.
        while sources.len() < MAX_SOURCES {
            let Some(peer) = candidates.pop_front() else {
                break;
            };
            let timeout =
                CHUNK_TIMEOUT + rtts.get(&peer).copied().flatten().unwrap_or_default() * 4;
            sources.push(Source {
                peer,
                in_flight: HashMap::new(),
                window: WINDOW,
                bad: 0,
                timeout,
            });
        }
        if sources.is_empty() {
            bail!("no peer left to download {id} from ({done} of {total} chunks)");
        }
        let now = Instant::now();
        let assigned: HashSet<u32> = sources
            .iter()
            .flat_map(|source| source.in_flight.keys().copied())
            .collect();
        let mut wanted = (0..total as u32)
            .filter(|index| chunks[*index as usize].is_none() && !assigned.contains(index));
        for source in &mut sources {
            let free = source.window.saturating_sub(source.in_flight.len());
            let batch: Vec<u32> = wanted.by_ref().take(free).collect();
            if batch.is_empty() {
                continue;
            }
            for index in &batch {
                source.in_flight.insert(*index, now);
            }
            request(state, &source.peer, id, batch).await;
        }

        match tokio::time::timeout(Duration::from_millis(500), answers.recv()).await {
            Ok(Some((from, BlobBody::Chunk { index, data }))) => {
                let Some(source) = sources.iter_mut().find(|source| source.peer == from) else {
                    continue;
                };
                let sent_at = source.in_flight.remove(&index);
                let slot = index as usize;
                if slot >= total || chunks[slot].is_some() {
                    continue;
                }
                let expected = expected_len(&manifest, slot);
                if data.len() != expected
                    || *blake3::hash(&data).as_bytes() != manifest.hashes[slot]
                {
                    source.bad += 1;
                    warn!(peer = %from, blob = %id, index, "Bad blob chunk");
                    continue;
                }
                if let Some(sent_at) = sent_at {
                    debug!(peer = %from, blob = %id, index, ms = sent_at.elapsed().as_millis() as u64, "Blob chunk");
                }
                chunks[slot] = Some(data);
                done += 1;
                last_progress = Instant::now();
                // A peer that delivers gets its full window back.
                source.window = WINDOW;
            }
            Ok(Some((from, BlobBody::Missing))) => {
                sources.retain(|source| source.peer != from);
            }
            Ok(Some(_)) => {}
            Ok(None) => bail!("download of {id} cancelled"),
            Err(_) => {}
        }

        // Chunks that took too long go to another peer; the slow peer gets a smaller
        // window. Peers that sent bad chunks are dropped.
        let now = Instant::now();
        for source in &mut sources {
            let timeout = source.timeout;
            let before = source.in_flight.len();
            source
                .in_flight
                .retain(|_, sent| now.saturating_duration_since(*sent) < timeout);
            if source.in_flight.len() < before {
                source.window = 1;
                debug!(peer = %source.peer, blob = %id, "Blob chunks timed out; asking others");
            }
        }
        sources.retain(|source| {
            let keep = source.bad < MAX_BAD;
            if !keep {
                warn!(peer = %source.peer, blob = %id, "Dropping a peer that sent bad blob chunks");
            }
            keep
        });
        if last_progress.elapsed() > STALL_TIMEOUT {
            bail!(
                "download of {id} stalled ({done} of {total} chunks, nothing for {}s)",
                STALL_TIMEOUT.as_secs()
            );
        }
    }

    let mut bytes = Vec::with_capacity(manifest.size as usize);
    for chunk in chunks.into_iter().flatten() {
        bytes.extend_from_slice(&chunk);
    }
    info!(
        blob = %id,
        bytes = bytes.len(),
        secs = started.elapsed().as_secs_f32(),
        "Downloaded blob"
    );
    Ok(bytes)
}

fn expected_len(manifest: &Manifest, slot: usize) -> usize {
    let chunk_size = manifest.chunk_size as u64;
    let start = slot as u64 * chunk_size;
    (manifest.size - start).min(chunk_size) as usize
}

/// Downloads change `hash` from `peers` and checks it.
/// Downloads the content with BLAKE3 hash `hash` from `peers` and checks it.
pub(crate) async fn fetch_content(
    state: &AppState,
    hash: [u8; 32],
    peers: Vec<String>,
) -> Result<Vec<u8>> {
    let id = BlobId::Content(hash);
    let bytes = fetch(state, id, peers).await?;
    if *blake3::hash(&bytes).as_bytes() != hash {
        bail!("downloaded {id} doesn't match its hash");
    }
    Ok(bytes)
}

pub(crate) async fn fetch_change(
    state: &AppState,
    hash: ChangeHash,
    peers: Vec<String>,
) -> Result<automerge::Change> {
    let bytes = fetch(state, BlobId::Change(hash), peers).await?;
    let change = automerge::Change::from_bytes(bytes)
        .map_err(|error| anyhow!("downloaded change {hash} doesn't load: {error}"))?;
    if change.hash() != hash {
        bail!("downloaded change has hash {}, not {hash}", change.hash());
    }
    Ok(change)
}

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

    #[test]
    fn chunking() {
        let bytes: Vec<u8> = (0..1000u32).map(|n| n as u8).collect();
        let hashes = chunk_hashes(&bytes, 300);
        assert_eq!(hashes.len(), 4);
        assert_eq!(chunk_count(1000, 300), 4);
        assert_eq!(chunk_count(900, 300), 3);
        let manifest = Manifest {
            size: 1000,
            chunk_size: 300,
            hashes,
        };
        assert_eq!(expected_len(&manifest, 0), 300);
        assert_eq!(expected_len(&manifest, 3), 100);
        assert_eq!(manifest.hashes[3], *blake3::hash(&bytes[900..]).as_bytes());
    }
}