using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Running;
using LanguageExt;
using PersistentMap;

namespace MapBenchmarks;

[MemoryDiagnoser]
public class IntMapBenchmarks
{
    [Params(100, 1000, 10000, 100000)]
    public int N { get; set; }

    private int[] _allKeys;
    private int[] _retrieveKeys;
    private int[] _updateKeys;
    private int[] _removeKeys;
    private int[] _mixedKeys; // Half existing, half new

    // Pre-built collections for read/update/remove tests
    private ImmutableDictionary<int, int> _immDict;
    private ImmutableSortedDictionary<int, int> _immSortedDict;
    private LanguageExt.Map<int, int> _extMap;
    private LanguageExt.HashMap<int, int> _extHashMap;
    private PersistentMap<int, int, IntStrategy> _persistentMap;

    private readonly IntStrategy _intStrategy = new IntStrategy();

    [GlobalSetup]
    public void Setup()
    {
        var rnd = new Random(42);
        
        // Build integer keys (inserted sorted)
        _allKeys = Enumerable.Range(0, N).ToArray();

        int subsetSize = Math.Max(1, N / 10);
        
        // Subsets for different operations
        var shuffled = _allKeys.OrderBy(x => rnd.Next()).ToArray();
        _retrieveKeys = shuffled.Take(subsetSize).ToArray();
        _updateKeys = shuffled.Skip(subsetSize).Take(subsetSize).ToArray();
        _removeKeys = shuffled.Skip(subsetSize * 2).Take(subsetSize).ToArray();

        // Mixed keys: half existing, half completely new (N + 1 to N + subsetSize/2)
        var existingHalf = shuffled.Skip(subsetSize * 3).Take(subsetSize / 2).ToArray();
        var newHalf = Enumerable.Range(N + 1, subsetSize - (subsetSize / 2)).ToArray();
        _mixedKeys = existingHalf.Concat(newHalf).OrderBy(x => rnd.Next()).ToArray();

        // Pre-build collections
        _immDict = ImmutableDictionary.CreateRange(_allKeys.Select(k => new KeyValuePair<int, int>(k, k)));
        _immSortedDict = ImmutableSortedDictionary.CreateRange(_allKeys.Select(k => new KeyValuePair<int, int>(k, k)));
        
        _extMap = LanguageExt.Map.empty<int, int>();
        _extHashMap = LanguageExt.HashMap.empty<int, int>();
        foreach (var k in _allKeys)
        {
            _extMap = _extMap.AddOrUpdate(k, k);
            _extHashMap = _extHashMap.AddOrUpdate(k, k);
        }

        var transient = BaseOrderedMap<int, int, IntStrategy>.CreateTransient(_intStrategy);
        foreach (var k in _allKeys) transient.Set(k, k);
        _persistentMap = transient.ToPersistent();
    }

    // --- 1. BUILD ---

    [Benchmark]
    public ImmutableDictionary<int, int> Build_ImmDict()
    {
        var map = ImmutableDictionary<int, int>.Empty;
        foreach (var k in _allKeys) map = map.Add(k, k);
        return map;
    }

    [Benchmark]
    public ImmutableSortedDictionary<int, int> Build_ImmSortedDict()
    {
        var map = ImmutableSortedDictionary<int, int>.Empty;
        foreach (var k in _allKeys) map = map.Add(k, k);
        return map;
    }

    [Benchmark]
    public LanguageExt.Map<int, int> Build_ExtMap()
    {
        var map = LanguageExt.Map.empty<int, int>();
        foreach (var k in _allKeys) map = map.AddOrUpdate(k, k);
        return map;
    }

    [Benchmark]
    public LanguageExt.HashMap<int, int> Build_ExtHashMap()
    {
        var map = LanguageExt.HashMap.empty<int, int>();
        foreach (var k in _allKeys) map = map.AddOrUpdate(k, k);
        return map;
    }

    [Benchmark]
    public PersistentMap<int, int, IntStrategy> Build_PersistentMap()
    {
        var map = PersistentMap<int, int, IntStrategy>.Empty(_intStrategy);
        foreach (var k in _allKeys) map = map.Set(k, k);
        return map;
    }

    [Benchmark]
    public PersistentMap<int, int, IntStrategy> Build_TransientMap()
    {
        var map = BaseOrderedMap<int, int, IntStrategy>.CreateTransient(_intStrategy);
        foreach (var k in _allKeys) map.Set(k, k);
        return map.ToPersistent();
    }

    // --- 2. RETRIEVAL ---

    [Benchmark]
    public int Retrieve_ImmDict()
    {
        int count = 0;
        foreach (var k in _retrieveKeys)
            if (_immDict.TryGetValue(k, out _)) count++;
        return count;
    }

    [Benchmark]
    public int Retrieve_ImmSortedDict()
    {
        int count = 0;
        foreach (var k in _retrieveKeys)
            if (_immSortedDict.TryGetValue(k, out _)) count++;
        return count;
    }

    [Benchmark]
    public int Retrieve_ExtMap()
    {
        int count = 0;
        foreach (var k in _retrieveKeys)
            if (_extMap.Find(k).IsSome) count++;
        return count;
    }

    [Benchmark]
    public int Retrieve_ExtHashMap()
    {
        int count = 0;
        foreach (var k in _retrieveKeys)
            if (_extHashMap.Find(k).IsSome) count++;
        return count;
    }

