package plog

import "sync"

// GroupOffsets tracks, per consumer group and partition, the next offset
// to read — so a group resumes where it left off after a restart, and
// different groups consume the same log independently at their own pace.
type GroupOffsets struct {
	mu        sync.Mutex
	committed map[string]map[int]int64 // group -> partition -> next offset
}

func NewGroupOffsets() *GroupOffsets {
	return &GroupOffsets{committed: map[string]map[int]int64{}}
}

// Commit records that `group` has processed through (and including) the
// record before nextOffset on `partition` (i.e. resume at nextOffset).
func (g *GroupOffsets) Commit(group string, partition int, nextOffset int64) {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.committed[group] == nil {
		g.committed[group] = map[int]int64{}
	}
	g.committed[group][partition] = nextOffset
}

// Committed returns the next offset to read for a group/partition (0 if
// the group has never committed — i.e. start from the beginning).
func (g *GroupOffsets) Committed(group string, partition int) int64 {
	g.mu.Lock()
	defer g.mu.Unlock()
	if m, ok := g.committed[group]; ok {
		return m[partition]
	}
	return 0
}

// AssignRange splits `partitions` across `consumers` in contiguous ranges
// (Kafka's RangeAssignor): earlier consumers get one extra partition when
// it doesn't divide evenly. Returns consumer index -> partition ids.
func AssignRange(partitions, consumers int) map[int][]int {
	out := map[int][]int{}
	if consumers <= 0 {
		return out
	}
	per := partitions / consumers
	extra := partitions % consumers
	start := 0
	for c := 0; c < consumers; c++ {
		count := per
		if c < extra {
			count++
		}
		for p := start; p < start+count; p++ {
			out[c] = append(out[c], p)
		}
		start += count
	}
	return out
}

// AssignRoundRobin spreads partitions across consumers one at a time
// (Kafka's RoundRobinAssignor): partition p -> consumer p % consumers.
// More even than range when partition counts vary across topics.
func AssignRoundRobin(partitions, consumers int) map[int][]int {
	out := map[int][]int{}
	if consumers <= 0 {
		return out
	}
	for p := 0; p < partitions; p++ {
		c := p % consumers
		out[c] = append(out[c], p)
	}
	return out
}