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

use std::{
    collections::{HashMap, HashSet},
    fmt,
    net::SocketAddr,
    path::Path,
    sync::Arc,
};

use futures::stream::FuturesUnordered;
use lanspread_db::content_manifest::ContentId;
use lanspread_proto::PeerEndpoint;
use crate::PeerEventSender;
use tokio_util::sync::CancellationToken;

use super::{
    DownloadTransferError,
    DownloadTransferErrorKind,
    confined_fs::ConfinedGameRoot,
    manifest::{ValidatedDownloadEntry, ValidatedDownloadManifest},
    ownership::{DownloadOwnershipTransaction, OwnershipJournalPublication},
    planning::{
        ChunkDownloadResult,
        DownloadChunk,
        PeerDownloadPlan,
        build_peer_plans,
        reconcile_chunk_results,
    },
    progress::{DownloadProgressTracker, sample_download_progress},
    retry::{RetryChunk, RetryContext, quarantine_if_integrity_failure, retry_failed_chunks},
    storage::{prepare_game_storage, sync_game_storage},
    task_drain::collect_or_drain_on_cancel,
    transport::{PeerDownloadRequest, download_from_peer},
    version_ini::{
        VersionIniBuffer,
        VersionIniCommit,
        begin_version_ini_transaction,
        commit_version_ini_buffer,
        restore_unjournaled_version_ini_transaction,
    },
};
use crate::{
    DownloadFailureReason,
    DownloadVerificationActivity,
    PeerEvent,
    content_quarantine::ContentQuarantine,
    peer_db::PeerId,
    quic_runtime::QuicConnector,
    transfer_status::{DownloadAttemptReporter, DownloadAttemptStatus},
};

/// Terminal state of a complete payload transfer.
#[derive(Debug)]
pub(crate) enum DownloadCompletion {
    /// Payload, sentinel, and ownership state are durably settled.
    Durable,
    /// The committed payload is visible, but recovery must settle its metadata
    /// before it can be advertised, served, or installed.
    RecoveryRequired(eyre::Report),
}

/// Typed owner-facing failure from one complete ordinary download operation.
#[derive(Debug)]
pub(crate) struct DownloadOperationError {
    reason: Option<DownloadFailureReason>,
    error: eyre::Report,
}

impl DownloadOperationError {
    pub(super) fn cancelled(error: impl Into<eyre::Report>) -> Self {
        Self {
            reason: None,
            error: error.into(),
        }
    }

    pub(super) fn sources_exhausted(error: impl Into<eyre::Report>) -> Self {
        Self {
            reason: Some(DownloadFailureReason::VerifiedCatalogSourcesExhausted),
            error: error.into(),
        }
    }

    pub(super) fn operation_failed(error: impl Into<eyre::Report>) -> Self {
        Self {
            reason: Some(DownloadFailureReason::OperationFailed),
            error: error.into(),
        }
    }

    pub(crate) const fn reason(&self) -> Option<DownloadFailureReason> {
        self.reason
    }

    pub(crate) fn into_report(self) -> eyre::Report {
        self.error
    }
}

/// Aggregates independent chunk failures without letting a later retryable
/// source failure downgrade an already-observed local operation failure.
#[derive(Default)]
struct TransferFailureAggregate {
    operation_failed: Option<DownloadTransferError>,
    cancelled: Option<DownloadTransferError>,
    sources_exhausted: Option<DownloadTransferError>,
}

impl TransferFailureAggregate {
    fn record(&mut self, error: DownloadTransferError) {
        match error.kind() {
            DownloadTransferErrorKind::LocalIo => {
                self.operation_failed.get_or_insert(error);
            }
            DownloadTransferErrorKind::Cancelled => {
                self.cancelled.get_or_insert(error);
            }
            DownloadTransferErrorKind::Integrity | DownloadTransferErrorKind::Transport => {
                self.sources_exhausted.get_or_insert(error);
            }
        }
    }

    fn into_operation_error(self) -> Option<DownloadOperationError> {
        if let Some(error) = self.operation_failed {
            return Some(DownloadOperationError::operation_failed(error));
        }
        if let Some(error) = self.cancelled {
            return Some(DownloadOperationError::cancelled(error));
        }
        self.sources_exhausted
            .map(DownloadOperationError::sources_exhausted)
    }

