/*
 * Copyright 2016 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 System.Collections;

namespace CartographerSharp.Mapping;

/// <summary>
/// Reminiscent of Dictionary, but indexed by 'IdType' which can be 'NodeId' or 'SubmapId'.
/// Note: This container will only ever contain non-empty trajectories. Trimming
/// the last remaining node of a trajectory drops the trajectory.
/// </summary>
public class MapById<TId, TData> : IEnumerable<MapById<TId, TData>.IdDataReference>
    where TId : struct, IIdType
{
    private class MapByIndex
    {
        public bool CanAppend { get; set; } = true;
        public SortedDictionary<int, TData> Data { get; } = [];
    }

    /// <summary>
    /// Reference to an ID and its associated data.
    /// </summary>
    public readonly struct IdDataReference(TId id, TData data)
    {
        public TId Id { get; } = id;
        public TData Data { get; } = data;
    }

    private readonly SortedDictionary<int, MapByIndex> _trajectories = [];

    /// <summary>
    /// Appends data to a 'trajectory_id', creating trajectories as needed.
    /// </summary>
    public TId Append(int trajectoryId, TData data)
    {
        if (trajectoryId < 0)
        {
            throw new ArgumentException("Trajectory ID must be non-negative", nameof(trajectoryId));
        }

        if (!_trajectories.TryGetValue(trajectoryId, out var trajectory))
        {
            trajectory = new MapByIndex();
            _trajectories[trajectoryId] = trajectory;
        }

        if (!trajectory.CanAppend)
        {
            throw new InvalidOperationException($"Cannot append to trajectory {trajectoryId} that has been modified. Trajectory has {trajectory.Data.Count} existing entries.");
        }

        var index = trajectory.Data.Count == 0 
            ? 0 
            : trajectory.Data.Keys.Max() + 1;
        
        trajectory.Data[index] = data;
        
        // Create ID instance
        if (typeof(TId) == typeof(NodeId))
        {
            return (TId)(object)new NodeId(trajectoryId, index);
        }
        else if (typeof(TId) == typeof(SubmapId))
        {
            return (TId)(object)new SubmapId(trajectoryId, index);
        }
        throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
    }

    /// <summary>
    /// Inserts data (which must not exist already) into a trajectory.
    /// </summary>
    public void Insert(TId id, TData data)
    {
        var trajectoryId = id.GetTrajectoryId();
        var index = id.GetIndex();

        if (trajectoryId < 0 || index < 0)
        {
            throw new ArgumentException("ID components must be non-negative", nameof(id));
        }

        if (!_trajectories.TryGetValue(trajectoryId, out var trajectory))
        {
            trajectory = new MapByIndex();
            _trajectories[trajectoryId] = trajectory;
        }

        trajectory.CanAppend = false;
        
        if (trajectory.Data.ContainsKey(index))
        {
            throw new ArgumentException($"Data already exists for ID {id}", nameof(id));
        }

        trajectory.Data[index] = data;
    }

    /// <summary>
    /// Removes the data for 'id' which must exist.
    /// </summary>
    public void Trim(TId id)
    {
        var trajectoryId = id.GetTrajectoryId();
        var index = id.GetIndex();

        if (!_trajectories.TryGetValue(trajectoryId, out var trajectory))
        {
            throw new KeyNotFoundException($"Trajectory {trajectoryId} not found");
        }
        
        if (!trajectory.Data.TryGetValue(index, out _))
        {
            throw new KeyNotFoundException($"ID {id} not found");
        }

        // Check if this is the highest index in the trajectory
        // Since SortedDictionary maintains sorted order, we can check if index equals the max key
        var maxKey = trajectory.Data.Keys.Max();
        if (index == maxKey)
        {
            // We are removing the data with the highest index from this trajectory.
            // We assume that we will never append to it anymore.
            trajectory.CanAppend = false;
        }

        trajectory.Data.Remove(index);
        
        if (trajectory.Data.Count == 0)
        {
            _trajectories.Remove(trajectoryId);
        }
    }

    /// <summary>
    /// Checks if the container contains the given ID.
    /// </summary>
    public bool Contains(TId id)
    {
        var trajectoryId = id.GetTrajectoryId();
        var index = id.GetIndex();
        
        return _trajectories.TryGetValue(trajectoryId, out var trajectory) &&
               trajectory.Data.ContainsKey(index);
    }

    /// <summary>
    /// Gets the data for the given ID.
    /// </summary>
    public TData this[TId id]
    {
        get
        {
            var trajectoryId = id.GetTrajectoryId();
            var index = id.GetIndex();
            
            if (!_trajectories.TryGetValue(trajectoryId, out var trajectory))
            {
                throw new KeyNotFoundException($"Trajectory {trajectoryId} not found");
            }
            
            if (!trajectory.Data.TryGetValue(index, out var data))
            {
                throw new KeyNotFoundException($"ID {id} not found");
            }
            
            return data;
        }
        set
        {
            var trajectoryId = id.GetTrajectoryId();
            var index = id.GetIndex();
            
            if (!_trajectories.TryGetValue(trajectoryId, out var trajectory))
            {
                trajectory = new MapByIndex();
                _trajectories[trajectoryId] = trajectory;
            }
            
            trajectory.Data[index] = value;
        }
    }

    /// <summary>
    /// Returns an iterator to the first element in the trajectory.
    /// </summary>
    public IEnumerable<IdDataReference> BeginOfTrajectory(int trajectoryId)
    {
        if (!_trajectories.TryGetValue(trajectoryId, out var trajectory))
        {
            yield break;
        }

        // Create snapshot to avoid concurrent modification exception
        foreach (var kvp in trajectory.Data.ToList())
        {
            TId id;
            if (typeof(TId) == typeof(NodeId))
            {
                id = (TId)(object)new NodeId(trajectoryId, kvp.Key);
            }
            else if (typeof(TId) == typeof(SubmapId))
            {
                id = (TId)(object)new SubmapId(trajectoryId, kvp.Key);
            }
            else
            {
                throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
            }
            yield return new IdDataReference(id, kvp.Value);
        }
    }

    /// <summary>
    /// Returns the size of a trajectory, or 0 if it doesn't exist.
    /// </summary>
    public int SizeOfTrajectoryOrZero(int trajectoryId)
    {
        return _trajectories.TryGetValue(trajectoryId, out var trajectory)
            ? trajectory.Data.Count
            : 0;
    }

    /// <summary>
    /// Returns the first element in the trajectory whose time is not before 'time',
    /// or null if the trajectory is empty or all elements are before 'time'.
    /// Match C++ MapById::lower_bound(trajectory_id, time) using internal::GetTime(data).
    /// </summary>
    /// <param name="trajectoryId">Trajectory ID.</param>
    /// <param name="time">Time (e.g. common::Time in C++).</param>
    /// <param name="getTime">Function to get time from TData (e.g. node => node.Time).</param>
    /// <returns>First IdDataReference with getTime(data) >= time, or null if none.</returns>
    public IdDataReference? LowerBound(int trajectoryId, long time, Func<TData, long> getTime)
    {
        if (getTime == null)
        {
            throw new ArgumentNullException(nameof(getTime));
        }

        if (!_trajectories.TryGetValue(trajectoryId, out var trajectory) || trajectory.Data.Count == 0)
        {
            return null;
        }

        // SortedDictionary orders by key (index); get ordered indices
        var indices = trajectory.Data.Keys.ToList();

        // If last element's time is before 'time', return null (past end)
        var lastIndex = indices[^1];
        var lastData = trajectory.Data[lastIndex];
        if (getTime(lastData) < time)
        {
            return null;
        }

        // Binary search: first index i such that getTime(Data[indices[i]]) >= time
        int left = 0;
        int right = indices.Count - 1;
        while (left < right)
        {
            int mid = left + (right - left) / 2;
            var midIndex = indices[mid];
            var midData = trajectory.Data[midIndex];
            if (getTime(midData) < time)
            {
                left = mid + 1;
            }
            else
            {
                right = mid;
            }
        }

        var index = indices[left];
        var data = trajectory.Data[index];
        TId id;
        if (typeof(TId) == typeof(NodeId))
        {
            id = (TId)(object)new NodeId(trajectoryId, index);
        }
        else if (typeof(TId) == typeof(SubmapId))
        {
            id = (TId)(object)new SubmapId(trajectoryId, index);
        }
        else
        {
            throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
        }
        return new IdDataReference(id, data);
    }

    /// <summary>
    /// Gets the last element in a trajectory (matching C++ std::prev(end_it)).
    /// Returns null if trajectory doesn't exist or is empty.
    /// CRITICAL FIX: Use Last() on SortedDictionary.Keys instead of Max() for better performance.
    /// </summary>
    public IdDataReference? GetLastOfTrajectory(int trajectoryId)
    {
        if (!_trajectories.TryGetValue(trajectoryId, out var trajectory) || trajectory.Data.Count == 0)
        {
            return null;
        }

        // CRITICAL FIX: SortedDictionary maintains sorted order, so use Last() instead of Max()
        // Max() iterates through all keys, while Last() is O(1) for SortedDictionary
        // However, SortedDictionary.Keys doesn't have Last() directly, so we need to use Max()
        // But we can optimize by checking if Count > 0 first (already done above)
        var lastKey = trajectory.Data.Keys.Max(); // This is O(n) but necessary for SortedDictionary
        var lastData = trajectory.Data[lastKey];
        
        TId id;
        if (typeof(TId) == typeof(NodeId))
        {
            id = (TId)(object)new NodeId(trajectoryId, lastKey);
        }
        else if (typeof(TId) == typeof(SubmapId))
        {
            id = (TId)(object)new SubmapId(trajectoryId, lastKey);
        }
        else
        {
            throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
        }
        
        return new IdDataReference(id, lastData);
    }

    /// <summary>
    /// Returns the total count of all elements.
    /// </summary>
    public int Count
    {
        get
        {
            return _trajectories.Values.Sum(t => t.Data.Count);
        }
    }

    /// <summary>
    /// Checks if the container is empty.
    /// </summary>
    public bool IsEmpty => _trajectories.Count == 0;

    /// <summary>
    /// Gets all trajectory IDs.
    /// </summary>
    public IEnumerable<int> TrajectoryIds => _trajectories.Keys;

    public IEnumerator<IdDataReference> GetEnumerator()
    {
        // Create snapshot to avoid concurrent modification exception
        // This prevents "Collection was modified after the enumerator was instantiated" error
        // when the collection is modified during iteration (e.g., from another thread)
        List<IdDataReference> snapshot = new List<IdDataReference>();

        foreach (var trajectoryKvp in _trajectories.ToList())
        {
            var trajectoryId = trajectoryKvp.Key;
            foreach (var dataKvp in trajectoryKvp.Value.Data.ToList())
            {
                TId id;
                if (typeof(TId) == typeof(NodeId))
                {
                    id = (TId)(object)new NodeId(trajectoryId, dataKvp.Key);
                }
                else if (typeof(TId) == typeof(SubmapId))
                {
                    id = (TId)(object)new SubmapId(trajectoryId, dataKvp.Key);
                }
                else
                {
                    throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
                }
                snapshot.Add(new IdDataReference(id, dataKvp.Value));
            }
        }

        return snapshot.GetEnumerator();
    }

    /// <summary>
    /// Enumerates all elements without creating a snapshot.
    /// ONLY safe to use when the caller holds the lock that protects this collection
    /// from concurrent modification (e.g., inside lock(_dataLock)).
    /// Avoids O(N) allocation that GetEnumerator() incurs.
    /// </summary>
    public IEnumerable<IdDataReference> EnumerateUnsafe()
    {
        foreach (var trajectoryKvp in _trajectories)
        {
            var trajectoryId = trajectoryKvp.Key;
            foreach (var dataKvp in trajectoryKvp.Value.Data)
            {
                TId id;
                if (typeof(TId) == typeof(NodeId))
                {
                    id = (TId)(object)new NodeId(trajectoryId, dataKvp.Key);
                }
                else if (typeof(TId) == typeof(SubmapId))
                {
                    id = (TId)(object)new SubmapId(trajectoryId, dataKvp.Key);
                }
                else
                {
                    throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
                }
                yield return new IdDataReference(id, dataKvp.Value);
            }
        }
    }

    IEnumerator IEnumerable.GetEnumerator()
    {
        return GetEnumerator();
    }

    #region Copy-on-Write Support Methods

    /// <summary>
    /// Creates a shallow clone of this MapById.
    /// The new container has its own data structures but references the same TData objects.
    /// For value types (structs), this effectively creates independent copies.
    /// Used for Copy-on-Write optimization pattern.
    /// </summary>
    /// <returns>A new MapById with the same data.</returns>
    public MapById<TId, TData> ShallowClone()
    {
        var clone = new MapById<TId, TData>();
        foreach (var trajectoryKvp in _trajectories)
        {
            var trajectoryId = trajectoryKvp.Key;
            foreach (var dataKvp in trajectoryKvp.Value.Data)
            {
                TId id;
                if (typeof(TId) == typeof(NodeId))
                {
                    id = (TId)(object)new NodeId(trajectoryId, dataKvp.Key);
                }
                else if (typeof(TId) == typeof(SubmapId))
                {
                    id = (TId)(object)new SubmapId(trajectoryId, dataKvp.Key);
                }
                else
                {
                    throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
                }
                clone.Insert(id, dataKvp.Value);
            }
        }
        return clone;
    }

    /// <summary>
    /// Clears all data from this MapById.
    /// </summary>
    public void Clear()
    {
        _trajectories.Clear();
    }

    /// <summary>
    /// Replaces all data in this MapById with data from another MapById.
    /// This is an atomic operation for Copy-on-Write pattern - clears existing data
    /// and copies all entries from the source.
    /// </summary>
    /// <param name="source">The source MapById to copy from.</param>
    public void ReplaceWith(MapById<TId, TData> source)
    {
        _trajectories.Clear();
        foreach (var trajectoryKvp in source._trajectories)
        {
            var trajectoryId = trajectoryKvp.Key;
            foreach (var dataKvp in trajectoryKvp.Value.Data)
            {
                TId id;
                if (typeof(TId) == typeof(NodeId))
                {
                    id = (TId)(object)new NodeId(trajectoryId, dataKvp.Key);
                }
                else if (typeof(TId) == typeof(SubmapId))
                {
                    id = (TId)(object)new SubmapId(trajectoryId, dataKvp.Key);
                }
                else
                {
                    throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
                }
                Insert(id, dataKvp.Value);
            }
        }
    }

    /// <summary>
    /// Gets all IDs that exist in this MapById but not in another MapById.
    /// Useful for finding nodes added during Copy-on-Write preparation phase.
    /// </summary>
    /// <param name="other">The other MapById to compare against.</param>
    /// <returns>IDs that exist in this but not in other.</returns>
    public IEnumerable<TId> GetIdsDifference(MapById<TId, TData> other)
    {
        foreach (var trajectoryKvp in _trajectories)
        {
            var trajectoryId = trajectoryKvp.Key;
            foreach (var dataKvp in trajectoryKvp.Value.Data)
            {
                TId id;
                if (typeof(TId) == typeof(NodeId))
                {
                    id = (TId)(object)new NodeId(trajectoryId, dataKvp.Key);
                }
                else if (typeof(TId) == typeof(SubmapId))
                {
                    id = (TId)(object)new SubmapId(trajectoryId, dataKvp.Key);
                }
                else
                {
                    throw new NotSupportedException($"ID type {typeof(TId)} is not supported");
                }

                if (!other.Contains(id))
                {
                    yield return id;
                }
            }
        }
    }

    /// <summary>
    /// Estimates memory usage in bytes for this MapById.
    /// Used for memory overhead analysis.
    /// </summary>
    /// <param name="dataSize">Size of each TData element in bytes.</param>
    /// <returns>Estimated memory usage in bytes.</returns>
    public long EstimateMemoryUsageBytes(int dataSize)
    {
        // Base overhead for MapById object + _trajectories SortedDictionary
        long baseOverhead = 64;

        // Per-trajectory overhead: MapByIndex object + SortedDictionary
        long perTrajectoryOverhead = 48;

        // Per-entry overhead: SortedDictionary node (~40 bytes) + key (4 bytes)
        long perEntryOverhead = 44;

        long total = baseOverhead;
        total += _trajectories.Count * perTrajectoryOverhead;
        total += Count * (perEntryOverhead + dataSize);

        return total;
    }

    #endregion
}

/// <summary>
/// Interface for ID types used with MapById.
/// </summary>
public interface IIdType
{
    int GetTrajectoryId();
    int GetIndex();
}