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

//! Client-side chunk streaming around the camera.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use renderer::vertex::Vertex;
use renderer::{MeshKey, RendererError};
use shared::protocol::chunk::ChunkMessage;
use shared::world::{CHUNK_SIZE, Chunk, ChunkPos};
use tracing::error;

use crate::mesh_pool::{JobGen, MeshJob, MeshPool, MeshResult};

/// Radius, in chunks, of the region kept resident around the camera center. Also the radius the client subscribes with, so the server's resident set matches the client's.
// TODO: make configurable / drive from view-distance setting.
pub const LOAD_RADIUS: i32 = 8;

/// Maximum number of chunk deliveries materialized in a single call to [`ChunkManager::update`], bounding per-frame materialization work. Deliveries beyond the budget remain queued in the transport for the next frame.
const LOADS_PER_UPDATE: usize = 4;

/// Maximum number of mesh jobs dispatched to the worker pool per call to [`ChunkManager::update`], draining the pending re-mesh set under a bound so a burst of deliveries does not flood the pool in a single frame. One delivery can enqueue up to seven mesh jobs (itself plus six neighbours), so this budget exceeds [`LOADS_PER_UPDATE`]. Finished meshes are ingested without a per-frame bound, since uploading already-computed geometry is cheap relative to generating it.
const MESHES_PER_UPDATE: usize = 16;

/// The six face-adjacent neighbour offsets, in chunk coordinates. The order matches the neighbour array carried by [`MeshJob`]: `[+X, -X, +Y, -Y, +Z, -Z]`.
const NEIGHBOR_OFFSETS: [(i32, i32, i32); 6] = [
    (1, 0, 0),
    (-1, 0, 0),
    (0, 1, 0),
    (0, -1, 0),
    (0, 0, 1),
    (0, 0, -1),
];

/// Sink that receives finished chunk meshes for upload.
///
/// The production sink is the Vulkan [`Renderer`](renderer::Renderer); the abstraction exists so the ingest pipeline can be exercised against a recording double in tests, which have no GPU. Method signatures mirror the renderer's exactly so the production `impl` is a direct forward.
pub trait MeshSink {
    /// Uploads (or replaces) the mesh identified by `key`.
    ///
    /// # Errors
    ///
    /// Returns [`RendererError`] when the underlying implementation fails to allocate or write the GPU buffers for the mesh.
    fn insert_mesh(
        &mut self,
        key: MeshKey,
        vertices: &[Vertex],
        indices: &[u32],
        world_offset: [f32; 3],
    ) -> Result<(), RendererError>;

    /// Removes any mesh currently associated with `key`; a no-op when none exists.
    fn remove_mesh(&mut self, key: MeshKey);
}

impl MeshSink for renderer::Renderer {
    fn insert_mesh(
        &mut self,
        key: MeshKey,
        vertices: &[Vertex],
        indices: &[u32],
        world_offset: [f32; 3],
    ) -> Result<(), RendererError> {
        renderer::Renderer::insert_mesh(self, key, vertices, indices, world_offset)
    }

    fn remove_mesh(&mut self, key: MeshKey) {
        renderer::Renderer::remove_mesh(self, key);
    }
}

/// Approximate resident footprint of one chunk, in bytes: a full `CHUNK_SIZE³` volume of `u16` block identifiers.
///
/// Ignores the `HashMap` entry and [`Arc`] header overheads, which are negligible beside the volume itself.
const CHUNK_RESIDENT_BYTES: u64 = (CHUNK_SIZE as u64).pow(3) * 2;