    fn cancellation_error(&mut self, game_id: &str) -> DownloadOperationError {
        self.operation_failed.take().map_or_else(
            || cancelled_download(game_id),
            DownloadOperationError::operation_failed,
        )
    }
}

impl fmt::Display for DownloadOperationError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.error.fmt(formatter)
    }
}

/// Complete authority and runtime input for one ordinary catalog download.
pub(crate) struct DownloadGameRequest<'a> {
    pub(crate) attempt: &'a DownloadAttemptStatus,
    pub(crate) manifest: ValidatedDownloadManifest,
    pub(crate) state_dir: &'a Path,
    pub(crate) sources: &'a [PeerEndpoint],
    pub(crate) content_id: ContentId,
    pub(crate) quarantine: &'a ContentQuarantine,
    pub(crate) tx_notify_ui: PeerEventSender,
    pub(crate) cancel_token: CancellationToken,
    pub(crate) quic: QuicConnector,
}

/// Downloads all game files from available peers.
#[allow(clippy::too_many_lines)]
pub(crate) async fn download_game_files(
    request: DownloadGameRequest<'_>,
) -> Result<DownloadCompletion, DownloadOperationError> {
    let DownloadGameRequest {
        attempt,
        manifest,
        state_dir,
        sources,
        content_id,
        quarantine,
        tx_notify_ui,
        cancel_token,
        quic,
    } = request;
    let game_id = manifest.game_id().to_owned();
    let manifest_content_id = manifest.catalog_manifest().content_id();
    if manifest_content_id != content_id {
        return Err(DownloadOperationError::operation_failed(eyre::eyre!(
            "download content ID does not match catalog authority for game {game_id}: requested {content_id}, catalog {manifest_content_id}"
        )));
    }
    let sources = eligible_content_sources(sources, content_id, quarantine)
        .map_err(DownloadOperationError::operation_failed)?;
    if sources.is_empty() {
        return Err(DownloadOperationError::sources_exhausted(eyre::eyre!(
            "no peers available for game {game_id}"
        )));
    }

    if cancel_token.is_cancelled() {
        return Err(cancelled_download(&game_id));
    }

    let version_entry = manifest.version_entry();
    let version_buffer = match VersionIniBuffer::new(
        version_entry.destination().canonical(),
        version_entry.size(),
    ) {
        Ok(buffer) => Arc::new(buffer),
        Err(err) => return Err(DownloadOperationError::operation_failed(err)),
    };
    let confined_root = ConfinedGameRoot::open_or_create(manifest.games_folder(), &game_id)
        .map_err(DownloadOperationError::operation_failed)?;
    let ownership = DownloadOwnershipTransaction::prepare(state_dir, &manifest, &confined_root)
        .await
        .map_err(DownloadOperationError::operation_failed)?;

    if let Err(error) = begin_version_ini_transaction(&confined_root) {
        if let Err(restore_error) = restore_before_ownership_journal(&confined_root) {
            return Err(DownloadOperationError::operation_failed(error.wrap_err(
                format!("sentinel parking failed and rollback also failed: {restore_error}"),
            )));
        }
        return Err(DownloadOperationError::operation_failed(error));
    }
    if cancel_token.is_cancelled() {
        restore_before_ownership_journal(&confined_root)
            .map_err(DownloadOperationError::operation_failed)?;
        return Err(cancelled_download(&game_id));
    }
    match ownership.journal_pending().await {
        Ok(OwnershipJournalPublication::Durable) => {}
        Ok(OwnershipJournalPublication::NeedsRecovery(error)) => {
            // The pending record is visible, so restoring the old sentinel
            // would make recovery mistake it for a landed new commit. Stop
            // before payload mutation and leave the phase unambiguous.
            return Err(DownloadOperationError::operation_failed(eyre::eyre!(
                "pending download ownership was renamed but its durability could not be established: {error}"
            )));
        }
        Err(error) => {
            if let Err(restore_error) = restore_before_ownership_journal(&confined_root) {
                return Err(DownloadOperationError::operation_failed(error.wrap_err(
                    format!(
                        "ownership journal failed and sentinel restore also failed: {restore_error}"
                    ),
                )));
            }
            return Err(DownloadOperationError::operation_failed(error));
        }
    }
    if let Err(err) = prepare_game_storage(&manifest, &confined_root) {
        let failure = DownloadOperationError::operation_failed(err);
        return Err(abort_download(&ownership, &game_id, failure).await);
    }
    if cancel_token.is_cancelled() {
        let failure = cancelled_download(&game_id);
        return Err(abort_download(&ownership, &game_id, failure).await);
    }

    let progress_tracker = DownloadProgressTracker::new(total_download_bytes(manifest.entries()));
    let attempt_reporter = attempt.reporter();
    let transfer_ctx = TransferContext {
        game_id: &game_id,
        game_root: &confined_root,
        sources: &sources,
        content_id,
        quarantine,
        tx_notify_ui: &tx_notify_ui,
        cancel_token: &cancel_token,
        quic: &quic,
        version_buffer: version_buffer.clone(),
        progress_tracker: progress_tracker.clone(),
        attempt: attempt_reporter.clone(),
    };
    let plans = match build_initial_transfer_plans(&transfer_ctx, &manifest) {
        Ok(plans) => plans,
        Err(error) => {
            attempt.close_source_admission();
            return Err(abort_download(&ownership, &game_id, error).await);
        }
    };
    attempt.emit_begin();
    attempt.set_activity(DownloadVerificationActivity::VerifyingDownloadedChunks);
    let transfer_result = sample_download_progress(
        attempt_reporter,
        progress_tracker,
        download_transfer_chunks(&transfer_ctx, plans),
    )
    .await;
    attempt.close_source_admission();
    attempt.clear_activity();

    if let Err(err) = transfer_result {
        return Err(abort_download(&ownership, &game_id, err).await);
    }

    if cancel_token.is_cancelled() {
        let failure = cancelled_download(&game_id);
        return Err(abort_download(&ownership, &game_id, failure).await);
    }

    if let Err(error) = sync_game_storage(&manifest, &confined_root) {
        let failure = DownloadOperationError::operation_failed(
            error.wrap_err("failed to make downloaded payload durable"),
        );
        return Err(abort_download(&ownership, &game_id, failure).await);
    }
    if cancel_token.is_cancelled() {
        let failure = cancelled_download(&game_id);
        return Err(abort_download(&ownership, &game_id, failure).await);
    }

    if let Err(error) = ownership.remove_stale() {
        let failure = DownloadOperationError::operation_failed(
            error.wrap_err("failed to remove stale download-owned files"),
        );
        return Err(abort_download(&ownership, &game_id, failure).await);
    }
    if cancel_token.is_cancelled() {
        let failure = cancelled_download(&game_id);
        return Err(abort_download(&ownership, &game_id, failure).await);
    }

    match commit_version_ini_buffer(&confined_root, &version_buffer).await {
        Ok(VersionIniCommit::Durable) => {}
        Ok(VersionIniCommit::NeedsRecovery(error)) => {
            // The visible sentinel makes rollback unsafe. Keep pending ownership
            // so startup recovery can decide from the durable filesystem state.
            return Ok(DownloadCompletion::RecoveryRequired(eyre::eyre!(
                "version.ini was renamed but its durability could not be established: {error}"
            )));
        }
        Err(error) => {
            let failure = DownloadOperationError::operation_failed(error);
            return Err(abort_download(&ownership, &game_id, failure).await);
        }
    }
    match ownership.finalize().await {
        Ok(OwnershipJournalPublication::Durable) => {}
        Ok(OwnershipJournalPublication::NeedsRecovery(error)) => {
            return Ok(DownloadCompletion::RecoveryRequired(error.wrap_err(
                "downloaded payload is visible, but ownership durability must be recovered",
            )));
        }
        Err(error) => {
            return Ok(DownloadCompletion::RecoveryRequired(error.wrap_err(
                "downloaded payload is visible, but ownership finalization must be recovered",
            )));
        }
    }
    log::info!("all files downloaded for game: {game_id}");
    Ok(DownloadCompletion::Durable)
}

