// Exercise 09 — Concurrency-Safe Counters (three ways)
//
// ASSIGNMENT
// ----------
// A warm-up screener that interviewers run with -race. Implement THREE
// versions of a keyed counter and one plain counter:
//
//  1. MutexCounter: map[string]int64 guarded by a sync.Mutex.
//     Inc(key), Get(key), Snapshot() (a COPY of the map — why a copy? if you
//     return the internal map, the caller can race on it).
//  2. AtomicCounter: a single int64 total using sync/atomic only.
//     Inc(), Get(). No mutex allowed.
//  3. ShardedCounter: N shards, each its own mutex+map, key routed by hash.
//     Same interface as MutexCounter. This is how you fix mutex contention
//     on hot maps.
//
// Run the tests:   go test -race ./09-safe-counter/
// Then benchmark:  go test -bench=. -cpu=1,4,8 ./09-safe-counter/
package main

import (
	"fmt"
	"hash/fnv"
	"sync"
	"sync/atomic"
)

// ---- 1. Mutex version ----

type MutexCounter struct {
	mu sync.Mutex
	m  map[string]int64
}

func NewMutexCounter() *MutexCounter {
	return &MutexCounter{
		m: make(map[string]int64),
	}
}
func (c *MutexCounter) Inc(key string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.m[key]++

}
func (c *MutexCounter) Get(key string) int64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	return (c.m[key])
}
func (c *MutexCounter) Snapshot() map[string]int64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	// possibly could use maps.clone
	outmap := make(map[string]int64, len(c.m))
	for k, v := range c.m {
		outmap[k] = v
	}
	return outmap
}

// ---- 2. Atomic version ----

type AtomicCounter struct {
	n atomic.Int64
}

func NewAtomicCounter() *AtomicCounter { return &AtomicCounter{} }
func (c *AtomicCounter) Inc()          { c.n.Add(1) }
func (c *AtomicCounter) Get() int64    { return c.n.Load() }

// ---- 3. Sharded version ----

const shards = 16

type shard struct {
	mu sync.Mutex
	m  map[string]int64
}
type ShardedCounter struct {
	s [shards]shard
	// route with fnv-1a: hash(key) % shards (hash/fnv or hand-rolled)
}

func NewShardedCounter() *ShardedCounter {
	c := &ShardedCounter{}
	// loop over shards and initialize each shard's map
	for i := 0; i < shards; i++ {
		c.s[i].m = make(map[string]int64)
	}
	return c
}
func (c *ShardedCounter) Inc(key string) {
	// find shard for the key
	hashF := fnv.New32a()
	hashF.Write([]byte(key))
	shard := hashF.Sum32() % shards

	c.s[shard].mu.Lock()
	defer c.s[shard].mu.Unlock()
	// fmt.Println("shard", shard, "key", key, "value", c.s[shard].m[key])
	c.s[shard].m[key]++
}
func (c *ShardedCounter) Get(key string) int64 {
	// find shard for the key
	hashF := fnv.New32a()
	hashF.Write([]byte(key))
	shard := hashF.Sum32() % shards

	c.s[shard].mu.Lock()
	defer c.s[shard].mu.Unlock()
	return c.s[shard].m[key]
}
func (c *ShardedCounter) Snapshot() map[string]int64 {
	out := make(map[string]int64)
	for i := 0; i < shards; i++ {
		c.s[i].mu.Lock()
		for k, v := range c.s[i].m {
			out[k] = v
		}
		c.s[i].mu.Unlock()
	}
	return out
}

func main() {
	c := NewMutexCounter()
	c.Inc("requests")
	c.Inc("requests")
	fmt.Println(c.Get("requests"), c.Snapshot())
}