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

use tokio_util::sync::CancellationToken;

use super::{
    planning::{ChunkDownloadResult, DownloadChunk, PeerDownloadPlan, resolve_file_peers},
    transport::download_from_peer,
    version_ini::VersionIniBuffer,
};
use crate::config::MAX_RETRY_COUNT;

/// Selects a peer for retrying a failed chunk.
fn select_retry_peer(
    peers: &[SocketAddr],
    last_peer: Option<SocketAddr>,
    attempt_offset: usize,
) -> Option<SocketAddr> {
    if peers.is_empty() {
        return None;
    }

    if peers.len() > 1
        && let Some(last) = last_peer
        && let Some(pos) = peers.iter().position(|addr| *addr == last)
    {
        let next_index = (pos + 1 + attempt_offset) % peers.len();
        return Some(peers[next_index]);
    }

    Some(peers[attempt_offset % peers.len()])
}

/// Returns a fallback peer address for error reporting.
fn fallback_peer_addr(peers: &[SocketAddr], last_peer: Option<SocketAddr>) -> SocketAddr {
    last_peer
        .or_else(|| peers.first().copied())
        .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)))
}

/// Retries downloading failed chunks.
pub(super) async fn retry_failed_chunks(
    failed_chunks: Vec<DownloadChunk>,
    peers: &[SocketAddr],
    base_dir: &Path,
    game_id: &str,
    file_peer_map: &HashMap<String, Vec<SocketAddr>>,
    cancel_token: &CancellationToken,
    version_buffer: Option<Arc<VersionIniBuffer>>,
) -> eyre::Result<Vec<ChunkDownloadResult>> {
    let mut exhausted = Vec::new();
    let mut queue: VecDeque<DownloadChunk> = failed_chunks.into_iter().collect();

    while let Some(mut chunk) = queue.pop_front() {
        if cancel_token.is_cancelled() {
            return Ok(exhausted);
        }

        let eligible_peers = resolve_file_peers(&chunk.relative_path, file_peer_map, peers);

        if chunk.retry_count >= MAX_RETRY_COUNT {
            exhausted.push(ChunkDownloadResult {
                chunk: chunk.clone(),
                result: Err(eyre::eyre!(
                    "Retry budget exhausted for chunk: {}",
                    chunk.relative_path
                )),
                peer_addr: fallback_peer_addr(eligible_peers, chunk.last_peer),
            });
            continue;
        }

        let retry_offset = chunk.retry_count.saturating_sub(1);
        let Some(peer_addr) = select_retry_peer(eligible_peers, chunk.last_peer, retry_offset)
        else {
            exhausted.push(ChunkDownloadResult {
                chunk: chunk.clone(),
                result: Err(eyre::eyre!(
                    "No peers available to retry chunk: {}",
                    chunk.relative_path
                )),
                peer_addr: fallback_peer_addr(eligible_peers, chunk.last_peer),
            });
            continue;
        };

        let mut attempt_chunk = chunk.clone();
        attempt_chunk.last_peer = Some(peer_addr);

        let plan = PeerDownloadPlan {
            chunks: vec![attempt_chunk.clone()],
            whole_files: Vec::new(),
        };

        match download_from_peer(
            peer_addr,
            game_id,
            plan,
            base_dir.to_path_buf(),
            cancel_token,
            version_buffer.clone(),
        )
        .await
        {
            Ok(results) => {
                if cancel_token.is_cancelled() {
                    return Ok(exhausted);
                }

                for result in results {
                    match result.result {
                        Ok(()) => {}
                        Err(e) => {
                            let mut retry_chunk = result.chunk.clone();
                            retry_chunk.retry_count = chunk.retry_count + 1;
                            retry_chunk.last_peer = Some(result.peer_addr);

                            if retry_chunk.retry_count >= MAX_RETRY_COUNT {
                                let context = format!(
                                    "Retry budget exhausted for chunk: {}",
                                    result.chunk.relative_path
                                );
                                exhausted.push(ChunkDownloadResult {
                                    chunk: retry_chunk,
                                    result: Err(e.wrap_err(context)),
                                    peer_addr: result.peer_addr,
                                });
                            } else {
                                queue.push_back(retry_chunk);
                            }
                        }
                    }
                }
            }
            Err(e) => {
                if cancel_token.is_cancelled() {
                    return Ok(exhausted);
                }

                chunk.retry_count += 1;
                chunk.last_peer = Some(peer_addr);

                if chunk.retry_count >= MAX_RETRY_COUNT {
                    exhausted.push(ChunkDownloadResult {
                        chunk: chunk.clone(),
                        result: Err(e.wrap_err(format!(
                            "Retry budget exhausted for chunk after connection failure: {}",
                            chunk.relative_path
                        ))),
                        peer_addr: fallback_peer_addr(eligible_peers, chunk.last_peer),
                    });
                } else {
                    queue.push_back(chunk);
                }
            }
        }
    }

    Ok(exhausted)
}