fn restore_before_ownership_journal(game_root: &ConfinedGameRoot) -> eyre::Result<()> {
    restore_unjournaled_version_ini_transaction(game_root)
}

fn cancelled_download(game_id: &str) -> DownloadOperationError {
    DownloadOperationError::cancelled(eyre::eyre!("download cancelled for game {game_id}"))
}

async fn abort_download(
    ownership: &DownloadOwnershipTransaction,
    game_id: &str,
    failure: DownloadOperationError,
) -> DownloadOperationError {
    match ownership.abort().await {
        Ok(()) => failure,
        Err(abort_error) => DownloadOperationError::operation_failed(eyre::eyre!(
            "download failed for {game_id}: {failure}; ownership rollback also failed: {abort_error}"
        )),
    }
}

struct TransferContext<'a> {
    game_id: &'a str,
    game_root: &'a ConfinedGameRoot,
    sources: &'a [PeerEndpoint],
    content_id: ContentId,
    quarantine: &'a ContentQuarantine,
    tx_notify_ui: &'a PeerEventSender,
    cancel_token: &'a CancellationToken,
    quic: &'a QuicConnector,
    version_buffer: Arc<VersionIniBuffer>,
    progress_tracker: Arc<DownloadProgressTracker>,
    attempt: DownloadAttemptReporter,
}

