Fishing2/Packages/com.nwh.dynamicwaterphysics/Samples~/KWSIntegration/KWSWaterDataProvider.cs

// ╔════════════════════════════════════════════════════════════════╗
// ║    Copyright © 2025 NWH Coding d.o.o.  All rights reserved.    ║
// ║    Licensed under Unity Asset Store Terms of Service:          ║
// ║        https://unity.com/legal/as-terms                        ║
// ║    Use permitted only in compliance with the License.          ║
// ║    Distributed "AS IS", without warranty of any kind.          ║
// ╚════════════════════════════════════════════════════════════════╝

using KWS;
using NWH.DWP2.WaterObjects;
using UnityEngine;
using System.Collections.Generic;


namespace NWH.DWP2.WaterData
{
    /// <summary>
    /// KWS only supports one active request at a time. This provider batches all WaterObject
    /// queries into a single request per frame, then distributes results back.
    /// </summary>
    [DefaultExecutionOrder(-50)]
    public class KWSWaterDataProvider : WaterDataProvider
    {
        // Single shared request for batch queries
        private WaterSurfaceRequestArray _request = new WaterSurfaceRequestArray();

        // Separate request for single-point queries (Engine.Submerged, etc.)
        private WaterSurfaceRequestArray _singleRequest = new WaterSurfaceRequestArray();
        private Vector3[] _singlePointArray = new Vector3[1];
        private float _lastSingleHeight;

        // Double-buffered batching: collect current frame, submit previous frame
        private List<Vector3> _currentPoints = new List<Vector3>();
        private List<Vector3> _pendingPoints = new List<Vector3>();
        private List<(int id, int start, int count)> _currentRanges = new List<(int, int, int)>();
        private List<(int id, int start, int count)> _pendingRanges = new List<(int, int, int)>();

        // Cached results per WaterObject
        private Dictionary<int, (float[] heights, Vector3[] flows, bool valid)> _cachedResults
            = new Dictionary<int, (float[], Vector3[], bool)>();

        public override bool SupportsWaterHeightQueries() => true;
        public override bool SupportsWaterNormalQueries() => false;
        public override bool SupportsWaterFlowQueries() => true;

        // Override to prevent single-point queries from corrupting batch cache
        public override float GetWaterHeightSingle(WaterObject waterObject, Vector3 point)
        {
            _singlePointArray[0] = point;
            _singleRequest.SetNewPositions(_singlePointArray);
            WaterSystem.TryGetWaterSurfaceData(_singleRequest);

            if (_singleRequest.IsDataReady)
            {
                _lastSingleHeight = _singleRequest.Result[0].Position.y;
            }
            return _lastSingleHeight;
        }

        private void FixedUpdate()
        {
            // Check if pending request has results
            if (_pendingPoints.Count > 0 && _request.IsDataReady)
            {
                DistributeResults();
            }

            // Submit current batch (becomes pending)
            if (_currentPoints.Count > 0)
            {
                // Swap buffers
                (_currentPoints, _pendingPoints) = (_pendingPoints, _currentPoints);
                (_currentRanges, _pendingRanges) = (_pendingRanges, _currentRanges);

                _request.SetNewPositions(_pendingPoints.ToArray());
                WaterSystem.TryGetWaterSurfaceData(_request);
            }

            // Clear current for new collection
            _currentPoints.Clear();
            _currentRanges.Clear();
        }

        public override void GetWaterHeights(WaterObject waterObject, ref Vector3[] points, ref float[] waterHeights)
        {
            int id = waterObject.instanceID;

            // Add this WaterObject's points to current batch
            int startIndex = _currentPoints.Count;
            for (int i = 0; i < points.Length; i++)
            {
                _currentPoints.Add(points[i]);
            }
            _currentRanges.Add((id, startIndex, points.Length));

            // Ensure cache exists and is large enough (only grow, never shrink)
            if (!_cachedResults.TryGetValue(id, out var cached) ||
                cached.heights == null || cached.heights.Length < points.Length)
            {
                _cachedResults[id] = (new float[points.Length], new Vector3[points.Length], false);
                cached = _cachedResults[id];
            }

            // Return cached values
            if (cached.valid)
            {
                int count = Mathf.Min(cached.heights.Length, waterHeights.Length);
                for (int i = 0; i < count; i++)
                {
                    waterHeights[i] = cached.heights[i];
                }
            }
        }

        private void DistributeResults()
        {
            var results = _request.Result;
            if (results == null) return;

            // Use pending ranges (matches the pending request that just completed)
            foreach (var (id, start, count) in _pendingRanges)
            {
                if (!_cachedResults.TryGetValue(id, out var cached)) continue;
                if (start >= results.Length) continue;

                var heights = cached.heights;
                var flows = cached.flows;

                int resultCount = Mathf.Min(count, results.Length - start);
                for (int i = 0; i < resultCount; i++)
                {
                    heights[i] = results[start + i].Position.y;
                    flows[i] = results[start + i].Velocity;
                }

                _cachedResults[id] = (heights, flows, true);
            }
        }

        public override void GetWaterFlows(WaterObject waterObject, ref Vector3[] points, ref Vector3[] waterFlows)
        {
            int id = waterObject.instanceID;
            if (!_cachedResults.TryGetValue(id, out var cached) || !cached.valid || cached.flows == null)
                return;

            int count = Mathf.Min(cached.flows.Length, waterFlows.Length);
            for (int i = 0; i < count; i++)
            {
                waterFlows[i] = cached.flows[i];
            }
        }
    }
}


#if UNITY_EDITOR

namespace NWH.DWP2.WaterData
{
    using UnityEditor;

    [CustomEditor(typeof(KWSWaterDataProvider))]
    [CanEditMultipleObjects]
    public class KWSWaterDataProviderEditor : WaterDataProviderEditor
    {
        protected override void DrawStatus(WaterDataProvider provider)
        {
            drawer.BeginSubsection("Status");
            bool hasWaterSystem = WaterSystem.Instance != null;
            if (hasWaterSystem)
            {
                drawer.Info("KWS Water System found in scene.");
            }
            else
            {
                drawer.Info("KWS Water System not found in scene.", MessageType.Warning);
            }
            drawer.EndSubsection();
        }
    }
}

#endif