/// A snapshot of the streaming pipeline's state and cumulative throughput.
///
/// Every field is a plain value copied out of [`ChunkManager`] at the moment of the call; nothing is retained or shared, so a reader on a slower cadence than the frame loop observes one self-consistent instant.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct ChunkStats {
    /// Chunks currently held in memory, each retaining its full voxel volume.
    pub resident: usize,
    /// Positions queued for re-meshing but not yet dispatched to the worker pool.
    pub pending_remesh: usize,
    /// Positions with a mesh job outstanding in the worker pool.
    pub in_flight: usize,
    /// Streaming radius, in chunks, currently in force.
    pub load_radius: i32,
    /// Size of the set of positions the streaming region wants resident. Residency lagging behind this figure indicates the server has not yet delivered the remainder.
    pub desired: usize,
    /// Chunks materialized from server deliveries since startup.
    pub loaded_total: u64,
    /// Chunks removed on the server's authoritative drop since startup.
    pub dropped_total: u64,
    /// Chunks evicted by the client's own radius check since startup, independently of any server drop.
    pub evicted_total: u64,
    /// Mesh jobs handed to the worker pool since startup.
    pub dispatched_total: u64,
    /// Finished meshes that were still current on return and were therefore uploaded since startup. The shortfall against `dispatched_total` is work superseded by a newer job or invalidated by eviction.
    pub applied_total: u64,
    /// Worker threads in the meshing pool.
    pub mesh_workers: usize,
    /// Estimated memory held by the resident chunk set, in bytes.
    pub resident_bytes: u64,
}

/// Tracks which server-streamed chunks are resident and orchestrates neighbour-aware background meshing.
pub struct ChunkManager {
    /// Resident chunks keyed by position, retained so the mesher can sample voxels across chunk boundaries. Stored behind [`Arc`] so a chunk can be handed to a worker thread without copying its 64 KiB volume.
    // TODO: a resident Chunk is 32³ × 2 bytes = 64 KiB; at LOAD_RADIUS = 8 the resident set is thousands of chunks (hundreds of MiB). A follow-up can store only the six 32×32 boundary planes per chunk instead of the full volume.
    resident: HashMap<ChunkPos, Arc<Chunk>>,
    /// Positions whose mesh must be rebuilt, accumulated across frames and dispatched under [`MESHES_PER_UPDATE`]. Held as a set so a burst of deliveries re-meshes each affected neighbour at most once.
    pending_remesh: HashSet<ChunkPos>,
    /// Positions with a mesh job currently outstanding, mapped to the generation of that job. A returned mesh is applied only when its generation still matches, so meshes superseded by a re-dispatch (or by eviction) are discarded rather than uploaded stale.
    in_flight: HashMap<ChunkPos, JobGen>,
    /// The generation stamped on the next dispatched job. Global and strictly increasing across all positions, so no two dispatches ever share a token; see [`JobGen`].
    next_gen: JobGen,
    /// Background worker pool that turns chunks into CPU geometry off the winit thread.
    pool: MeshPool,
    /// Reused all-air baseline that server [`ChunkData`](shared::world::ChunkData) diffs are materialized against.
    baseline: Chunk,
    /// Running totals of pipeline throughput since startup, reported through [`ChunkManager::stats`].
    totals: ChunkTotals,
}

/// Cumulative counts of the work the streaming pipeline has performed since startup.
///
/// Kept as a separate struct so the per-frame counters already computed inside [`ChunkManager::update`] fold into one place rather than becoming five loose fields on the manager.
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
struct ChunkTotals {
    /// Chunks materialized from server deliveries.
    loaded: u64,
    /// Chunks removed on the server's authoritative drop.
    dropped: u64,
    /// Chunks evicted by the client's own radius check.
    evicted: u64,
    /// Mesh jobs handed to the worker pool.
    dispatched: u64,
    /// Finished meshes uploaded because they were still current on return.
    applied: u64,
}

impl ChunkTotals {
    /// Folds one frame's per-category counts into the running totals.
    ///
    /// Saturating addition is used throughout: these are monotonic diagnostic counters, and pinning them at `u64::MAX` is preferable to an overflow panic in the frame loop. Reaching the bound would require more chunk operations than any session performs.
    fn accumulate(
        &mut self,
        loaded: usize,
        dropped: usize,
        evicted: usize,
        dispatched: usize,
        applied: usize,
    ) {
        self.loaded = self.loaded.saturating_add(loaded as u64);
        self.dropped = self.dropped.saturating_add(dropped as u64);
        self.evicted = self.evicted.saturating_add(evicted as u64);
        self.dispatched = self.dispatched.saturating_add(dispatched as u64);
        self.applied = self.applied.saturating_add(applied as u64);
    }
}