struct InitialAttempt {
    source: PeerEndpoint,
    planned_chunks: Vec<DownloadChunk>,
    result: Result<Vec<ChunkDownloadResult>, DownloadTransferError>,
}

fn eligible_content_sources(
    sources: &[PeerEndpoint],
    content_id: ContentId,
    quarantine: &ContentQuarantine,
) -> eyre::Result<Vec<PeerEndpoint>> {
    let mut seen_peer_ids = HashSet::<PeerId>::new();
    let mut peer_by_addr = HashMap::<SocketAddr, PeerId>::new();
    let mut eligible = Vec::with_capacity(sources.len());

    for source in sources {
        if !seen_peer_ids.insert(source.peer_id) {
            continue;
        }
        if let Some(previous_peer_id) = peer_by_addr.insert(source.addr, source.peer_id) {
            eyre::bail!(
                "content source address {} is ambiguously assigned to peers {previous_peer_id} and {}",
                source.addr,
                source.peer_id
            );
        }
        if !quarantine.is_quarantined(source, content_id) {
            eligible.push(*source);
        }
    }

    Ok(eligible)
}

async fn download_transfer_chunks(
    ctx: &TransferContext<'_>,
    plans: HashMap<PeerEndpoint, PeerDownloadPlan>,
) -> Result<(), DownloadOperationError> {
    let tasks = FuturesUnordered::new();
    for (endpoint, plan) in plans {
        let source = endpoint;
        let planned_chunks = plan.chunks.clone();
        let game_root = ctx.game_root.clone();
        let game_id = ctx.game_id.to_string();
        let cancel_token = ctx.cancel_token.clone();
        let version_buffer = ctx.version_buffer.clone();
        let progress_tracker = ctx.progress_tracker.clone();
        let quic = ctx.quic.clone();
        tasks.push(async move {
            let result = download_from_peer(PeerDownloadRequest {
                quic,
                endpoint,
                game_id,
                plan,
                game_root,
                cancel_token,
                version_buffer,
                progress_tracker,
            })
            .await;
            InitialAttempt {
                source,
                planned_chunks,
                result,
            }
        });
    }

    let mut failed_chunks = Vec::new();
    let mut failures = TransferFailureAggregate::default();

    let attempts = collect_or_drain_on_cancel(tasks, ctx.cancel_token, ctx.game_id)
        .await
        .map_err(DownloadOperationError::cancelled)?;
    for attempt in attempts {
        if ctx.cancel_token.is_cancelled() {
            return Err(failures.cancellation_error(ctx.game_id));
        }
        collect_initial_attempt(ctx, attempt, &mut failed_chunks, &mut failures)
            .map_err(DownloadOperationError::operation_failed)?;
    }

    if !failed_chunks.is_empty() {
        retry_chunks(ctx, failed_chunks, &mut failures).await?;
    }

    if ctx.cancel_token.is_cancelled() {
        return Err(failures.cancellation_error(ctx.game_id));
    }

    if let Some(error) = failures.into_operation_error() {
        return Err(error);
    }

    Ok(())
}

