OrekiWoofsBeehives/OrekiWoofsBees.Common/PlantPositionRegistryModSystem2.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using Vintagestory.API.Common;
using Vintagestory.API.MathTools;
using Vintagestory.GameContent;

namespace OrekiWoofsBees.Common;

/// <summary>
/// tracks plant positions for all registered beehives.
/// Instead of each beehive scanning its radius every tick, this registry
/// incrementally scans blocks across all beehives to keep performance consistent
/// </summary>
public class PlantPositionRegistryModSystem2 : ModSystem, IPlantPositionRegistry
{
    private readonly static BlockPos staticBlockPos = new(0);

    private const int default_blocks_per_tick = 20;
    private const int tick_interval_ms = 20;

    private static readonly Dictionary<int, ScanOffsetTable> offsetTables = [];

    private long? tickListenerId;
    private readonly Dictionary<StructVec3i, BeehiveScanCursor> beehives = [];
    private readonly HashSet<StructVec3i> flowerPositions = [];
    private readonly HashSet<StructVec3i> cropPositions = [];

    // blocks below this are skipped
    private readonly Dictionary<(int X, int Z), int> soilFloorCache = [];

    private int lastScannedBeehiveIndex = 0;
    private int blocksPerTick = default_blocks_per_tick;

    public ICoreAPI? Api { get; private set; }

    public int BlocksPerTick
    {
        get => blocksPerTick;
        set => blocksPerTick = Math.Clamp(value, 0, 1000);
    }

    public event Action<BlockPos, int>? FlowerEvent;
    public event Action<BlockPos, int>? CropEvent;

    public override double ExecuteOrder() => 0.10;

    public override void Start(ICoreAPI api)
    {
        Api = api;
        if (api.Side.IsServer())
            tickListenerId = api.Event.RegisterGameTickListener(OnTick, tick_interval_ms);
    }

    public override void Dispose()
    {
        if (Api != null && tickListenerId.HasValue)
            Api.Event.UnregisterGameTickListener(tickListenerId.Value);
        beehives.Clear();
        flowerPositions.Clear();
        cropPositions.Clear();
        FlowerEvent = null;
        CropEvent = null;
    }

    public void RegisterBeehive(BlockPos pos, int radius)
    {
        var key = StructVec3i.FromBlockPos(pos);
        if (beehives.ContainsKey(key))
            return;

        if (!offsetTables.ContainsKey(radius))
            offsetTables[radius] = new ScanOffsetTable(radius);

        beehives[key] = new BeehiveScanCursor(pos, radius);
    }

    public void UnregisterBeehive(BlockPos pos)
    {
        var key = StructVec3i.FromBlockPos(pos);
        beehives.Remove(key);
    }

    public (
        float InitialScanProgress,
        float RescanProgress
    ) GetPlantsNearPosition(BlockPos hivePos, int radius, List<StructVec3i> flowerPositionsBuffer, List<StructVec3i> cropPositionsBuffer)
    {
        flowerPositionsBuffer.Clear();
        cropPositionsBuffer.Clear();

        var key = StructVec3i.FromBlockPos(hivePos);

        foreach (var flowerPos in flowerPositions)
        {
            if (!Overlaps.IsWithinSphericalRadius(hivePos, flowerPos, radius))
                continue;
            flowerPositionsBuffer.Add(flowerPos);
        }

        foreach (var cropPos in cropPositions)
        {
            if (!Overlaps.IsWithinSphericalRadius(hivePos, cropPos, radius))
                continue;
            cropPositionsBuffer.Add(cropPos);
        }

        float initialProgress = 1.0f;
        float rescanProgress = 0.0f;
        if (!beehives.TryGetValue(key, out var cursor))
            return (initialProgress, rescanProgress);

        var table = offsetTables[cursor.Radius];
        if (table.Count > 0)
        {
            initialProgress = Math.Min(1.0f, (float)cursor.BlocksCheckedCount / table.Count);
            rescanProgress = cursor.GetRescanProgress(offsetTables);
        }

        return (initialProgress, rescanProgress);
    }

    public (int FlowerCount, int CropCount, float InitialScanProgress, float RescanProgress) GetPlantCountsNearPosition(BlockPos hivePos, int radius)
    {
        int flowers = 0;
        foreach (var flowerPos in flowerPositions)
        {
            if (Overlaps.IsWithinSphericalRadius(hivePos, flowerPos, radius))
                flowers++;
        }

        int crops = 0;
        foreach (var cropPos in cropPositions)
        {
            if (Overlaps.IsWithinSphericalRadius(hivePos, cropPos, radius))
                crops++;
        }

        float initialProgress = 0.0f;
        float rescanProgress = 0.0f;
        var key = StructVec3i.FromBlockPos(hivePos);
        if (!beehives.TryGetValue(key, out var cursor))
            return (flowers, crops, initialProgress, rescanProgress);

        var table = offsetTables[cursor.Radius];
        if (table.Count > 0)
        {
            initialProgress = Math.Min(1.0f, (float)cursor.BlocksCheckedCount / table.Count);
            rescanProgress = cursor.GetRescanProgress(offsetTables);
        }

        return (flowers, crops, initialProgress, rescanProgress);
    }

