Contractless/src/records/memory/torrentmap.rs

use crate::torrent::structs::PieceStatus;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

pub enum PieceReservation {
    // The requested piece was assigned to this peer.
    Reserved,
    // The piece is already downloading or complete.
    PieceUnavailable,
    // This peer is not a good candidate for this piece right now.
    PeerUnavailable,
}

#[derive(Debug)]
pub struct Files {
    pub status: PieceStatus,
    pub ip: Option<String>,
    pub failed_ips: Vec<String>,
}

#[derive(Debug)]
pub struct TorrentMap {
    pub pieces: HashMap<u8, Files>,
}

impl TorrentMap {
    pub fn add_piece(
        &mut self,
        piece_number: u8,
        status: PieceStatus,
        ip: Option<String>,
        failed_ips: Vec<String>,
    ) {
        // Every piece starts with its current status plus the peer history used
        // to avoid repeatedly asking the same failing peer.
        let file = Files {
            status,
            ip,
            failed_ips,
        };
        self.pieces.insert(piece_number, file);
    }

    pub fn update_piece(
        &mut self,
        piece_number: u8,
        status: Option<PieceStatus>,
        ip: Option<String>,
        failed_ip: Option<String>,
    ) -> Result<(), String> {
        if let Some(piece) = self.pieces.get_mut(&piece_number) {
            // Only fields supplied by the caller are changed, which lets the
            // downloader update status, owner IP, and failure history separately.
            if let Some(new_status) = status {
                piece.status = new_status;
            }
            if let Some(new_ip) = ip {
                piece.ip = Some(new_ip);
            }
            if let Some(new_failed_ip) = failed_ip {
                piece.failed_ips.push(new_failed_ip);
            }
            Ok(())
        } else {
            Err(format!("Piece number {piece_number} not found"))
        }
    }

    pub fn get_piece_data(&self, piece_number: u8) -> Option<&Files> {
        self.pieces.get(&piece_number)
    }

    pub fn try_reserve_piece_for_ip(
        &mut self,
        piece_number: u8,
        ip: &str,
    ) -> Result<PieceReservation, String> {
        if let Some(piece) = self.pieces.get(&piece_number) {
            // Do not reassign a piece to a peer that already failed it.
            if piece.failed_ips.iter().any(|failed_ip| failed_ip == ip) {
                return Ok(PieceReservation::PeerUnavailable);
            }
            // Completed and actively downloading pieces are not available for
            // another reservation.
            match piece.status {
                PieceStatus::Downloading | PieceStatus::Complete => {
                    return Ok(PieceReservation::PieceUnavailable);
                }
                PieceStatus::Pending | PieceStatus::Failed => {}
            }
        } else {
            return Err(format!("Piece number {piece_number} not found"));
        }

        for (existing_piece_number, piece) in &self.pieces {
            if *existing_piece_number == piece_number {
                continue;
            }
            // A peer should only carry one active piece at a time so one slow
            // connection cannot stall several pieces at once.
            if piece.status == PieceStatus::Downloading && piece.ip.as_deref() == Some(ip) {
                return Ok(PieceReservation::PeerUnavailable);
            }
        }

        let piece = self
            .pieces
            .get_mut(&piece_number)
            .ok_or_else(|| format!("Piece number {piece_number} not found"))?;
        // Reservation marks ownership immediately; completion/failure will
        // clear the IP when the piece finishes.
        piece.status = PieceStatus::Downloading;
        piece.ip = Some(ip.to_string());
        Ok(PieceReservation::Reserved)
    }

    pub fn mark_piece_complete(&mut self, piece_number: u8) -> Result<(), String> {
        let piece = self
            .pieces
            .get_mut(&piece_number)
            .ok_or_else(|| format!("Piece number {piece_number} not found"))?;
        // A completed piece no longer belongs to a peer.
        piece.status = PieceStatus::Complete;
        piece.ip = None;
        Ok(())
    }

    pub fn mark_piece_failed(&mut self, piece_number: u8, failed_ip: &str) -> Result<(), String> {
        let piece = self
            .pieces
            .get_mut(&piece_number)
            .ok_or_else(|| format!("Piece number {piece_number} not found"))?;
        // Failed pieces return to the pool, but remember which peer failed.
        piece.status = PieceStatus::Failed;
        piece.ip = None;
        if !piece.failed_ips.iter().any(|ip| ip == failed_ip) {
            piece.failed_ips.push(failed_ip.to_string());
        }
        Ok(())
    }
}

// Type alias for shared TorrentMap
pub type SharedTorrentMap = Arc<Mutex<TorrentMap>>;