// SPDX-License-Identifier: AGPL-3.0-only

//! Authoritative chunk storage and generation logic for the server.

use bevy_ecs::prelude::Resource;
use crossbeam_channel::{Receiver, Sender};
use shared::{
    generator::VoxelGenerator,
    world::{Chunk, ChunkPos},
};
use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
    thread::JoinHandle,
};

/// Outcome of a single streaming reconcile pass, for logging and tests.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct StreamStats {
    /// Number of chunks generated and inserted this pass.
    pub loaded: usize,
    /// Number of chunks evicted this pass.
    pub unloaded: usize,
    /// Number of chunks resident after the pass.
    pub resident: usize,
    /// Number of chunks in flight.
    pub in_flight: usize,
}

/// The server's authoritative representation of the world.
#[derive(Resource)]
pub struct ServerWorld {
    /// Currently resident chunks, keyed by chunk-space position.
    chunks: HashMap<ChunkPos, Chunk>,
    /// Sending end of the job channel; the main thread pushes positions that require generation.
    job_tx: Sender<ChunkPos>,
    /// Receiving end of the result channel; the main thread drains finished chunks returned by workers.
    result_rx: Receiver<(ChunkPos, Chunk)>,
    /// Positions dispatched to a worker but not yet returned, preventing the same chunk from being re-dispatched on subsequent passes.
    in_flight: HashSet<ChunkPos>,
    /// Handles to the generation worker threads, retained so they can be joined on shutdown.
    // Retained ahead of a dedicated shutdown path; not yet read because the server has no graceful-stop sequence.
    #[expect(dead_code)]
    workers: Vec<JoinHandle<()>>,
}

impl ServerWorld {
    /// Initializes a new authoritative server world with the provided generator.
    #[must_use]
    pub fn new(generator: VoxelGenerator) -> Self {
        let generator = Arc::new(generator);
        let (job_tx, job_rx) = crossbeam_channel::unbounded::<ChunkPos>();
        let (result_tx, result_rx) = crossbeam_channel::unbounded::<(ChunkPos, Chunk)>();

        // One worker per logical core, falling back to a small fixed pool if the platform cannot report its parallelism.
        let worker_count = std::thread::available_parallelism().map_or(4, std::num::NonZero::get);

        let workers = (0..worker_count)
            .map(|_| {
                // Each worker owns its own clones: a shared handle to the read-only generator, its own view of the shared job queue, and its own sender back into the result channel.
                let generator = Arc::clone(&generator);
                let job_rx = job_rx.clone();
                let result_tx = result_tx.clone();

                std::thread::spawn(move || {
                    // Block until a job arrives.
                    while let Ok(pos) = job_rx.recv() {
                        let chunk = generator.generate_chunk(pos);
                        // A send error means the main thread has gone away; nothing left to do but let the worker wind down.
                        if result_tx.send((pos, chunk)).is_err() {
                            break;
                        }
                    }
                })
            })
            .collect();

        // Drop the template ends left over after cloning.
        drop(job_rx);
        drop(result_tx);

        Self {
            chunks: HashMap::new(),
            job_tx,
            result_rx,
            in_flight: HashSet::new(),
            workers,
        }
    }

    /// Number of chunks currently resident in memory.
    #[must_use]
    pub fn loaded_count(&self) -> usize {
        self.chunks.len()
    }

    /// Number of chunks dispatched to the worker pool but not yet returned.
    #[must_use]
    pub fn in_flight_count(&self) -> usize {
        self.in_flight.len()
    }

    /// Returns `true` when the worker pool has no outstanding generation work, i.e. every dispatched chunk has been returned. A loading gate can poll this to decide when an initial region has finished streaming.
    #[must_use]
    pub fn streaming_idle(&self) -> bool {
        self.in_flight.is_empty()
    }

    /// Reconciles resident chunks against a desired set without blocking the caller: finished chunks are drained from the worker pool, resident chunks absent from `desired` are evicted, and still-missing chunks are dispatched to the pool.
    pub fn reconcile(&mut self, desired: &HashSet<ChunkPos>) -> StreamStats {
        // Drain: absorb every chunk the workers have finished since the last pass.
        let mut loaded = 0;
        while let Ok((pos, chunk)) = self.result_rx.try_recv() {
            // The position is no longer dispatched now that its chunk has returned.
            self.in_flight.remove(&pos);
            // Guard against the eviction race: the anchor may have moved away while this chunk was generating, so a returned chunk is only kept if it is still wanted.
            if desired.contains(&pos) {
                self.chunks.insert(pos, chunk);
                loaded += 1;
            }
        }

        // Evict: drop resident chunks that no anchor wants any more.
        let stale: Vec<ChunkPos> = self
            .chunks
            .keys()
            .filter(|pos| !desired.contains(pos))
            .copied()
            .collect();
        for pos in &stale {
            self.chunks.remove(pos);
        }

        // Dispatch: request generation for every wanted position that is neither resident nor already in flight.
        for &pos in desired {
            if !self.chunks.contains_key(&pos) && !self.in_flight.contains(&pos) {
                self.in_flight.insert(pos);
                // A send error means every worker has shut down; there is nothing useful to do with the position, so the failure is ignored.
                let _ = self.job_tx.send(pos);
            }
        }

        StreamStats {
            loaded,
            unloaded: stale.len(),
            resident: self.chunks.len(),
            in_flight: self.in_flight.len(),
        }
    }
}