impl ChunkManager {
    /// Creates a manager with no chunks yet resident, spawning the background mesh worker pool.
    #[must_use]
    pub fn new() -> Self {
        Self {
            resident: HashMap::new(),
            pending_remesh: HashSet::new(),
            in_flight: HashMap::new(),
            next_gen: JobGen::FIRST,
            pool: MeshPool::new(),
            baseline: Chunk::default(),
            totals: ChunkTotals::default(),
        }
    }

    /// Snapshots the streaming pipeline's current state and cumulative throughput.
    ///
    /// `center` is the chunk the streaming region is currently anchored to, and is needed only to size the desired set; it is not retained.
    #[must_use]
    pub fn stats(&self, center: ChunkPos) -> ChunkStats {
        ChunkStats {
            resident: self.resident.len(),
            pending_remesh: self.pending_remesh.len(),
            in_flight: self.in_flight.len(),
            load_radius: LOAD_RADIUS,
            desired: desired_chunks(center, LOAD_RADIUS).len(),
            loaded_total: self.totals.loaded,
            dropped_total: self.totals.dropped,
            evicted_total: self.totals.evicted,
            dispatched_total: self.totals.dispatched,
            applied_total: self.totals.applied,
            mesh_workers: self.pool.worker_count(),
            resident_bytes: self.resident.len() as u64 * CHUNK_RESIDENT_BYTES,
        }
    }

    /// Advances the streaming pipeline for one frame: ingests finished meshes from the pool, evicts chunks outside the load radius around `center`, applies queued server deliveries under a materialization budget, then dispatches pending re-mesh jobs under a dispatch budget.
    ///
    /// The client's own radius eviction runs independently of the server's authoritative `Drop`, so memory stays bounded even if the server is slow to drop chunks that leave the region.
    pub fn update(
        &mut self,
        center: ChunkPos,
        deliveries: &mut net::ChunkStream,
        sink: &mut impl MeshSink,
    ) {
        let applied = self.drain_results(sink);
        let unloaded = self.unload_outside(center, sink);
        let (loaded, dropped) = self.apply_deliveries(deliveries, sink);
        let dispatched = self.dispatch_pending();

        self.totals
            .accumulate(loaded, dropped, unloaded, dispatched, applied);
    }

    /// Ingests every finished mesh currently available from the pool, uploading the ones that are still current and discarding superseded or evicted ones. Returns the number uploaded.
    fn drain_results(&mut self, sink: &mut impl MeshSink) -> usize {
        let mut applied = 0;
        while let Some(result) = self.pool.poll() {
            if self.apply_result(&result, sink) {
                applied += 1;
            }
        }
        applied
    }

    /// Uploads a single finished mesh when it is still current, returning whether it was uploaded.
    ///
    /// A result is current when its position is still resident and its generation matches the latest job dispatched for that position (see [`should_apply`]). On a match the in-flight entry is cleared; otherwise the result is dropped and any newer outstanding job for the position is left untouched.
    fn apply_result(&mut self, result: &MeshResult, sink: &mut impl MeshSink) -> bool {
        if !should_apply(
            result.pos,
            result.generation,
            |pos| self.resident.contains_key(&pos),
            &self.in_flight,
        ) {
            return false;
        }
        self.in_flight.remove(&result.pos);
        upload_result(result, sink);
        true
    }

