package utils

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

const shardCount = 256

type LockManager struct {
	shards [shardCount]sync.RWMutex
}

type LockType int

const (
	ReadLock LockType = iota
	WriteLock
)

type Lock struct {
	Path     string
	shard    uint32
	lockType LockType
}

func NewLockManager() *LockManager {
	return &LockManager{}
}

func hashPath(path string) uint32 {
	h := fnv.New32a()
	h.Write([]byte(path))
	return h.Sum32() % shardCount
}

func (lm *LockManager) AcquireRLock(path string) *Lock {
	shard := hashPath(path)
	lm.shards[shard].RLock()

	return &Lock{
		Path:     path,
		shard:    shard,
		lockType: ReadLock,
	}
}

func (lm *LockManager) AcquireLock(path string) *Lock {
	shard := hashPath(path)
	lm.shards[shard].Lock()

	return &Lock{
		Path:     path,
		shard:    shard,
		lockType: WriteLock,
	}
}

func (lm *LockManager) ReleaseLock(lock *Lock) error {
	if lock == nil {
		return fmt.Errorf("lock is nil")
	}
	if lock.shard >= shardCount {
		return fmt.Errorf("invalid shard index")
	}

	switch lock.lockType {
	case ReadLock:
		lm.shards[lock.shard].RUnlock()
	case WriteLock:
		lm.shards[lock.shard].Unlock()
	default:
		return fmt.Errorf("unknown lock type")
	}

	return nil
}

func (lm *LockManager) AcquireOrderedLock(paths ...string) []*Lock {
	sortedPaths := make([]string, len(paths))
	copy(sortedPaths, paths)
	sort.Strings(sortedPaths)

	locks := make([]*Lock, 0, len(sortedPaths))
	for _, path := range sortedPaths {
		lock := lm.AcquireLock(path)
		locks = append(locks, lock)
	}
	return locks
}

func (lm *LockManager) AcquireOrderedRLock(paths ...string) []*Lock {
	sortedPaths := make([]string, len(paths))
	copy(sortedPaths, paths)
	sort.Strings(sortedPaths)

	locks := make([]*Lock, 0, len(sortedPaths))
	for _, path := range sortedPaths {
		lock := lm.AcquireRLock(path)
		locks = append(locks, lock)
	}
	return locks
}

func (lm *LockManager) ReleaseLocks(locks []*Lock) error {
	for i := len(locks) - 1; i >= 0; i-- {
		if err := lm.ReleaseLock(locks[i]); err != nil {
			return err
		}
	}
	return nil
}