fn build_initial_transfer_plans(
    ctx: &TransferContext<'_>,
    manifest: &ValidatedDownloadManifest,
) -> Result<HashMap<PeerEndpoint, PeerDownloadPlan>, DownloadOperationError> {
    let sources = eligible_content_sources(ctx.sources, ctx.content_id, ctx.quarantine)
        .map_err(DownloadOperationError::operation_failed)?;
    if sources.is_empty() {
        return Err(DownloadOperationError::sources_exhausted(eyre::eyre!(
            "no nonquarantined sources remain for game {}",
            ctx.game_id
        )));
    }

    // Local catalog identity defines the exact content and canonical path carried
    // by every request. Planning only distributes those immutable chunks over
    // authenticated, exact-content, quarantine-filtered endpoints.
    let plans =
        build_peer_plans(&sources, manifest).map_err(DownloadOperationError::operation_failed)?;
    if plans.is_empty() {
        return Err(DownloadOperationError::operation_failed(eyre::eyre!(
            "catalog download plan contains no chunks for game {}",
            ctx.game_id
        )));
    }
    Ok(plans)
}

fn collect_initial_attempt(
    ctx: &TransferContext<'_>,
    attempt: InitialAttempt,
    failed_chunks: &mut Vec<RetryChunk>,
    failures: &mut TransferFailureAggregate,
) -> eyre::Result<()> {
    let InitialAttempt {
        source,
        planned_chunks,
        result,
    } = attempt;
    match result {
        Ok(results) => {
            let expected_endpoint = source;
            for (planned_chunk, mut result) in reconcile_chunk_results(
                planned_chunks,
                results,
                expected_endpoint,
                |chunk| chunk,
                "initial download",
            )? {
                result.chunk = planned_chunk;
                collect_initial_chunk_result(ctx, &source, result, failed_chunks, failures);
            }
        }
        Err(error) => {
            for chunk in planned_chunks {
                collect_initial_chunk_result(
                    ctx,
                    &source,
                    ChunkDownloadResult {
                        chunk,
                        result: Err(error.clone()),
                        peer_endpoint: source,
                    },
                    failed_chunks,
                    failures,
                );
            }
        }
    }
    Ok(())
}

fn collect_initial_chunk_result(
    ctx: &TransferContext<'_>,
    source: &PeerEndpoint,
    result: ChunkDownloadResult,
    failed_chunks: &mut Vec<RetryChunk>,
    failures: &mut TransferFailureAggregate,
) {
    match result.result {
        Ok(()) => notify_chunk_finished(ctx, &result.chunk, result.peer_endpoint),
        Err(error) => {
            log::warn!("Failed to download chunk from {}: {error}", source.addr);
            quarantine_if_integrity_failure(ctx.quarantine, source, ctx.content_id, &error);
            match error.kind() {
                DownloadTransferErrorKind::Integrity | DownloadTransferErrorKind::Transport => {
                    failed_chunks.push(RetryChunk::after_failure(result.chunk, *source, error));
                }
                DownloadTransferErrorKind::LocalIo | DownloadTransferErrorKind::Cancelled => {
                    failures.record(error);
                }
            }
        }
    }
}

fn notify_chunk_finished(
    ctx: &TransferContext<'_>,
    chunk: &DownloadChunk,
    peer_endpoint: PeerEndpoint,
) {
    let _ = ctx
        .tx_notify_ui
        .send(PeerEvent::DownloadGameFileChunkFinished {
            id: ctx.game_id.to_string(),
            peer_id: peer_endpoint.peer_id,
            peer_addr: peer_endpoint.addr,
            content_id: chunk.content_id,
            relative_path: chunk.canonical_path().clone(),
            offset: chunk.offset,
            length: chunk.length,
        });
}

