//! Download pipeline for game files from peers.

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

use lanspread_db::db::GameFileDescription;
use tokio::{
    fs::OpenOptions,
    io::{AsyncSeekExt, AsyncWriteExt},
    sync::mpsc::UnboundedSender,
};
use tokio_util::codec::{FramedWrite, LengthDelimitedCodec};

use crate::{
    PeerEvent,
    config::{CHUNK_SIZE, MAX_RETRY_COUNT},
    network::connect_to_peer,
    path_validation::validate_game_file_path,
};

// =============================================================================
// Download data structures
// =============================================================================

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

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

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

// =============================================================================
// Storage preparation
// =============================================================================

/// Prepares storage for game files by creating directories and pre-allocating files.
pub async fn prepare_game_storage(
    games_folder: &Path,
    file_descs: &[GameFileDescription],
) -> eyre::Result<()> {
    for desc in file_descs {
        // Validate the path to prevent directory traversal
        let validated_path = validate_game_file_path(games_folder, &desc.relative_path)?;

        if desc.is_dir {
            tokio::fs::create_dir_all(&validated_path).await?;
        } else {
            if let Some(parent) = validated_path.parent() {
                tokio::fs::create_dir_all(parent).await?;
            }

            // Create and pre-allocate the file with the expected size
            let file = OpenOptions::new()
                .create(true)
                .truncate(true)
                .write(true)
                .open(&validated_path)
                .await?;

            // Pre-allocate the file with the expected size
            let size = desc.size;
            if let Err(e) = file.set_len(size).await {
                log::warn!(
                    "Failed to pre-allocate file {} (size: {}): {}",
                    desc.relative_path,
                    size,
                    e
                );
                // Continue without pre-allocation - the file will grow as chunks are written
            } else {
                log::debug!(
                    "Pre-allocated file {} with {} bytes",
                    desc.relative_path,
                    size
                );
            }
        }
    }
    Ok(())
}

// =============================================================================
// Peer plan building
// =============================================================================

/// Resolves which peers have a specific file.
pub 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 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 peer_index = 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 = eligible_peers[peer_index % eligible_peers.len()];
            peer_index += 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 = eligible_peers[peer_index % eligible_peers.len()];
            peer_index += 1;
            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
}

// =============================================================================
// Chunk downloading
// =============================================================================

/// Downloads a single chunk from a peer.
pub async fn download_chunk(
    conn: &mut s2n_quic::Connection,
    base_dir: &Path,
    game_id: &str,
    chunk: &DownloadChunk,
) -> eyre::Result<()> {
    use futures::SinkExt;
    use lanspread_proto::{Message, Request};

    let stream = conn.open_bidirectional_stream().await?;
    let (mut rx, tx) = stream.split();
    let mut framed_tx = FramedWrite::new(tx, LengthDelimitedCodec::new());

    let request = Request::GetGameFileChunk {
        game_id: game_id.to_string(),
        relative_path: chunk.relative_path.clone(),
        offset: chunk.offset,
        length: chunk.length,
    };
    framed_tx.send(request.encode()).await?;

    framed_tx.close().await?;

    // Validate the path to prevent directory traversal
    let validated_path = validate_game_file_path(base_dir, &chunk.relative_path)?;
    let mut file = OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(false)
        .open(&validated_path)
        .await?;
    if chunk.length == 0 && chunk.offset == 0 {
        // fallback-to-whole-file path replaces any existing partial data
        file.set_len(0).await?;
    }
    file.seek(std::io::SeekFrom::Start(chunk.offset)).await?;

    let mut remaining = chunk.length;
    let mut received_bytes = 0u64;

    while let Some(bytes) = rx.receive().await? {
        file.write_all(&bytes).await?;
        received_bytes += bytes.len() as u64;

        if remaining == 0 {
            continue;
        }
        remaining = remaining.saturating_sub(bytes.len() as u64);
        if remaining == 0 {
            break;
        }
    }

    // Verify we received the expected amount of data
    if chunk.length > 0 && received_bytes != chunk.length {
        eyre::bail!(
            "Incomplete chunk download: expected {} bytes, received {} bytes for file {} at offset {}",
            chunk.length,
            received_bytes,
            chunk.relative_path,
            chunk.offset
        );
    }

    file.flush().await?;

    // Verify file integrity by checking the file size
    verify_chunk_integrity(&validated_path, chunk.offset, chunk.length).await?;

    Ok(())
}

/// Verifies that a chunk was written correctly.
async fn verify_chunk_integrity(
    file_path: &Path,
    offset: u64,
    expected_length: u64,
) -> eyre::Result<()> {
    if expected_length == 0 {
        return Ok(()); // Skip verification for whole files or zero-length chunks
    }

    let metadata = tokio::fs::metadata(file_path).await?;
    let file_size = metadata.len();

    if file_size < offset + expected_length {
        eyre::bail!(
            "File integrity check failed: file size {} is less than expected {} (offset: {})",
            file_size,
            offset + expected_length,
            offset
        );
    }

    Ok(())
}

/// Downloads a whole file from a peer.
pub async fn download_whole_file(
    conn: &mut s2n_quic::Connection,
    base_dir: &Path,
    desc: &GameFileDescription,
) -> eyre::Result<()> {
    use futures::SinkExt;
    use lanspread_proto::{Message, Request};

    let stream = conn.open_bidirectional_stream().await?;
    let (mut rx, tx) = stream.split();
    let mut framed_tx = FramedWrite::new(tx, LengthDelimitedCodec::new());

    framed_tx
        .send(Request::GetGameFileData(desc.clone()).encode())
        .await?;
    framed_tx.close().await?;

    // Validate the path to prevent directory traversal
    let validated_path = validate_game_file_path(base_dir, &desc.relative_path)?;
    let mut file = OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .open(&validated_path)
        .await?;
    file.seek(std::io::SeekFrom::Start(0)).await?;

    while let Some(bytes) = rx.receive().await? {
        file.write_all(&bytes).await?;
    }

    file.flush().await?;
    Ok(())
}