/// Inserts every chunk position inside the streaming cylinder around `center` into `out`.
pub fn cylinder_chunks<S: std::hash::BuildHasher>(
    center: ChunkPos,
    radius: i32,
    out: &mut HashSet<ChunkPos, S>,
) {
    for x in center.x - radius..=center.x + radius {
        for z in center.z - radius..=center.z + radius {
            let dx = x - center.x;
            let dz = z - center.z;

            // Keep only the columns whose XZ distance falls within the disc.
            if dx * dx + dz * dz <= radius * radius {
                for y in center.y - radius / 2..=center.y + radius / 2 {
                    out.insert(ChunkPos::new(x, y, z));
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{ChunkPos, ServerWorld, StreamStats, cylinder_chunks};
    use shared::generator::{VoxelGenerator, WorldGenConfig};
    use shared::world::BlockId;
    use std::collections::HashSet;
    use std::time::{Duration, Instant};

    /// Builds a server world backed by a real worker pool for streaming tests.
    fn test_world() -> ServerWorld {
        let config = WorldGenConfig {
            base_height: 8,
            noise_scale: 0.05,
            surface_block: BlockId(1),
            subsurface_block: BlockId(2),
            stone_block: BlockId(3),
        };
        ServerWorld::new(VoxelGenerator::new(config, 42))
    }

    /// Repeatedly reconciles against `desired` until the worker pool reports no outstanding work, returning the final pass's stats. Fails the test if the pool does not drain within a fixed timeout.
    fn drain_to_idle(world: &mut ServerWorld, desired: &HashSet<ChunkPos>) -> StreamStats {
        let deadline = Instant::now() + Duration::from_secs(5);
        loop {
            let stats = world.reconcile(desired);
            if stats.in_flight == 0 {
                return stats;
            }
            assert!(
                Instant::now() < deadline,
                "worker pool did not drain in time"
            );
            std::thread::sleep(Duration::from_millis(1));
        }
    }

    #[test]
    fn reconcile_converges_over_multiple_passes() {
        let mut world = test_world();
        let mut desired = HashSet::new();
        cylinder_chunks(ChunkPos::new(0, 0, 0), 2, &mut desired);

        // The first pass only dispatches work; because generation is off-thread, nothing is resident yet and every position is in flight.
        let first = world.reconcile(&desired);
        assert_eq!(first.loaded, 0);
        assert_eq!(first.resident, 0);
        assert!(first.in_flight > 0);

        // Later passes drain finished chunks until the pool is idle, at which point every desired position must be resident.
        let final_stats = drain_to_idle(&mut world, &desired);
        assert_eq!(final_stats.in_flight, 0);
        assert_eq!(final_stats.resident, desired.len());
    }

    #[test]
    fn evicted_chunk_is_not_repopulated_on_arrival() {
        let mut world = test_world();
        let target = ChunkPos::new(0, 0, 0);
        let mut desired = HashSet::new();
        desired.insert(target);

        // Dispatch the chunk, then immediately stop wanting it.
        world.reconcile(&desired);

        // Every subsequent pass reconciles against an empty desired set, so whenever the worker returns the chunk the drain guard discards it rather than resurrecting an unwanted chunk.
        let empty = HashSet::new();
        let final_stats = drain_to_idle(&mut world, &empty);
        assert_eq!(final_stats.in_flight, 0);
        assert_eq!(
            final_stats.resident, 0,
            "a chunk that is no longer wanted must not become resident when it arrives"
        );
    }

    #[test]
    fn cylinder_is_symmetric_and_bounded() {
        let mut set = HashSet::new();
        cylinder_chunks(ChunkPos::new(0, 0, 0), 2, &mut set);

        // Center column is always included.
        assert!(set.contains(&ChunkPos::new(0, 0, 0)));
        // A corner outside the disc (dx=2, dz=2 -> 8 > 4) is excluded.
        assert!(!set.contains(&ChunkPos::new(2, 0, 2)));
        // An axis cell at exactly the radius is included (dx=2, dz=0 -> 4 == 4).
        assert!(set.contains(&ChunkPos::new(2, 0, 0)));
        // Vertical extent is radius/2 = 1, so y=2 is out of range.
        assert!(!set.contains(&ChunkPos::new(0, 2, 0)));
        assert!(set.contains(&ChunkPos::new(0, 1, 0)));
    }

    #[test]
    fn cylinder_translates_with_center() {
        let mut origin = HashSet::new();
        cylinder_chunks(ChunkPos::new(0, 0, 0), 3, &mut origin);

        let mut shifted = HashSet::new();
        cylinder_chunks(ChunkPos::new(10, 0, -5), 3, &mut shifted);

        // Shape is translation-invariant: same count regardless of center.
        assert_eq!(origin.len(), shifted.len());
    }
}