async fn retry_chunks(
    ctx: &TransferContext<'_>,
    failed_chunks: Vec<RetryChunk>,
    failures: &mut TransferFailureAggregate,
) -> Result<(), DownloadOperationError> {
    if ctx.cancel_token.is_cancelled() {
        return Err(failures.cancellation_error(ctx.game_id));
    }

    log::info!("Retrying {} failed chunks", failed_chunks.len());

    let retry_ctx = RetryContext {
        sources: ctx.sources,
        content_id: ctx.content_id,
        quarantine: ctx.quarantine,
        game_root: ctx.game_root,
        game_id: ctx.game_id,
        cancel_token: ctx.cancel_token,
        quic: ctx.quic,
        version_buffer: ctx.version_buffer.clone(),
        progress_tracker: ctx.progress_tracker.clone(),
        attempt: ctx.attempt.clone(),
    };
    let retry_results = match retry_failed_chunks(failed_chunks, &retry_ctx).await {
        Ok(results) => results,
        Err(_) if ctx.cancel_token.is_cancelled() => {
            return Err(failures.cancellation_error(ctx.game_id));
        }
        Err(err) => {
            return Err(DownloadOperationError::operation_failed(err));
        }
    };

    for chunk_result in retry_results {
        if ctx.cancel_token.is_cancelled() {
            return Err(failures.cancellation_error(ctx.game_id));
        }

        match chunk_result.result {
            Ok(()) => {
                notify_chunk_finished(ctx, &chunk_result.chunk, chunk_result.peer_endpoint);
            }
            Err(e) => {
                log::error!("Retry failed for chunk: {e}");
                failures.record(e);
            }
        }
    }

    Ok(())
}

fn total_download_bytes(file_descs: &[ValidatedDownloadEntry]) -> u64 {
    file_descs
        .iter()
        .filter(|entry| !entry.is_dir())
        .fold(0u64, |total, entry| total.saturating_add(entry.size()))
}

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

    fn source(peer_id: &str, port: u16) -> PeerEndpoint {
        PeerEndpoint::new(
            PeerId::from_bytes(*blake3::hash(peer_id.as_bytes()).as_bytes()),
            SocketAddr::from(([127, 0, 0, 1], port)),
        )
    }

    fn content(seed: u8) -> ContentId {
        ContentId::from_bytes([seed; 32])
    }

    #[test]
    fn initial_source_filter_skips_bad_source_and_keeps_good_source() {
        let bad = source("bad", 12000);
        let good = source("good", 12001);
        let sources = vec![bad, good];
        let quarantine = ContentQuarantine::default();
        let content_id = content(1);
        quarantine.record_integrity_failure(&bad, content_id);

        let eligible = eligible_content_sources(&sources, content_id, &quarantine)
            .expect("unambiguous content sources should validate");

        assert_eq!(eligible, vec![good]);
    }

    #[test]
    fn initial_source_filter_rejects_ambiguous_address_identity() {
        let sources = vec![source("first", 12000), source("second", 12000)];

        let error = eligible_content_sources(&sources, content(2), &ContentQuarantine::default())
            .expect_err("one address must not represent two authenticated identities");

        assert!(error.to_string().contains("ambiguously assigned"));
    }

    #[test]
    fn local_failure_is_not_downgraded_by_later_retryable_exhaustion() {
        let mut failures = TransferFailureAggregate::default();
        failures.record(DownloadTransferError::local_io("destination write failed"));
        failures.record(DownloadTransferError::transport(
            "retry source disconnected",
        ));

        let error = failures
            .into_operation_error()
            .expect("recorded failures should produce an operation error");

        assert_eq!(error.reason(), Some(DownloadFailureReason::OperationFailed));
        assert_eq!(error.to_string(), "destination write failed");
    }

    #[test]
    fn local_failure_is_not_downgraded_by_later_cancellation() {
        let mut failures = TransferFailureAggregate::default();
        failures.record(DownloadTransferError::local_io("destination write failed"));

        let error = failures.cancellation_error("game");

        assert_eq!(error.reason(), Some(DownloadFailureReason::OperationFailed));
        assert_eq!(error.to_string(), "destination write failed");
    }
}