    /// Applies up to [`LOADS_PER_UPDATE`] chunk deliveries plus any interleaved drops, returning the counts of chunks loaded and dropped.
    ///
    /// Delivered chunks are materialized and retained; drops remove the chunk from residency, from any in-flight tracking, and from the renderer. Both kinds enqueue the affected neighbourhood for re-meshing.
    fn apply_deliveries(
        &mut self,
        deliveries: &mut net::ChunkStream,
        sink: &mut impl MeshSink,
    ) -> (usize, usize) {
        let mut loaded = Vec::new();
        let mut dropped = Vec::new();
        // Only chunk deliveries count against the budget; drops are cheap and always applied.
        while loaded.len() < LOADS_PER_UPDATE {
            match deliveries.try_recv() {
                Ok(ChunkMessage::Chunk { pos, data }) => {
                    let chunk = Arc::new(data.materialize(&self.baseline));
                    self.resident.insert(pos, chunk);
                    loaded.push(pos);
                }
                Ok(ChunkMessage::Drop { pos }) => {
                    if self.resident.remove(&pos).is_some() {
                        self.in_flight.remove(&pos);
                        sink.remove_mesh((pos.x, pos.y, pos.z));
                        dropped.push(pos);
                    }
                }
                // Empty or disconnected: nothing more to apply this frame.
                Err(_) => break,
            }
        }

        self.queue_remesh(&loaded, &dropped);
        (loaded.len(), dropped.len())
    }

    /// Evicts every resident chunk outside the load radius around `center`, returning the number removed.
    ///
    /// Each evicted chunk is removed from residency, from in-flight tracking, and from the renderer; its resident neighbours have a boundary toward it that is now exposed, so they are enqueued for re-meshing.
    fn unload_outside(&mut self, center: ChunkPos, sink: &mut impl MeshSink) -> usize {
        let desired = desired_chunks(center, LOAD_RADIUS);
        let stale: Vec<ChunkPos> = self
            .resident
            .keys()
            .filter(|pos| !desired.contains(pos))
            .copied()
            .collect();
        for pos in &stale {
            sink.remove_mesh((pos.x, pos.y, pos.z));
            self.resident.remove(pos);
            self.in_flight.remove(pos);
        }
        self.queue_remesh(&[], &stale);
        stale.len()
    }

    /// Adds the chunks affected by `loaded` and `dropped` to the pending re-mesh set.
    fn queue_remesh(&mut self, loaded: &[ChunkPos], dropped: &[ChunkPos]) {
        let targets = remesh_targets(loaded, dropped, |pos| self.resident.contains_key(&pos));
        self.pending_remesh.extend(targets);
    }

    /// Dispatches up to [`MESHES_PER_UPDATE`] pending re-mesh jobs to the worker pool, returning the number dispatched.
    ///
    /// Each dispatched position snapshots its chunk and current resident neighbours behind [`Arc`]s, is stamped with a fresh generation, and is recorded as in-flight (superseding any previous outstanding job for it). Positions no longer resident (dropped after being enqueued) are skipped without dispatch.
    fn dispatch_pending(&mut self) -> usize {
        // Take a bounded batch out of the set; the remainder stays queued for later frames.
        let batch: Vec<ChunkPos> = self
            .pending_remesh
            .iter()
            .take(MESHES_PER_UPDATE)
            .copied()
            .collect();

        let mut dispatched = 0;
        for pos in batch {
            self.pending_remesh.remove(&pos);
            let Some(chunk) = self.resident.get(&pos) else {
                continue;
            };
            let chunk = Arc::clone(chunk);
            let neighbors = self.neighbor_arcs(pos);
            let generation = self.bump_gen();
            self.in_flight.insert(pos, generation);
            self.pool.dispatch(MeshJob {
                pos,
                generation,
                chunk,
                neighbors,
            });
            dispatched += 1;
        }
        dispatched
    }

    /// Snapshots the six face-adjacent resident chunks of `pos` as [`Arc`] handles, ordered to match [`NEIGHBOR_OFFSETS`]. Absent neighbours are `None`.
    fn neighbor_arcs(&self, pos: ChunkPos) -> [Option<Arc<Chunk>>; 6] {
        NEIGHBOR_OFFSETS.map(|(dx, dy, dz)| {
            self.resident
                .get(&ChunkPos::new(pos.x + dx, pos.y + dy, pos.z + dz))
                .map(Arc::clone)
        })
    }