    [Benchmark]
    public int Retrieve_PersistentMap()
    {
        int count = 0;
        foreach (var k in _retrieveKeys)
            if (_persistentMap.TryGetValue(k, out _)) count++;
        return count;
    }

    // --- 3. UPDATING ---

    [Benchmark]
    public ImmutableDictionary<int, int> Update_ImmDict()
    {
        var map = _immDict;
        foreach (var k in _updateKeys) map = map.SetItem(k, 999);
        return map;
    }

[Benchmark]
    public PersistentMap<int, int, IntStrategy> Update_PersistentMap()
    {
        var map = _persistentMap;
        foreach (var k in _updateKeys) map = map.Set(k, 999);
        return map;
    }

    [Benchmark]
    public PersistentMap<int, int, IntStrategy> Update_TransientMap()
    {
        var transient = _persistentMap.ToTransient();
        foreach (var k in _updateKeys) transient.Set(k, 999);
        return transient.ToPersistent();
    }
    
    [Benchmark]
    public ImmutableSortedDictionary<int, int> Update_ImmSortedDict()
    {
        var map = _immSortedDict;
        foreach (var k in _updateKeys) map = map.SetItem(k, 999);
        return map;
    }

    [Benchmark]
    public LanguageExt.Map<int, int> Update_ExtMap()
    {
        var map = _extMap;
        foreach (var k in _updateKeys) map = map.SetItem(k, 999);
        return map;
    }

    [Benchmark]
    public LanguageExt.HashMap<int, int> Update_ExtHashMap()
    {
        var map = _extHashMap;
        foreach (var k in _updateKeys) map = map.SetItem(k, 999);
        return map;
    }

    // --- 4. UPDATE & SET (MIXED) ---

    [Benchmark]
    public ImmutableDictionary<int, int> UpdateSet_ImmDict()
    {
        var map = _immDict;
        foreach (var k in _mixedKeys) map = map.SetItem(k, 999);
        return map;
    }

[Benchmark]
    public PersistentMap<int, int, IntStrategy> UpdateSet_PersistentMap()
    {
        var map = _persistentMap;
        foreach (var k in _mixedKeys) map = map.Set(k, 999);
        return map;
    }

    [Benchmark]
    public PersistentMap<int, int, IntStrategy> UpdateSet_TransientMap()
    {
        var transient = _persistentMap.ToTransient();
        foreach (var k in _mixedKeys) transient.Set(k, 999);
        return transient.ToPersistent();
    }
    
    [Benchmark]
    public ImmutableSortedDictionary<int, int> UpdateSet_ImmSortedDict()
    {
        var map = _immSortedDict;
        foreach (var k in _mixedKeys) map = map.SetItem(k, 999);
        return map;
    }

    [Benchmark]
    public LanguageExt.Map<int, int> UpdateSet_ExtMap()
    {
        var map = _extMap;
        foreach (var k in _mixedKeys) map = map.AddOrUpdate(k, 999);
        return map;
    }

    [Benchmark]
    public LanguageExt.HashMap<int, int> UpdateSet_ExtHashMap()
    {
        var map = _extHashMap;
        foreach (var k in _mixedKeys) map = map.AddOrUpdate(k, 999);
        return map;
    }

    // --- 5. ITERATION ---

    [Benchmark]
    public int Iterate_ImmDict()
    {
        int sum = 0;
        foreach (var kvp in _immDict) sum += kvp.Value;
        return sum;
    }

    [Benchmark]
    public int Iterate_PersistentMap()
    {
        int sum = 0;
        foreach (var kvp in _persistentMap) sum += kvp.Value;
        return sum;
    }
    
    [Benchmark]
    public int Iterate_ImmSortedDict()
    {
        int sum = 0;
        foreach (var kvp in _immSortedDict) sum += kvp.Value;
        return sum;
    }

    [Benchmark]
    public int Iterate_ExtMap()
    {
        int sum = 0;
        foreach (var kvp in _extMap) sum += kvp.Value;
        return sum;
    }

    [Benchmark]
    public int Iterate_ExtHashMap()
    {
        int sum = 0;
        foreach (var kvp in _extHashMap) sum += kvp.Value;
        return sum;
    }

    
    // --- 6. REMOVAL ---

    [Benchmark]
    public ImmutableDictionary<int, int> Remove_ImmDict()
    {
        var map = _immDict;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }

[Benchmark]
    public PersistentMap<int, int, IntStrategy> Remove_PersistentMap()
    {
        var map = _persistentMap;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }
    [Benchmark]
    public PersistentMap<int, int, IntStrategy> Remove_TransientMap()
    {
        var transient = _persistentMap.ToTransient();
        foreach (var k in _removeKeys) transient.Remove(k);
        return transient.ToPersistent();
    }
    
    [Benchmark]
    public ImmutableSortedDictionary<int, int> Remove_ImmSortedDict()
    {
        var map = _immSortedDict;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }

    [Benchmark]
    public LanguageExt.Map<int, int> Remove_ExtMap()
    {
        var map = _extMap;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }

    [Benchmark]
    public LanguageExt.HashMap<int, int> Remove_ExtHashMap()
    {
        var map = _extHashMap;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }
    
    
    
    public static void Main(string[] args)
    {
        BenchmarkSwitcher
            .FromAssembly(typeof(IntMapBenchmarks).Assembly)
            .Run(args);
    }
}