lanspread/crates/lanspread-peer/src/download/planning.rs

use std::{collections::HashMap, net::SocketAddr};

use lanspread_db::db::GameFileDescription;

use crate::config::CHUNK_SIZE;

/// Represents a chunk of a file to be downloaded.
#[derive(Debug, Clone)]
pub(super) struct DownloadChunk {
    pub(super) relative_path: String,
    pub(super) offset: u64,
    pub(super) length: u64,
    pub(super) retry_count: usize,
    pub(super) last_peer: Option<SocketAddr>,
}

/// Download plan for a single peer.
#[derive(Debug, Default)]
pub(super) struct PeerDownloadPlan {
    pub(super) chunks: Vec<DownloadChunk>,
}

/// Result of downloading a chunk.
#[derive(Debug)]
pub(super) struct ChunkDownloadResult {
    pub(super) chunk: DownloadChunk,
    pub(super) result: eyre::Result<()>,
    pub(super) peer_addr: SocketAddr,
}

/// Extracts the root `version.ini` descriptor while keeping every descriptor in
/// the transfer list. The chunk writer diverts the sentinel bytes into memory.
pub(super) fn extract_version_descriptor(
    game_id: &str,
    game_file_descs: Vec<GameFileDescription>,
) -> eyre::Result<(GameFileDescription, Vec<GameFileDescription>)> {
    let mut version_descs = Vec::new();
    let mut transfer_descs = Vec::new();

    for desc in game_file_descs {
        if desc.is_version_ini() {
            version_descs.push(desc.clone());
        }
        transfer_descs.push(desc);
    }

    if version_descs.len() != 1 {
        eyre::bail!(
            "expected exactly one root-level version.ini sentinel for {game_id}, found {}",
            version_descs.len()
        );
    }

    let version_desc = version_descs.remove(0);
    Ok((version_desc, transfer_descs))
}

/// Resolves which peers have a specific file.
pub(super) fn resolve_file_peers<'a>(
    relative_path: &str,
    file_peer_map: &'a HashMap<String, Vec<SocketAddr>>,
    fallback: &'a [SocketAddr],
) -> &'a [SocketAddr] {
    if let Some(peers) = file_peer_map.get(relative_path)
        && !peers.is_empty()
    {
        return peers;
    }

    fallback
}

/// Builds download plans distributing files across peers.
pub(super) fn build_peer_plans(
    peers: &[SocketAddr],
    file_descs: &[GameFileDescription],
    file_peer_map: &HashMap<String, Vec<SocketAddr>>,
) -> HashMap<SocketAddr, PeerDownloadPlan> {
    let mut plans: HashMap<SocketAddr, PeerDownloadPlan> = HashMap::new();
    if peers.is_empty() {
        return plans;
    }

    let mut planned_bytes: HashMap<SocketAddr, u64> = HashMap::new();
    let mut tie_breaker = 0usize;

    for desc in file_descs.iter().filter(|d| !d.is_dir) {
        let size = desc.file_size();
        let eligible_peers = resolve_file_peers(&desc.relative_path, file_peer_map, peers);
        if eligible_peers.is_empty() {
            continue;
        }

        if size == 0 {
            let peer = select_least_loaded_peer(eligible_peers, &planned_bytes, &mut tie_breaker);
            *planned_bytes.entry(peer).or_default() += 1;
            plans.entry(peer).or_default().chunks.push(DownloadChunk {
                relative_path: desc.relative_path.clone(),
                offset: 0,
                length: 0,
                retry_count: 0,
                last_peer: Some(peer),
            });
            continue;
        }

        let mut offset = 0u64;
        while offset < size {
            let length = std::cmp::min(CHUNK_SIZE, size - offset);
            let peer = select_least_loaded_peer(eligible_peers, &planned_bytes, &mut tie_breaker);
            *planned_bytes.entry(peer).or_default() += length;
            plans.entry(peer).or_default().chunks.push(DownloadChunk {
                relative_path: desc.relative_path.clone(),
                offset,
                length,
                retry_count: 0,
                last_peer: Some(peer),
            });
            offset += length;
        }
    }

    plans
}