    public void AddPlantPosition(BlockPos pos, Block block)
    {
        if (Api is null)
            return;

        var structPos = new StructVec3i(pos.X, pos.Y, pos.Z);

        if (PlantRecognitionUtilities.IsCrop(block))
        {
            cropPositions.Add(structPos);
            flowerPositions.Remove(structPos);
            CropEvent?.Invoke(pos, 1);
        }
        else if (PlantRecognitionUtilities.IsFlower(block, Api.World.BlockAccessor, pos))
        {
            flowerPositions.Add(structPos);
            cropPositions.Remove(structPos);
            FlowerEvent?.Invoke(pos, 1);
        }
    }

    public void RemovePlantPosition(BlockPos pos, Block block)
    {
        if (Api is null)
            return;

        var structPos = new StructVec3i(pos.X, pos.Y, pos.Z);
        flowerPositions.Remove(structPos);
        cropPositions.Remove(structPos);

        if (PlantRecognitionUtilities.IsCrop(block))
            CropEvent?.Invoke(pos, -1);
        else if (PlantRecognitionUtilities.IsFlower(block, Api.World.BlockAccessor, pos))
            FlowerEvent?.Invoke(pos, -1);
    }

    public int CountBeehivesInRadius(BlockPos pos, int radius)
    {
        var plantPos = StructVec3i.FromBlockPos(pos);

        int count = 0;
        foreach (var (beehivePos, _) in beehives)
        {
            if (Overlaps.IsWithinSphericalRadius(beehivePos, plantPos, radius))
                count++;
        }

        return count;
    }

    private void OnTick(float dt)
    {
        if (Api is null)
            return;

        var stopwatch = Stopwatch.StartNew();
        if (beehives.Count == 0)
            return;

        var accessor = Api.World.BlockAccessor;

        for (int i = 0; i < BlocksPerTick; i++)
        {
            var nextBlock = GetNextBlockToCheck();
            if (nextBlock == null)
                break;

            var (cursor, blockPos) = nextBlock.Value;

            // check if this position is below the soil floor for this X/Z
            var xz = (blockPos.X, blockPos.Z);
            if (soilFloorCache.TryGetValue(xz, out int soilFloorY) && blockPos.Y < soilFloorY)
            {
                // it's below the soil floor, skip
                cursor.Advance();
                continue;
            }

            cursor.Advance();

            var block = accessor.GetBlock(staticBlockPos.Set(blockPos.X, blockPos.Y, blockPos.Z));

            if (block == null || block.BlockId == 0)
            {
                // block is air or unloaded - remove from caches if present
                flowerPositions.Remove(blockPos);
                cropPositions.Remove(blockPos);
                continue;
            }

            // check if this is soil with soil below - mark as soil floor
            if (block is BlockSoil)
            {
                var blockBelow = accessor.GetBlockRaw(blockPos.X, blockPos.Y - 1, blockPos.Z);
                if (blockBelow is BlockSoil)
                {
                    // found soil floor - record it
                    bool isNewFloor = false;
                    if (!soilFloorCache.TryGetValue(xz, out int existingFloor) || blockPos.Y < existingFloor)
                    {
                        soilFloorCache[xz] = blockPos.Y;
                        isNewFloor = true;
                    }

                    if (isNewFloor)
                        cursor.CountAndSkipBlocksBelowY(blockPos.Y, xz, offsetTables);
                }
            }

            if (PlantRecognitionUtilities.IsCrop(block))
            {
                cropPositions.Add(blockPos);
                flowerPositions.Remove(blockPos);
            }
            else if (PlantRecognitionUtilities.IsFlower(block, accessor, blockPos))
            {
                flowerPositions.Add(blockPos);
                cropPositions.Remove(blockPos);
            }
            else
            {
                flowerPositions.Remove(blockPos);
                cropPositions.Remove(blockPos);
            }
        }

        stopwatch.Stop();
        if (stopwatch.Elapsed.TotalSeconds > 0.2)
        {
            Mod.Logger.Warning($"{nameof(PlantPositionRegistryModSystem2)} {nameof(OnTick)} took {stopwatch.Elapsed.TotalSeconds:F2}s (beehives: {beehives.Count}).");
        }
    }

