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

namespace MapBenchmarks;

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

    [Params(8, 50)]
    public int StringLength { get; set; }

    private string[] _allKeys;
    private string[] _retrieveKeys;
    private string[] _updateKeys;
    private string[] _removeKeys;
    private string[] _mixedKeys;

    private ImmutableDictionary<string, int> _immDict;
    private ImmutableSortedDictionary<string, int> _immSortedDict;
    private LanguageExt.Map<string, int> _extMap;
    private LanguageExt.HashMap<string, int> _extHashMap;
    
    private PersistentMap<string, int, StandardStrategy<string>> _persistentMapStandard;
    private PersistentMap<string, int, UnicodeStrategy> _persistentMapUnicode;

    private readonly StandardStrategy<string> _stdStrategy = new StandardStrategy<string>();
    private readonly UnicodeStrategy _uniStrategy = new UnicodeStrategy();

    [GlobalSetup]
    public void Setup()
    {
        var rnd = new Random(42);
        
        // Build random strings
        _allKeys = Enumerable.Range(0, N).Select(_ => GenerateRandomString(StringLength, rnd)).Distinct().ToArray();
        
        // Regenerate if Distinct() reduced array size (highly unlikely with length 8/50, but safe)
        while (_allKeys.Length < N)
        {
            _allKeys = _allKeys.Concat(new[] { GenerateRandomString(StringLength, rnd) }).Distinct().ToArray();
        }

        int subsetSize = Math.Max(1, N / 10);
        
        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();

        var existingHalf = shuffled.Skip(subsetSize * 3).Take(subsetSize / 2).ToArray();
        var newHalf = Enumerable.Range(0, subsetSize - (subsetSize / 2)).Select(_ => GenerateRandomString(StringLength, rnd)).ToArray();
        _mixedKeys = existingHalf.Concat(newHalf).OrderBy(x => rnd.Next()).ToArray();

        // Pre-build collections
        _immDict = ImmutableDictionary.CreateRange(_allKeys.Select((k, i) => new KeyValuePair<string, int>(k, i)));
        _immSortedDict = ImmutableSortedDictionary.CreateRange(_allKeys.Select((k, i) => new KeyValuePair<string, int>(k, i)));
        
        _extMap = LanguageExt.Map.empty<string, int>();
        _extHashMap = LanguageExt.HashMap.empty<string, int>();
        for (int i = 0; i < _allKeys.Length; i++)
        {
            _extMap = _extMap.AddOrUpdate(_allKeys[i], i);
            _extHashMap = _extHashMap.AddOrUpdate(_allKeys[i], i);
        }

        var transStd = BaseOrderedMap<string, int, StandardStrategy<string>>.CreateTransient(_stdStrategy);
        var transUni = BaseOrderedMap<string, int, UnicodeStrategy>.CreateTransient(_uniStrategy);
        for (int i = 0; i < _allKeys.Length; i++)
        {
            transStd.Set(_allKeys[i], i);
            transUni.Set(_allKeys[i], i);
        }
        _persistentMapStandard = transStd.ToPersistent();
        _persistentMapUnicode = transUni.ToPersistent();
    }

    private static string GenerateRandomString(int length, Random rnd)
    {
        const string chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
        return new string(Enumerable.Repeat(chars, length).Select(s => s[rnd.Next(s.Length)]).ToArray());
    }

    // --- 1. BUILD ---
    
    [Benchmark]
    public PersistentMap<string, int, StandardStrategy<string>> Build_TransientMap_Standard()
    {
        var map = BaseOrderedMap<string, int, StandardStrategy<string>>.CreateTransient(_stdStrategy);
        for (int i = 0; i < _allKeys.Length; i++) map.Set(_allKeys[i], i);
        return map.ToPersistent();
    }

    [Benchmark]
    public PersistentMap<string, int, UnicodeStrategy> Build_TransientMap_Unicode()
    {
        var map = BaseOrderedMap<string, int, UnicodeStrategy>.CreateTransient(_uniStrategy);
        for (int i = 0; i < _allKeys.Length; i++) map.Set(_allKeys[i], i);
        return map.ToPersistent();
    }

    [Benchmark]
    public ImmutableDictionary<string, int> Build_ImmDict()
    {
        var map = ImmutableDictionary<string, int>.Empty;
        for (int i = 0; i < _allKeys.Length; i++) map = map.Add(_allKeys[i], i);
        return map;
    }

    // --- 1. BUILD (Missing) ---

    [Benchmark]
    public ImmutableSortedDictionary<string, int> Build_ImmSortedDict()
    {
        var map = ImmutableSortedDictionary<string, int>.Empty;
        for (int i = 0; i < _allKeys.Length; i++) map = map.Add(_allKeys[i], i);
        return map;
    }

    [Benchmark]
    public LanguageExt.Map<string, int> Build_ExtMap()
    {
        var map = LanguageExt.Map.empty<string, int>();
        for (int i = 0; i < _allKeys.Length; i++) map = map.AddOrUpdate(_allKeys[i], i);
        return map;
    }

    [Benchmark]
    public LanguageExt.HashMap<string, int> Build_ExtHashMap()
    {
        var map = LanguageExt.HashMap.empty<string, int>();
        for (int i = 0; i < _allKeys.Length; i++) map = map.AddOrUpdate(_allKeys[i], i);
        return map;
    }

    // --- 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_PersistentMap_Standard()
    {
        int count = 0;
        foreach (var k in _retrieveKeys)
            if (_persistentMapStandard.TryGetValue(k, out _)) count++;
        return count;
    }
    
    [Benchmark]
    public int Retrieve_PersistentMap_Unicode()
    {
        int count = 0;
        foreach (var k in _retrieveKeys)
            if (_persistentMapUnicode.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;
    }
    

    // --- 3. UPDATING ---

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

    [Benchmark]
    public PersistentMap<string, int, StandardStrategy<string>> Update_PersistentMap_Standard()
    {
        var map = _persistentMapStandard;
        foreach (var k in _updateKeys) map = map.Set(k, 999);
        return map;
    }
    [Benchmark]
    public PersistentMap<string, int, UnicodeStrategy> Update_PersistentMap_Unicode()
    {
        var map = _persistentMapUnicode;
        foreach (var k in _updateKeys) map = map.Set(k, 999);
        return map;
    }

    [Benchmark]
    public PersistentMap<string, int, StandardStrategy<string>> Update_TransientMap_Standard()
    {
        var transient = _persistentMapStandard.ToTransient();
        foreach (var k in _updateKeys) transient.Set(k, 999);
        return transient.ToPersistent();
    }

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

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

    [Benchmark]
    public LanguageExt.HashMap<string, 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<string, int> UpdateSet_ImmDict()
    {
        var map = _immDict;
        foreach (var k in _mixedKeys) map = map.SetItem(k, 999);
        return map;
    }

    [Benchmark]
    public PersistentMap<string, int, StandardStrategy<string>> UpdateSet_PersistentMap_Standard()
    {
        var map = _persistentMapStandard;
        foreach (var k in _mixedKeys) map = map.Set(k, 999);
        return map;
    }

[Benchmark]
    public PersistentMap<string, int, UnicodeStrategy> UpdateSet_PersistentMap_Unicode()
    {
        var map = _persistentMapUnicode;
        foreach (var k in _mixedKeys) map = map.Set(k, 999);
        return map;
    }

    [Benchmark]
    public PersistentMap<string, int, StandardStrategy<string>> UpdateSet_TransientMap_Standard()
    {
        var transient = _persistentMapStandard.ToTransient();
        foreach (var k in _mixedKeys) transient.Set(k, 999);
        return transient.ToPersistent();
    }

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

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

    [Benchmark]
    public LanguageExt.HashMap<string, 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_Standard()
    {
        int sum = 0;
        foreach (var kvp in _persistentMapStandard) 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;
    }
    
    [Benchmark]
    public int Iterate_PersistentMap_Unicode()
    {
        int sum = 0;
        foreach (var kvp in _persistentMapUnicode) sum += kvp.Value;
        return sum;
    }

    // --- 6. REMOVAL ---

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

    [Benchmark]
    public PersistentMap<string, int, StandardStrategy<string>> Remove_PersistentMap_Standard()
    {
        var map = _persistentMapStandard;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }

    [Benchmark]
    public PersistentMap<string, int, UnicodeStrategy> Remove_PersistentMap_Unicode()
    {
        var map = _persistentMapUnicode;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }

    [Benchmark]
    public PersistentMap<string, int, StandardStrategy<string>> Remove_TransientMap_Standard()
    {
        var transient = _persistentMapStandard.ToTransient();
        foreach (var k in _removeKeys) transient.Remove(k);
        return transient.ToPersistent();
    }

[Benchmark]
    public PersistentMap<string, int, UnicodeStrategy> Remove_TransientMap_Unicode()
    {
        var transient = _persistentMapUnicode.ToTransient();
        foreach (var k in _removeKeys) transient.Remove(k);
        return transient.ToPersistent();
    }
    
    [Benchmark]
    public ImmutableSortedDictionary<string, int> Remove_ImmSortedDict()
    {
        var map = _immSortedDict;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }

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

    [Benchmark]
    public LanguageExt.HashMap<string, int> Remove_ExtHashMap()
    {
        var map = _extHashMap;
        foreach (var k in _removeKeys) map = map.Remove(k);
        return map;
    }
}