/// Downloads all assigned chunks and files from a single peer.
pub async fn download_from_peer(
    peer_addr: SocketAddr,
    game_id: &str,
    plan: PeerDownloadPlan,
    games_folder: PathBuf,
) -> eyre::Result<Vec<ChunkDownloadResult>> {
    if plan.chunks.is_empty() && plan.whole_files.is_empty() {
        return Ok(Vec::new());
    }

    let mut conn = connect_to_peer(peer_addr).await?;
    conn.keep_alive(true)?;
    conn.keep_alive(true)?;

    let base_dir = games_folder;
    let mut results = Vec::new();

    // Download chunks with error handling
    for chunk in &plan.chunks {
        log::info!(
            "Downloading chunk {} (offset {}, length {}) from {}",
            chunk.relative_path,
            chunk.offset,
            chunk.length,
            peer_addr
        );
        let result = download_chunk(&mut conn, &base_dir, game_id, chunk).await;
        results.push(ChunkDownloadResult {
            chunk: chunk.clone(),
            result,
            peer_addr,
        });
    }

    // Download whole files
    for desc in &plan.whole_files {
        let chunk = DownloadChunk {
            relative_path: desc.relative_path.clone(),
            offset: 0,
            length: 0, // Indicates whole file
            retry_count: 0,
            last_peer: Some(peer_addr),
        };

        let result = download_whole_file(&mut conn, &base_dir, desc).await;
        results.push(ChunkDownloadResult {
            chunk,
            result,
            peer_addr,
        });
    }

    Ok(results)
}

// =============================================================================
// Retry logic
// =============================================================================

/// 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 async fn retry_failed_chunks(
    failed_chunks: Vec<DownloadChunk>,
    peers: &[SocketAddr],
    base_dir: &Path,
    game_id: &str,
    file_peer_map: &HashMap<String, Vec<SocketAddr>>,
) -> 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() {
        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()).await {
            Ok(results) => {
                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) => {
                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);
                }
            }
        }
    }

    exhausted
}

// =============================================================================
// Main download orchestration
// =============================================================================

/// Downloads all game files from available peers.
pub async fn download_game_files(
    game_id: &str,
    game_file_descs: Vec<GameFileDescription>,
    games_folder: String,
    peers: Vec<SocketAddr>,
    file_peer_map: HashMap<String, Vec<SocketAddr>>,
    tx_notify_ui: UnboundedSender<PeerEvent>,
) -> eyre::Result<()> {
    if peers.is_empty() {
        eyre::bail!("no peers available for game {game_id}");
    }

    let base_dir = PathBuf::from(&games_folder);
    prepare_game_storage(&base_dir, &game_file_descs).await?;

    tx_notify_ui.send(PeerEvent::DownloadGameFilesBegin {
        id: game_id.to_string(),
    })?;

    let plans = build_peer_plans(&peers, &game_file_descs, &file_peer_map);

    let mut tasks = Vec::new();
    for (peer_addr, plan) in plans {
        let base_dir = base_dir.clone();
        let game_id = game_id.to_string();
        tasks.push(tokio::spawn(async move {
            download_from_peer(peer_addr, &game_id, plan, base_dir).await
        }));
    }

    let mut failed_chunks: Vec<DownloadChunk> = Vec::new();
    let mut last_err: Option<eyre::Report> = None;

    for handle in tasks {
        match handle.await {
            Ok(Ok(results)) => {
                for chunk_result in results {
                    if let Err(e) = chunk_result.result {
                        log::warn!(
                            "Failed to download chunk from {}: {e}",
                            chunk_result.peer_addr
                        );
                        if chunk_result.chunk.retry_count < MAX_RETRY_COUNT {
                            let mut retry_chunk = chunk_result.chunk;
                            retry_chunk.retry_count += 1;
                            retry_chunk.last_peer = Some(chunk_result.peer_addr);
                            failed_chunks.push(retry_chunk);
                        } else {
                            last_err = Some(eyre::eyre!(
                                "Max retries exceeded for chunk: {}",
                                chunk_result.chunk.relative_path
                            ));
                        }
                    }
                }
            }
            Ok(Err(e)) => last_err = Some(e),
            Err(e) => last_err = Some(eyre::eyre!("task join error: {e}")),
        }
    }

    // Retry failed chunks if any
    if !failed_chunks.is_empty() && !peers.is_empty() {
        log::info!("Retrying {} failed chunks", failed_chunks.len());

        let retry_results =
            retry_failed_chunks(failed_chunks, &peers, &base_dir, game_id, &file_peer_map).await;

        for chunk_result in retry_results {
            if let Err(e) = chunk_result.result {
                log::error!("Retry failed for chunk: {e}");
                last_err = Some(e);
            }
        }
    }

    if let Some(err) = last_err {
        tx_notify_ui.send(PeerEvent::DownloadGameFilesFailed {
            id: game_id.to_string(),
        })?;
        return Err(err);
    }

    log::info!("all files downloaded for game: {game_id}");
    tx_notify_ui.send(PeerEvent::DownloadGameFilesFinished {
        id: game_id.to_string(),
    })?;
    Ok(())
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    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_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}"),
                }
            }
        }
    }
}