fn select_least_loaded_peer(
    eligible_peers: &[SocketAddr],
    planned_bytes: &HashMap<SocketAddr, u64>,
    tie_breaker: &mut usize,
) -> SocketAddr {
    let start = *tie_breaker % eligible_peers.len();
    *tie_breaker = (*tie_breaker).wrapping_add(1);

    let mut selected = eligible_peers[start];
    let mut selected_load = planned_bytes.get(&selected).copied().unwrap_or_default();

    for offset in 1..eligible_peers.len() {
        let peer = eligible_peers[(start + offset) % eligible_peers.len()];
        let load = planned_bytes.get(&peer).copied().unwrap_or_default();
        if load < selected_load {
            selected = peer;
            selected_load = load;
        }
    }

    selected
}

#[cfg(test)]
mod tests {
    use lanspread_db::db::GameFileDescription;

    use super::*;
    fn loopback_addr(port: u16) -> SocketAddr {
        SocketAddr::from(([127, 0, 0, 1], port))
    }

    #[test]
    fn build_peer_plans_handles_partial_final_chunk() {
        let peers = vec![loopback_addr(12000), loopback_addr(12001)];
        let file_size = CHUNK_SIZE * 2 + CHUNK_SIZE / 4;
        let mut file_peer_map = HashMap::new();
        file_peer_map.insert("game/file.dat".to_string(), peers.clone());
        let file_descs = vec![GameFileDescription {
            game_id: "test".to_string(),
            relative_path: "game/file.dat".to_string(),
            is_dir: false,
            size: file_size,
        }];

        let plans = build_peer_plans(&peers, &file_descs, &file_peer_map);
        let mut chunks: Vec<_> = plans.values().flat_map(|plan| plan.chunks.iter()).collect();

        assert_eq!(chunks.len(), 3, "expected three chunks for 2.25 blocks");

        chunks.sort_by_key(|chunk| chunk.offset);
        let last_chunk = chunks.last().expect("last chunk exists");

        assert_eq!(last_chunk.offset, CHUNK_SIZE * 2);
        assert_eq!(last_chunk.length, file_size - last_chunk.offset);
        assert_eq!(last_chunk.length, CHUNK_SIZE / 4);
        assert_eq!(
            last_chunk.offset + last_chunk.length,
            file_size,
            "last chunk should finish the file"
        );
    }

    #[test]
    fn build_peer_plans_spreads_large_file_chunks_across_shared_peers() {
        let peers = vec![loopback_addr(12000), loopback_addr(12001)];
        let large_file = "game/large.eti";
        let file_size = CHUNK_SIZE * 3 + CHUNK_SIZE / 2;
        let mut file_peer_map = HashMap::new();
        file_peer_map.insert("game/version.ini".to_string(), peers.clone());
        file_peer_map.insert(large_file.to_string(), peers.clone());
        let file_descs = vec![
            GameFileDescription {
                game_id: "game".to_string(),
                relative_path: "game/version.ini".to_string(),
                is_dir: false,
                size: 9,
            },
            GameFileDescription {
                game_id: "game".to_string(),
                relative_path: large_file.to_string(),
                is_dir: false,
                size: file_size,
            },
        ];

        let plans = build_peer_plans(&peers, &file_descs, &file_peer_map);
        let mut chunk_counts = HashMap::new();
        let mut byte_counts = HashMap::new();

        for (peer, plan) in plans {
            for chunk in plan.chunks {
                if chunk.relative_path == large_file {
                    *chunk_counts.entry(peer).or_insert(0usize) += 1;
                    *byte_counts.entry(peer).or_insert(0u64) += chunk.length;
                }
            }
        }

        assert_eq!(chunk_counts.get(&peers[0]), Some(&2));
        assert_eq!(chunk_counts.get(&peers[1]), Some(&2));

        let peer_a_bytes = byte_counts.get(&peers[0]).copied().unwrap_or_default();
        let peer_b_bytes = byte_counts.get(&peers[1]).copied().unwrap_or_default();
        assert_eq!(peer_a_bytes + peer_b_bytes, file_size);
        assert!(
            peer_a_bytes.abs_diff(peer_b_bytes) <= CHUNK_SIZE,
            "large file bytes should be balanced within one chunk: {peer_a_bytes} vs {peer_b_bytes}"
        );
    }