    /// Returns a fresh, never-before-used generation and advances the counter.
    fn bump_gen(&mut self) -> JobGen {
        let current = self.next_gen;
        self.next_gen = self.next_gen.next();
        current
    }
}

impl Default for ChunkManager {
    fn default() -> Self {
        Self::new()
    }
}

/// Reports whether a finished mesh should be uploaded.
///
/// A mesh is current, and therefore applied, only when its position is still wanted (resident) and the generation recorded as in-flight for that position still equals the mesh's own generation. A missing in-flight entry (the position was evicted) or a mismatched generation (a newer job superseded this one) both mean the result is stale and must be discarded.
fn should_apply(
    pos: ChunkPos,
    generation: JobGen,
    is_wanted: impl Fn(ChunkPos) -> bool,
    in_flight: &HashMap<ChunkPos, JobGen>,
) -> bool {
    is_wanted(pos) && in_flight.get(&pos) == Some(&generation)
}

/// Uploads a finished mesh to the sink, or clears the slot when the mesh is empty.
///
/// A chunk that meshes to no geometry (all air, or fully enclosed by solid neighbours) is removed from the sink rather than uploaded, since a zero-length buffer is invalid; this also clears any mesh a previous state had left there.
fn upload_result(result: &MeshResult, sink: &mut impl MeshSink) {
    let pos = result.pos;
    let key = (pos.x, pos.y, pos.z);

    if result.indices.is_empty() {
        sink.remove_mesh(key);
        return;
    }

    // Chunk coordinates and CHUNK_SIZE are small and represent exactly as f32.
    #[expect(
        clippy::cast_precision_loss,
        reason = "chunk coordinates stay well within f32's exact-integer range"
    )]
    let world_offset = {
        let size = CHUNK_SIZE as f32;
        [
            pos.x as f32 * size,
            pos.y as f32 * size,
            pos.z as f32 * size,
        ]
    };

    if let Err(e) = sink.insert_mesh(key, &result.vertices, &result.indices, world_offset) {
        error!(?pos, "failed to upload chunk mesh: {e}");
    }
}

/// Returns the six face-adjacent neighbour positions of `pos`.
fn neighbor_positions(pos: ChunkPos) -> [ChunkPos; 6] {
    NEIGHBOR_OFFSETS.map(|(dx, dy, dz)| ChunkPos::new(pos.x + dx, pos.y + dy, pos.z + dz))
}

/// Computes the deduplicated set of resident chunks whose mesh must be rebuilt after a batch of loads and drops.
///
/// A newly-loaded chunk contributes itself (when resident) and each of its resident neighbours, whose boundary toward it may now be culled. A dropped chunk contributes only its resident neighbours, whose boundary toward it is re-exposed; the dropped chunk itself is gone and is never a target. `is_resident` reports whether a position is currently resident.
fn remesh_targets(
    loaded: &[ChunkPos],
    dropped: &[ChunkPos],
    is_resident: impl Fn(ChunkPos) -> bool,
) -> HashSet<ChunkPos> {
    let mut targets = HashSet::new();
    for &pos in loaded {
        if is_resident(pos) {
            targets.insert(pos);
        }
        for neighbor in neighbor_positions(pos) {
            if is_resident(neighbor) {
                targets.insert(neighbor);
            }
        }
    }
    for &pos in dropped {
        for neighbor in neighbor_positions(pos) {
            if is_resident(neighbor) {
                targets.insert(neighbor);
            }
        }
    }
    targets
}

/// Returns the set of chunk positions within the streaming cylinder around `center`.
///
/// The region is a disc of `radius` chunks in the horizontal XZ plane and half that extent in Y, matching the flatter vertical shape of the playable world. This mirrors the server's `world_server::cylinder_chunks`.
#[must_use]
pub fn desired_chunks(center: ChunkPos, radius: i32) -> HashSet<ChunkPos> {
    let mut out = HashSet::new();
    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));
                }
            }
        }
    }
    out
}

#[cfg(test)]
#[path = "tests/chunks.rs"]
mod tests;