using System;
using System.Collections.Generic;
using System.Linq;

namespace AibisDream.Framework
{
    public class RandomSelector<T>
    {
        private readonly List<T> _items;
        private readonly List<int> _prefixSums;
        private readonly int _totalWeight;
        private readonly Random _random;
        
        public RandomSelector(IEnumerable<T> items, Func<T, int> weightSelector)
        {
            if (items == null) throw new ArgumentNullException(nameof(items));
            if (weightSelector == null) throw new ArgumentNullException(nameof(weightSelector));

            var itemList = items.ToList();
            _items = new List<T>(itemList.Count);
            _prefixSums = new List<int>(itemList.Count);
        
            int sum = 0;
            foreach (var item in itemList)
            {
                int weight = weightSelector(item);
                if (weight < 0)
                    throw new ArgumentException("Weights cannot be negative", nameof(items));

                if (weight > 0)
                {
                    sum += weight;
                    _prefixSums.Add(sum);
                    _items.Add(item);
                }
            }

            _totalWeight = sum;
            _random = new Random();
        }
        
        public T Select()
        {
            if (_totalWeight == 0)
                throw new InvalidOperationException("All weights are zero");

            int randomValue = _random.Next(0, _totalWeight);
            int index = BinarySearch(randomValue);
            return _items[index];
        }

        private int BinarySearch(int value)
        {
            int left = 0;
            int right = _prefixSums.Count - 1;

            while (left < right)
            {
                int mid = left + (right - left) / 2;
                if (_prefixSums[mid] <= value)
                {
                    left = mid + 1;
                }
                else
                {
                    right = mid;
                }
            }

            return left;
        }
    }
}