    #[test]
    fn build_peer_plans_respects_file_peer_map() {
        let shared_a = loopback_addr(12010);
        let shared_b = loopback_addr(12011);
        let exclusive = loopback_addr(12012);
        let peers = vec![shared_a, shared_b, exclusive];

        let mut file_peer_map = HashMap::new();
        file_peer_map.insert("shared.bin".to_string(), vec![shared_a, shared_b]);
        file_peer_map.insert("exclusive.bin".to_string(), vec![exclusive]);

        let file_descs = vec![
            GameFileDescription {
                game_id: "test".to_string(),
                relative_path: "shared.bin".to_string(),
                is_dir: false,
                size: CHUNK_SIZE * 2,
            },
            GameFileDescription {
                game_id: "test".to_string(),
                relative_path: "exclusive.bin".to_string(),
                is_dir: false,
                size: CHUNK_SIZE,
            },
        ];

        let plans = build_peer_plans(&peers, &file_descs, &file_peer_map);
        let exclusive_plan = plans
            .get(&exclusive)
            .expect("exclusive peer should have a plan");
        assert!(
            exclusive_plan
                .chunks
                .iter()
                .all(|chunk| chunk.relative_path == "exclusive.bin"),
            "exclusive peer should only receive exclusive.bin chunks"
        );

        for (peer, plan) in plans {
            for chunk in plan.chunks {
                match chunk.relative_path.as_str() {
                    "exclusive.bin" => assert_eq!(
                        peer, exclusive,
                        "exclusive.bin chunks should only be assigned to the exclusive peer"
                    ),
                    "shared.bin" => assert!(
                        peer == shared_a || peer == shared_b,
                        "shared.bin chunks must stay within shared peers"
                    ),
                    other => panic!("unexpected file in plan: {other}"),
                }
            }
        }
    }

    #[test]
    fn version_descriptor_extraction_keeps_nested_decoy_in_transfer_list() {
        let nested_decoy = vec![
            GameFileDescription {
                game_id: "game".to_string(),
                relative_path: "game/version.ini".to_string(),
                is_dir: false,
                size: 8,
            },
            GameFileDescription {
                game_id: "game".to_string(),
                relative_path: "game/local/version.ini".to_string(),
                is_dir: false,
                size: 8,
            },
        ];

        let (version, transfer) =
            extract_version_descriptor("game", nested_decoy).expect("only one root sentinel");
        assert_eq!(version.relative_path, "game/version.ini");
        assert_eq!(transfer.len(), 2);
    }

    #[test]
    fn version_descriptor_extraction_requires_a_root_version_ini() {
        let missing = vec![GameFileDescription {
            game_id: "game".to_string(),
            relative_path: "game/archive.eti".to_string(),
            is_dir: false,
            size: 1,
        }];
        assert!(extract_version_descriptor("game", missing).is_err());
    }

    #[test]
    fn version_descriptor_extraction_rejects_duplicate_root_version_ini() {
        let multiple = vec![
            GameFileDescription {
                game_id: "game".to_string(),
                relative_path: "game/version.ini".to_string(),
                is_dir: false,
                size: 8,
            },
            GameFileDescription {
                game_id: "game".to_string(),
                relative_path: "game/version.ini".to_string(),
                is_dir: false,
                size: 8,
            },
        ];
        assert!(extract_version_descriptor("game", multiple).is_err());
    }
}