/*
 * Copyright 2018 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

using CartographerSharp.Common.Math;
using System;
using System.Runtime.InteropServices;
using  RobotNet10.Shared.Numbers;

namespace CartographerSharp.Mapping.D2D;

/// <summary>
/// Grid type enumeration.
/// </summary>
public enum GridType
{
    ProbabilityGrid,
    TSDF
}

/// <summary>
/// Base class for 2D grids.
/// </summary>
public abstract class Grid2D : IGrid
{
    protected MapLimits _limits;
    protected List<ushort> _correspondenceCostCells;
    protected double _minCorrespondenceCost;
    protected double _maxCorrespondenceCost;
    protected List<int> _updateIndices;
    protected (int minX, int minY, int maxX, int maxY) _knownCellsBox;
    protected double[] _valueToCorrespondenceCostTable;

    private const ushort kUnknownCorrespondenceValue = 0;

    protected Grid2D(
        MapLimits limits,
        double minCorrespondenceCost,
        double maxCorrespondenceCost,
        ValueConversionTables conversionTables)
    {
        if (minCorrespondenceCost >= maxCorrespondenceCost)
        {
            throw new ArgumentException("min_correspondence_cost must be less than max_correspondence_cost");
        }

        _limits = limits;
        _minCorrespondenceCost = minCorrespondenceCost;
        _maxCorrespondenceCost = maxCorrespondenceCost;
        _correspondenceCostCells = new List<ushort>(
            limits.CellLimits.NumXCells * limits.CellLimits.NumYCells);
        for (int i = 0; i < _correspondenceCostCells.Capacity; i++)
        {
            _correspondenceCostCells.Add(kUnknownCorrespondenceValue);
        }
        _updateIndices = [];
        // Match C++: AlignedBox2i is empty by default (min > max)
        // Use sentinel values to mark empty: minX > maxX indicates empty box
        _knownCellsBox = (int.MaxValue, int.MaxValue, int.MinValue, int.MinValue);
        _valueToCorrespondenceCostTable = conversionTables.GetConversionTable(
            maxCorrespondenceCost, minCorrespondenceCost, maxCorrespondenceCost);
    }

    protected Grid2D(
        Models.Mapping.Grid2D proto,
        ValueConversionTables conversionTables)
    {
        _limits = new MapLimits(proto.Limits);
        // Match C++: MinCorrespondenceCostFromProto/MaxCorrespondenceCostFromProto - older proto had 0,0
        if (proto.MinCorrespondenceCost == 0 && proto.MaxCorrespondenceCost == 0)
        {
            _minCorrespondenceCost = 0.1;
            _maxCorrespondenceCost = 0.9;
        }
        else
        {
            _minCorrespondenceCost = proto.MinCorrespondenceCost;
            _maxCorrespondenceCost = proto.MaxCorrespondenceCost;
        }
        if (_minCorrespondenceCost >= _maxCorrespondenceCost)
        {
            throw new ArgumentException("min_correspondence_cost must be less than max_correspondence_cost");
        }
        // Match C++: Copy cells from proto
        _correspondenceCostCells = new List<ushort>(proto.Cells?.Count ?? 0);
        if (proto.Cells != null)
        {
            foreach (var cell in proto.Cells)
            {
                _correspondenceCostCells.Add((ushort)cell);
            }
        }
        _updateIndices = [];
        // Match C++: Copy known_cells_box from proto if present
        var knownCellsBox = proto.KnownCellsBox;
        _knownCellsBox = (knownCellsBox.MinX, knownCellsBox.MinY, knownCellsBox.MaxX, knownCellsBox.MaxY);
        _valueToCorrespondenceCostTable = conversionTables.GetConversionTable(
            _maxCorrespondenceCost, _minCorrespondenceCost, _maxCorrespondenceCost);
    }

    /// <summary>
    /// Returns the limits of this Grid2D.
    /// </summary>
    public MapLimits Limits => _limits;

    /// <summary>
    /// Returns the correspondence cost cells.
    /// </summary>
    protected List<ushort> CorrespondenceCostCells => _correspondenceCostCells;

    /// <summary>
    /// Returns the update indices.
    /// </summary>
    protected List<int> UpdateIndices => _updateIndices;

    /// <summary>
    /// Returns the known cells box.
    /// </summary>
    protected (int minX, int minY, int maxX, int maxY) KnownCellsBox => _knownCellsBox;

    /// <summary>
    /// Returns mutable correspondence cost cells.
    /// </summary>
    protected List<ushort> MutableCorrespondenceCostCells => _correspondenceCostCells;

    /// <summary>
    /// Returns mutable update indices.
    /// </summary>
    protected List<int> MutableUpdateIndices => _updateIndices;

    /// <summary>
    /// Returns mutable known cells box.
    /// </summary>
    protected ref (int minX, int minY, int maxX, int maxY) MutableKnownCellsBox => ref _knownCellsBox;

    /// <summary>
    /// Returns the minimum possible correspondence cost.
    /// </summary>
    public double MinCorrespondenceCost => _minCorrespondenceCost;

    /// <summary>
    /// Returns the maximum possible correspondence cost.
    /// </summary>
    public double MaxCorrespondenceCost => _maxCorrespondenceCost;

    /// <summary>
    /// Gets the grid type.
    /// </summary>
    public abstract GridType GetGridType();

    /// <summary>
    /// Returns the correspondence cost of the cell with 'cell_index'.
    /// </summary>
    public double GetCorrespondenceCost(Array2i cellIndex)
    {
        if (!_limits.Contains(cellIndex))
        {
            return _maxCorrespondenceCost;
        }
        var flatIndex = ToFlatIndex(cellIndex);
        if (flatIndex >= _correspondenceCostCells.Count)
        {
            return _maxCorrespondenceCost;
        }
        var value = _correspondenceCostCells[flatIndex];
        if (value >= _valueToCorrespondenceCostTable.Length)
        {
            return _maxCorrespondenceCost;
        }
        return _valueToCorrespondenceCostTable[value];
    }

    /// <summary>
    /// Copies all correspondence cost values into the destination array using bulk access.
    /// Much faster than calling GetCorrespondenceCost per cell (avoids per-cell bounds checks,
    /// struct allocation, and virtual dispatch). Data is in row-major order matching the internal layout.
    /// </summary>
    /// <param name="destination">Pre-allocated array of size >= NumXCells * NumYCells.</param>
    public void CopyCorrespondenceCostData(double[] destination)
    {
        var numXCells = _limits.CellLimits.NumXCells;
        var numYCells = _limits.CellLimits.NumYCells;
        var totalCells = numXCells * numYCells;

        if (destination.Length < totalCells)
            throw new ArgumentException($"Destination array too small: {destination.Length} < {totalCells}", nameof(destination));

        var cells = CollectionsMarshal.AsSpan(_correspondenceCostCells);
        var table = _valueToCorrespondenceCostTable;
        var maxCost = _maxCorrespondenceCost;
        var tableLen = table.Length;
        var count = Math.Min(totalCells, cells.Length);

        for (int i = 0; i < count; i++)
        {
            var value = cells[i];
            destination[i] = value < tableLen ? table[value] : maxCost;
        }

        // Fill remaining if cells is shorter than expected (shouldn't happen normally)
        for (int i = count; i < totalCells; i++)
        {
            destination[i] = maxCost;
        }
    }

    /// <summary>
    /// Returns true if the probability at the specified index is known.
    /// </summary>
    public bool IsKnown(Array2i cellIndex)
    {
        if (!_limits.Contains(cellIndex))
        {
            return false;
        }
        
        var flatIndex = ToFlatIndex(cellIndex);
        var cellValue = _correspondenceCostCells[flatIndex];
        const ushort kUpdateMarker = (ushort)(1u << 15);
        var hasUpdateMarker = cellValue >= kUpdateMarker;
        var valueWithoutMarker = hasUpdateMarker ? (ushort)(cellValue - kUpdateMarker) : cellValue;
        var isKnown = valueWithoutMarker != kUnknownCorrespondenceValue;
        
        return isKnown;
    }

    /// <summary>
    /// Finishes the update sequence.
    /// Match C++: DCHECK_GE(correspondence_cost_cells_[update_indices_.back()], kUpdateMarker)
    /// Optimized: uses Span for batch processing instead of per-element RemoveAt.
    /// </summary>
    public void FinishUpdate()
    {
        const ushort kUpdateMarker = (ushort)(1u << 15);
        var indices = CollectionsMarshal.AsSpan(_updateIndices);
        var cells = CollectionsMarshal.AsSpan(_correspondenceCostCells);
        for (int i = indices.Length - 1; i >= 0; i--)
        {
            var idx = indices[i];
#if DEBUG
            if (cells[idx] < kUpdateMarker)
            {
                throw new InvalidOperationException($"Cell at index {idx} does not have update marker. Value: {cells[idx]}, Expected: >= {kUpdateMarker}");
            }
#endif
            cells[idx] -= kUpdateMarker;
        }
        _updateIndices.Clear();
    }

    /// <summary>
    /// Fills in 'offset' and 'limits' to define a subregion that contains all known cells.
    /// Match C++: if (known_cells_box_.isEmpty()) - check by min > max
    /// </summary>
    public void ComputeCroppedLimits(out Array2i offset, out CellLimits limits)
    {
        // Match C++: AlignedBox2i::isEmpty() returns true if min > max
        // Empty box is marked by minX > maxX (or minY > maxY)
        if (_knownCellsBox.minX > _knownCellsBox.maxX || _knownCellsBox.minY > _knownCellsBox.maxY)
        {
            offset = Array2i.Zero;
            limits = new CellLimits(1, 1);
            return;
        }

        offset = new Array2i(_knownCellsBox.minX, _knownCellsBox.minY);
        limits = new CellLimits(
            _knownCellsBox.maxX - _knownCellsBox.minX + 1,
            _knownCellsBox.maxY - _knownCellsBox.minY + 1);
    }

    /// <summary>
    /// Grows the map as necessary to include 'point'. This changes the meaning of
    /// these coordinates going forward. This method must be called immediately
    /// after 'FinishUpdate', before any calls to 'ApplyLookupTable'.
    /// </summary>
    public virtual void GrowLimits(Vector2 point)
    {
        GrowLimits(point, [_correspondenceCostCells], [kUnknownCorrespondenceValue]);
    }

    /// <summary>
    /// Grows limits for multiple grids.
    /// </summary>
    protected void GrowLimits(Vector2 point, List<ushort>[] grids, ushort[] gridsUnknownCellValues)
    {
        if (_updateIndices.Count > 0)
        {
            throw new InvalidOperationException("GrowLimits must be called after FinishUpdate");
        }

        // Log before resize if point is outside current limits
        var needsGrow = !_limits.Contains(_limits.GetCellIndex(point));
        
        while (!_limits.Contains(_limits.GetCellIndex(point)))
        {
            var xOffset = _limits.CellLimits.NumXCells / 2;
            var yOffset = _limits.CellLimits.NumYCells / 2;
            
            // CRITICAL FIX: Check for integer overflow before multiplying by 2
            // Use long to prevent overflow during calculation
            long newNumXCellsLong = 2L * _limits.CellLimits.NumXCells;
            long newNumYCellsLong = 2L * _limits.CellLimits.NumYCells;
            
            // Check if result exceeds int.MaxValue
            if (newNumXCellsLong > int.MaxValue || newNumYCellsLong > int.MaxValue)
            {
                throw new InvalidOperationException(
                    $"Cannot grow grid: new size would exceed int.MaxValue. " +
                    $"Current: NumXCells={_limits.CellLimits.NumXCells}, NumYCells={_limits.CellLimits.NumYCells}. " +
                    $"New would be: NumXCells={newNumXCellsLong}, NumYCells={newNumYCellsLong}");
            }
            
            // Match C++: limits_.max() + limits_.resolution() * Eigen::Vector2d(y_offset, x_offset)
            // C++ uses Eigen::Vector2d(y_offset, x_offset), which means:
            // newMax.x() = oldMax.x() + resolution * y_offset
            // newMax.y() = oldMax.y() + resolution * x_offset
            var newMax = new Vector2(
                (_limits.Max.X + _limits.Resolution * yOffset),
                (_limits.Max.Y + _limits.Resolution * xOffset));
            var newCellLimits = new CellLimits(
                (int)newNumXCellsLong,
                (int)newNumYCellsLong);
            var newLimits = new MapLimits(_limits.Resolution, newMax, newCellLimits);

            var stride = newLimits.CellLimits.NumXCells;
            var offset = xOffset + stride * yOffset;
            
            // CRITICAL FIX: Use long for newSize calculation to prevent overflow
            long newSizeLong = (long)newLimits.CellLimits.NumXCells * newLimits.CellLimits.NumYCells;
            if (newSizeLong <= 0 || newSizeLong > int.MaxValue)
            {
                throw new InvalidOperationException(
                    $"New size is invalid: {newSizeLong}. " +
                    $"NumXCells={newLimits.CellLimits.NumXCells}, NumYCells={newLimits.CellLimits.NumYCells}");
            }
            var newSize = (int)newSizeLong;

            for (int gridIndex = 0; gridIndex < grids.Length; gridIndex++)
            {
                var oldGrid = grids[gridIndex];
                if (oldGrid == null || oldGrid.Count == 0)
                {
                    throw new InvalidOperationException($"Grid at index {gridIndex} is null or empty");
                }

                // Allocate and initialize new cells without per-element Add() loop.
                // CollectionsMarshal.SetCount sets the internal _size field directly,
                // avoiding repeated bounds checks; the backing array is zero-initialised
                // by the runtime, so an explicit fill is only needed for non-zero values.
                var newCells = new List<ushort>(newSize);
                CollectionsMarshal.SetCount(newCells, newSize);
                var newSpan = CollectionsMarshal.AsSpan(newCells);
                var unknownValue = gridsUnknownCellValues[gridIndex];
                if (unknownValue != 0)
                    newSpan.Fill(unknownValue);

                // Copy old rows into the correct offset in the new (larger) grid.
                // Using Span.CopyTo() per row lets the JIT emit a single memcpy/memmove
                // for each row instead of indexing element-by-element.
                var oldSpan = CollectionsMarshal.AsSpan(oldGrid);
                int numXCells = _limits.CellLimits.NumXCells;
                int numYCells = _limits.CellLimits.NumYCells;
                for (int i = 0; i < numYCells; i++)
                {
                    var srcRow = oldSpan.Slice(i * numXCells, numXCells);
                    var dstRow = newSpan.Slice(offset + i * stride, numXCells);
                    srcRow.CopyTo(dstRow);
                }

                // THREAD-SAFETY FIX: Instead of Clear()+AddRange() which creates a
                // window where Count==0 (race with ConstraintBuilder2D ThreadPool readers),
                // assign the completed newCells to the array slot. The field references
                // are updated atomically via UpdateGridReferences below.
                grids[gridIndex] = newCells;
            }

            // Atomic field swap: update _correspondenceCostCells (and _weightCells for TSDF2D)
            // BEFORE updating _limits, so concurrent readers with old limits compute small
            // indices into the (larger) new cells list — always safe.
            UpdateGridReferences(grids);
            _limits = newLimits;

            // Match C++: if (!known_cells_box_.isEmpty()) { known_cells_box_.translate(...); }
            // Update known cells box offset only if box is not empty
            if (_knownCellsBox.minX <= _knownCellsBox.maxX && _knownCellsBox.minY <= _knownCellsBox.maxY)
            {
                _knownCellsBox = (
                    _knownCellsBox.minX + xOffset,
                    _knownCellsBox.minY + yOffset,
                    _knownCellsBox.maxX + xOffset,
                    _knownCellsBox.maxY + yOffset
                );
            }
        }
    }

    /// <summary>
    /// Updates field references after GrowLimits rebuilds the cells lists.
    /// Called with the new lists before _limits is updated, ensuring concurrent
    /// readers always see a valid (complete) cells list.
    /// Override in derived classes that hold additional grid lists (e.g., TSDF2D._weightCells).
    /// </summary>
    protected virtual void UpdateGridReferences(List<ushort>[] grids)
    {
        _correspondenceCostCells = grids[0];
    }

    /// <summary>
    /// Converts a 'cell_index' into an index into 'cells_'.
    /// </summary>
    protected int ToFlatIndex(Array2i cellIndex)
    {
        if (!_limits.Contains(cellIndex))
        {
            throw new ArgumentOutOfRangeException(nameof(cellIndex), "Cell index out of bounds");
        }
        return _limits.CellLimits.NumXCells * cellIndex.Y + cellIndex.X;
    }

    /// <summary>
    /// Lightweight, read-only snapshot of grid cell data used for thread-safe
    /// occupancy grid generation. Only the data needed for merging is captured;
    /// the underlying Grid2D remains unaffected.
    /// </summary>
    public sealed class GridCellSnapshot
    {
        public readonly ushort[] Cells;
        public readonly MapLimits Limits;
        public readonly (int minX, int minY, int maxX, int maxY) KnownCellsBox;
        public readonly double[] ValueToCorrespondenceCostTable;
        public readonly double MinCorrespondenceCost;
        public readonly double MaxCorrespondenceCost;

        internal GridCellSnapshot(
            ushort[] cells,
            MapLimits limits,
            (int minX, int minY, int maxX, int maxY) knownCellsBox,
            double[] valueToCorrespondenceCostTable,
            double minCorrespondenceCost,
            double maxCorrespondenceCost)
        {
            Cells = cells;
            Limits = limits;
            KnownCellsBox = knownCellsBox;
            ValueToCorrespondenceCostTable = valueToCorrespondenceCostTable;
            MinCorrespondenceCost = minCorrespondenceCost;
            MaxCorrespondenceCost = maxCorrespondenceCost;
        }

        public bool IsKnown(Array2i cellIndex)
        {
            if (!Limits.Contains(cellIndex)) return false;
            var flatIndex = Limits.CellLimits.NumXCells * cellIndex.Y + cellIndex.X;
            if (flatIndex < 0 || flatIndex >= Cells.Length) return false;
            const ushort kUpdateMarker = (ushort)(1u << 15);
            var value = Cells[flatIndex];
            var raw = value >= kUpdateMarker ? (ushort)(value - kUpdateMarker) : value;
            return raw != 0;
        }

        public void ComputeCroppedLimits(out Array2i offset, out CellLimits limits)
        {
            if (KnownCellsBox.minX > KnownCellsBox.maxX || KnownCellsBox.minY > KnownCellsBox.maxY)
            {
                offset = Array2i.Zero;
                limits = new CellLimits(1, 1);
                return;
            }
            offset = new Array2i(KnownCellsBox.minX, KnownCellsBox.minY);
            limits = new CellLimits(
                KnownCellsBox.maxX - KnownCellsBox.minX + 1,
                KnownCellsBox.maxY - KnownCellsBox.minY + 1);
        }
    }

    /// <summary>
    /// Creates a snapshot of the current cell data for thread-safe reading.
    /// The returned snapshot is decoupled from the live grid and will not be
    /// affected by concurrent InsertRangeData / GrowLimits calls.
    /// Intended for occupancy grid generation on active (non-finished) submaps.
    /// </summary>
    public GridCellSnapshot SnapshotCellData()
    {
        // Capture references / value copies before allocating the array.
        var limits = _limits;
        var knownCellsBox = _knownCellsBox;
        var cells = CollectionsMarshal.AsSpan(_correspondenceCostCells);
        var copy = new ushort[cells.Length];
        cells.CopyTo(copy);
        return new GridCellSnapshot(
            copy, limits, knownCellsBox,
            _valueToCorrespondenceCostTable,
            _minCorrespondenceCost, _maxCorrespondenceCost);
    }

    /// <summary>
    /// Computes a cropped grid containing only known cells.
    /// </summary>
    public abstract Grid2D ComputeCroppedGrid();

    /// <summary>
    /// Converts to proto representation.
    /// Match C++: CHECK(update_indices().empty()) before serializing
    /// </summary>
    public virtual Models.Mapping.Grid2D ToProto()
    {
        // Match C++: CHECK(update_indices().empty()) << "Serializing a grid during an update is not supported. Finish the update first.";
        if (_updateIndices.Count > 0)
        {
            throw new InvalidOperationException("Serializing a grid during an update is not supported. Finish the update first.");
        }

        var cells = new List<int>(_correspondenceCostCells.Count);
        foreach (var cell in _correspondenceCostCells)
        {
            cells.Add(cell);
        }

        var proto = new Models.Mapping.Grid2D
        {
            Limits = _limits.ToProto(),
            Cells = cells,
            MinCorrespondenceCost = _minCorrespondenceCost,
            MaxCorrespondenceCost = _maxCorrespondenceCost
        };

        // Match C++: if (!known_cells_box().isEmpty()) { set known_cells_box }
        if (_knownCellsBox.minX <= _knownCellsBox.maxX && _knownCellsBox.minY <= _knownCellsBox.maxY)
        {
            proto.KnownCellsBox = new Models.Mapping.Grid2D.CellBox(
                _knownCellsBox.maxX, _knownCellsBox.maxY,
                _knownCellsBox.minX, _knownCellsBox.minY);
        }

        return proto;
    }
}