package rangemap

import (
	"fmt"
	"strings"
)

type RangeMap[D any] struct {
	ranges []rangeEntry[D]
}

func New[D any]() RangeMap[D] {
	return RangeMap[D]{
		// First entry is a sentinel.
		ranges: []rangeEntry[D]{
			{
				from: 0,
				to:   -1,
				data: nil,
			},
		},
	}
}

func (rm *RangeMap[D]) AppendToLast(to int, data D) {
	lastEntry := &rm.ranges[len(rm.ranges)-1]

	rm.Set(lastEntry.from, to, data)
}

func (rm *RangeMap[D]) Set(from, to int, data D) {
	if to < from {
		return
	}

	newRanges := make([]rangeEntry[D], 0, len(rm.ranges)+2)

	leftIndex := rm.entry(from)
	rightIndex := rm.entry(to)

	newRanges = append(newRanges, rm.ranges[:leftIndex]...)

	leftEntry := rm.ranges[leftIndex]
	if leftEntry.to < from && leftEntry.to != -1 {
		newRanges = append(newRanges, leftEntry)
	} else if leftEntry.from < from {
		newRanges = append(newRanges, rangeEntry[D]{
			from: leftEntry.from,
			to:   from - 1,
			data: leftEntry.data,
		})
	}

	newRanges = append(newRanges, rangeEntry[D]{
		from: from,
		to:   to,
		data: &data,
	})

	rightEntry := rm.ranges[rightIndex]
	if rightEntry.from > to {
		newRanges = append(newRanges, rightEntry)
	} else if rightEntry.to > to || rightEntry.to == -1 {
		newRanges = append(newRanges, rangeEntry[D]{
			from: to + 1,
			to:   rightEntry.to,
			data: rightEntry.data,
		})
	}

	newRanges = append(newRanges, rm.ranges[rightIndex+1:]...)

	rm.ranges = newRanges
}

func (rm *RangeMap[D]) Get(i int) *D {
	return rm.ranges[rm.entry(i)].data
}

func (rm *RangeMap[D]) Len() int {
	return len(rm.ranges)
}

func (rm *RangeMap[D]) String() string {
	ranges := make([]string, 0, len(rm.ranges))

	for _, r := range rm.ranges {
		if r.data == nil {
			ranges = append(ranges, fmt.Sprintf("{%d..%d}", r.from, r.to))
		} else {
			ranges = append(ranges, fmt.Sprintf("{%d..%d: %v}", r.from, r.to, *r.data))
		}
	}

	return fmt.Sprintf("[%s]", strings.Join(ranges, ", "))
}

func (rm *RangeMap[D]) entry(i int) int {
	if i < 0 {
		return 0
	}

	lastRange := rm.ranges[len(rm.ranges)-1]
	if i >= lastRange.to && lastRange.to != -1 {
		return len(rm.ranges) - 1
	}

	left := 0
	right := len(rm.ranges) - 1
	for left <= right {
		mid := (left + right) / 2
		entry := rm.ranges[mid]
		if i < entry.from {
			right = mid - 1
		} else if i > entry.to && entry.to != -1 {
			left = mid + 1
		} else {
			return mid
		}
	}

	panic("unreachable")
}

type rangeEntry[D any] struct {
	from int
	to   int
	data *D
}