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]) Fill(from, to int, data D) {
	if to < from {
		return
	}

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

	for i := leftIndex; i <= rightIndex; i++ {
		r := rm.ranges[i]
		// Range is not inside fill zone or is already filled.
		if r.from > to || (r.to < from && r.to != -1) || r.data != nil {
			continue
		}

		if r.from >= from && r.to <= to && r.to != -1 {
			// The empty range is completely contained in the range to fill, overwrite it.
			rm.ranges[i].data = &data
		} else {
			// The empty range is partially contained in the range to fill, split it.
			rm.Set(max(r.from, from), min(r.to, to), data)
		}
	}
}

func (rm *RangeMap[D]) AppendRanges(other RangeMap[D]) {
	newRanges := make([]rangeEntry[D], 0, len(rm.ranges)+len(other.ranges))
	offset := rm.ranges[len(rm.ranges)-1].from

	// No need to copy the sentinel.
	newRanges = append(newRanges, rm.ranges[:len(rm.ranges)-1]...)

	for _, otherEntry := range other.ranges {
		to := otherEntry.to + offset
		if otherEntry.to == -1 {
			to = -1 // Sentinel not affected by offset.
		}

		newRanges = append(newRanges, rangeEntry[D]{
			from: otherEntry.from + offset,
			to:   to,
			data: otherEntry.data,
		})
	}

	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
	}

	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
}

func min(a, b int) int {
	if a < b && a != -1 {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b && b != -1 {
		return a
	}
	return b
}