    /// <summary>
    /// 1. pick the beehive with the least blocks checked in its radius
    /// 2. pick the closest horizontal space vertically
    /// 3. check the next closest block to the beehive in this horizontal space
    /// </summary>
    private (BeehiveScanCursor Cursor, StructVec3i BlockPos)? GetNextBlockToCheck()
    {
        if (beehives.Count == 0)
            return null;

        bool allCompleted = true;
        foreach (var cursor in beehives.Values)
        {
            var table = offsetTables[cursor.Radius];
            if (cursor.BlocksCheckedCount < table.Count)
            {
                allCompleted = false;
                break;
            }
        }

        BeehiveScanCursor? bestCursor = null;

        if (!allCompleted)
        {
            var lowestRelativeProgress = float.MaxValue;
            foreach (var cursor in beehives.Values)
            {
                var table = offsetTables[cursor.Radius];
                if (cursor.BlocksCheckedCount >= table.Count)
                    continue;

                var relativeProgress = (float)cursor.BlocksCheckedCount / table.Count;
                if (relativeProgress < lowestRelativeProgress)
                {
                    lowestRelativeProgress = relativeProgress;
                    bestCursor = cursor;
                }
            }
        }
        else
        {
            var beehivesList = beehives.Values;
            lastScannedBeehiveIndex = (lastScannedBeehiveIndex + 1) % beehivesList.Count;
            bestCursor = beehivesList.ElementAt(lastScannedBeehiveIndex);
        }

        if (bestCursor == null)
            return null;

        if (bestCursor.IsStartingNewCycle())
            soilFloorCache.Clear();

        var blockPos = bestCursor.GetCurrentBlockPos(offsetTables);
        return (bestCursor, blockPos);
    }

    private class BeehiveScanCursor(BlockPos hivePos, int radius)
    {
        public BlockPos HivePos { get; } = hivePos;
        public int Radius { get; } = radius;
        public int BlocksCheckedCount { get; private set; } = 0;

        private int currentIndex = 0;

        public StructVec3i GetCurrentBlockPos(Dictionary<int, ScanOffsetTable> tables)
        {
            var table = tables[Radius];
            var (X, Y, Z) = table.GetOffset(currentIndex);
            return new(
                HivePos.X + X,
                HivePos.Y + Y,
                HivePos.Z + Z
            );
        }

        public void Advance()
        {
            var table = offsetTables[Radius];
            currentIndex++;

            if (currentIndex >= table.Count)
                currentIndex = 0;

            if (BlocksCheckedCount < table.Count)
                BlocksCheckedCount++;
        }

        public bool IsStartingNewCycle()
        {
            return currentIndex == 0 && BlocksCheckedCount >= offsetTables[Radius].Count;
        }

        public int GetCurrentIndex()
        {
            return currentIndex;
        }

        public void CountAndSkipBlocksBelowY(int floorY, (int X, int Z) xz, Dictionary<int, ScanOffsetTable> tables)
        {
            var table = tables[Radius];
            int skippedCount = 0;

            for (int i = currentIndex + 1; i < table.Count; i++)
            {
                var (offsetX, offsetY, offsetZ) = table.GetOffset(i);
                int worldX = HivePos.X + offsetX;
                int worldZ = HivePos.Z + offsetZ;
                int worldY = HivePos.Y + offsetY;

                if (worldX == xz.X && worldZ == xz.Z && worldY < floorY)
                {
                    skippedCount++;
                    if (BlocksCheckedCount < table.Count)
                        BlocksCheckedCount++;
                }
            }
        }

        public float GetRescanProgress(Dictionary<int, ScanOffsetTable> tables)
        {
            var table = tables[Radius];
            if (table.Count == 0)
                return 0.0f;

            if (BlocksCheckedCount < table.Count)
                return 0.0f;

            return (float)currentIndex / table.Count;
        }
    }

    /// <summary>
    /// shared table of offsets for a given radius, sorted by priority.
    /// only one instance per radius is created and shared across all beehives.
    /// </summary>
    private class ScanOffsetTable
    {
        private readonly (int X, int Y, int Z)[] offsets;
        public int Count => offsets.Length;

        public ScanOffsetTable(int radius)
        {
            int radiusSq = radius * radius;
            var offsets = new List<(int X, int Y, int Z, int YDist, int HorizontalDistSq)>();

            for (int dy = -radius; dy <= radius; dy++)
            {
                for (int dx = -radius; dx <= radius; dx++)
                {
                    for (int dz = -radius; dz <= radius; dz++)
                    {
                        int distSq = dx * dx + dy * dy + dz * dz;
                        if (distSq <= radiusSq)
                        {
                            int yDist = Math.Abs(dy);
                            int horizontalDistSq = dx * dx + dz * dz;
                            offsets.Add((dx, dy, dz, yDist, horizontalDistSq));
                        }
                    }
                }
            }

            this.offsets = [.. offsets
                .OrderBy(o => o.YDist)
                .ThenBy(o => o.HorizontalDistSq)
                .Select(o => (o.X, o.Y, o.Z))];
        }

        public (int X, int Y, int Z) GetOffset(int index)
        {
            return offsets[index];
        }
    }
}