// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package nfs4server

import (
	"sync"

	"sourcedock.dev/petrbalvin/nfs/internal/nfs4"
	"sourcedock.dev/petrbalvin/nfs/internal/nfsfs"
)

// A lockRange is one held byte range. A length of zero means to the end
// of the file, which the store normalises to the largest offset.
type lockRange struct {
	offset uint64
	length uint64
	write  bool
}

// A lockState is one lock owner's state on one file: the ranges it holds
// and the stateid the client uses for further LOCK and LOCKU calls.
type lockState struct {
	stateid  nfs4.Stateid
	fileKey  string
	clientID uint64
	owner    []byte
	ranges   []lockRange
}

// lockStore keeps the byte range locks of the server. Locks of one owner
// never conflict with each other; overlapping locks of different owners
// do, when either side is a write lock.
type lockStore struct {
	mu      sync.Mutex
	next    uint64
	locks   map[string]*lockState    // by stateid other
	byFile  map[string][]*lockState  // by file key
	byOwn   map[string]*lockState    // by file key + clientid + owner
	waiters map[string][]*lockWaiter // denied locks by file key
}

func newLockStore() *lockStore {
	return &lockStore{
		next:   randCounter(),
		locks:  make(map[string]*lockState),
		byFile: make(map[string][]*lockState),
		byOwn:  make(map[string]*lockState),
	}
}

// ownerKey builds the identity key of one lock owner on one file.
func ownerKey(fileKey string, clientid uint64, owner []byte) string {
	return fileKey + "|" + clientidKey(clientid, owner)
}

// clientidKey renders the client identity part of an owner key.
func clientidKey(clientid uint64, owner []byte) string {
	return ownerKeyHex(clientid) + "|" + string(owner)
}

func ownerKeyHex(clientid uint64) string {
	return hexClientid(clientid)
}

func hexClientid(clientid uint64) string {
	const digits = "0123456789abcdef"
	out := make([]byte, 16)
	for i := range 16 {
		out[15-i] = digits[(clientid>>(4*i))&0xf]
	}
	return string(out)
}

// rangeEnd reports the exclusive end of a range, treating the zero length
// as to the end of the file.
func rangeEnd(r lockRange) uint64 {
	if r.length == 0 {
		return ^uint64(0)
	}
	return r.offset + r.length
}

// rangesConflict reports whether two ranges of the given kinds fight.
// Write locks fight everything they overlap; read locks fight only write
// locks.
func rangesConflict(a, b lockRange) bool {
	if a.offset >= rangeEnd(b) || b.offset >= rangeEnd(a) {
		return false
	}
	return a.write || b.write
}

// lock applies a LOCK of the given lock owner against the file. Locks of
// one owner coexist; overlapping locks of different owners are denied.
func (s *lockStore) lock(fh nfsfs.Handle, lockClientid uint64, lockOwner []byte, write bool, offset, length uint64) (nfs4.Stateid, uint32) {
	key := fileKey(fh)
	s.mu.Lock()
	defer s.mu.Unlock()
	probe := lockRange{offset: offset, length: length, write: write}
	// A lock owner never conflicts with itself; every other live lock over
	// an overlapping range denies the lock.
	for _, other := range s.byFile[key] {
		if sameIdentity(other, lockClientid, lockOwner) {
			continue
		}
		for _, r := range other.ranges {
			if rangesConflict(probe, r) {
				return nfs4.Stateid{}, nfs4.ErrDenied
			}
		}
	}
	oid := ownerKey(key, lockClientid, lockOwner)
	ls := s.byOwn[oid]
	if ls == nil {
		s.next++
		ls = &lockState{stateid: lockStateid(s.next), fileKey: key, clientID: lockClientid, owner: lockOwner}
		s.locks[string(ls.stateid[:])] = ls
		s.byOwn[oid] = ls
		s.byFile[key] = append(s.byFile[key], ls)
	}
	ls.ranges = addRange(ls.ranges, lockRange{offset: offset, length: length, write: write})
	return ls.stateid, nfs4.ErrOK
}

// sameIdentity reports whether a lock state belongs to the given lock
// owner identity.
func sameIdentity(ls *lockState, clientid uint64, owner []byte) bool {
	return ls != nil && ls.clientID == clientid && string(ls.owner) == string(owner)
}

// unlock releases the range of a lock stateid, splitting held ranges as
// needed, and bumps the stateid sequence. Only the client the lock
// state belongs to may release it. A cleared sequence names the
// current version, RFC 8881 section 8.2.2. Releasing a range that is
// not held is not an error.
func (s *lockStore) unlock(st nfs4.Stateid, clientid, offset, length uint64) (nfs4.Stateid, uint32) {
	s.mu.Lock()
	defer s.mu.Unlock()
	ls, ok := s.locks[string(st[:])]
	if !ok && stateidSeq(st) == 0 {
		// A cleared sequence names the current version: resolve by the
		// other field among the client's lock states.
		for _, cand := range s.locks {
			if string(cand.stateid[4:]) == string(st[4:]) && cand.clientID == clientid {
				ls, ok = cand, true
				break
			}
		}
	}
	if !ok {
		return nfs4.Stateid{}, nfs4.ErrBadStateid
	}
	if ls.clientID != clientid {
		return nfs4.Stateid{}, nfs4.ErrBadStateid
	}
	if length == 0 {
		length = ^uint64(0) - offset
	}
	ls.ranges = subtractRange(ls.ranges, lockRange{offset: offset, length: length})
	setStateidSeq(&ls.stateid, stateidSeq(ls.stateid)+1)
	return ls.stateid, nfs4.ErrOK
}

// locksHeldOn reports whether any lock state of the file still holds
// ranges, which is the CLOSE guard.
func (s *lockStore) locksHeldOn(fh nfsfs.Handle) bool {
	key := fileKey(fh)
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, ls := range s.byFile[key] {
		if len(ls.ranges) > 0 {
			return true
		}
	}
	return false
}

// dropClient releases every lock state of the client, which DESTROY_CLIENT
// ID and lease expiry require.
func (s *lockStore) dropClient(clientid uint64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for oid, ls := range s.byOwn {
		if ls.clientID != clientid {
			continue
		}
		delete(s.locks, string(ls.stateid[:]))
		delete(s.byOwn, oid)
		list := s.byFile[ls.fileKey]
		for i, cand := range list {
			if cand == ls {
				s.byFile[ls.fileKey] = append(list[:i], list[i+1:]...)
				break
			}
		}
	}
}

// lockStateid builds a fresh lock stateid: sequence one and the counter in
// the other field.
func lockStateid(counter uint64) nfs4.Stateid {
	var st nfs4.Stateid
	setStateidSeq(&st, 1)
	st[4], st[5], st[6], st[7] = 'L', 'O', 'C', 'K'
	for i := range 8 {
		st[15-i] = byte(counter >> (8 * i))
	}
	return st
}

// addRange inserts a range into the held list, merging it with the
// overlapping and adjacent ranges of the same kind, RFC 8881 section
// 18.10: a new lock over a range the owner already holds replaces it,
// never piles a second entry on top. A mixed read and write overlap is
// resolved by subtraction first: the write form wins the intersection.
func addRange(ranges []lockRange, r lockRange) []lockRange {
	// Carve the new range out of every held range of the other kind,
	// then absorb it into the held ranges of its own kind.
	var carved []lockRange
	for _, held := range ranges {
		if held.write == r.write {
			carved = append(carved, held)
			continue
		}
		end := rangeEnd(r)
		heldEnd := rangeEnd(held)
		if r.offset >= heldEnd || held.offset >= end {
			carved = append(carved, held)
			continue
		}
		if held.offset < r.offset {
			carved = append(carved, lockRange{held.offset, r.offset - held.offset, held.write})
		}
		if heldEnd > end {
			carved = append(carved, lockRange{end, heldEnd - end, held.write})
		}
	}
	var out []lockRange
	merged := r
	for _, held := range carved {
		if held.write != r.write {
			out = append(out, held)
			continue
		}
		if held.offset > rangeEnd(merged) || merged.offset > rangeEnd(held) {
			out = append(out, held) // disjoint: keep both
			continue
		}
		// Overlapping or adjacent: widen the merge, keeping both ends.
		end := rangeEnd(merged)
		if rangeEnd(held) > end {
			end = rangeEnd(held)
		}
		if held.offset < merged.offset {
			merged.offset = held.offset
		}
		if merged.length == 0 || held.length == 0 {
			merged.length = 0 // spans to the end of the file
		} else {
			merged.length = end - merged.offset
		}
	}
	return append(out, merged)
}

// subtractRange removes the interval of r from every held range, splitting
// the ranges it cuts through.
func subtractRange(ranges []lockRange, r lockRange) []lockRange {
	end := rangeEnd(r)
	var out []lockRange
	for _, held := range ranges {
		heldEnd := rangeEnd(held)
		if r.offset >= rangeEnd(held) || held.offset >= end {
			out = append(out, held)
			continue
		}
		if held.offset < r.offset {
			out = append(out, lockRange{held.offset, r.offset - held.offset, held.write})
		}
		if heldEnd > end {
			out = append(out, lockRange{end, heldEnd - end, held.write})
		}
	}
	return out
}

// byStateid resolves a lock state by its stateid, verifies that it
// names the current file and belongs to the asking client. A stateid
// presented with a cleared sequence names the current version, RFC
// 8881 section 8.2.2.
func (s *lockStore) byStateid(st nfs4.Stateid, fh nfsfs.Handle, clientid uint64) (*lockState, uint32) {
	key := fileKey(fh)
	s.mu.Lock()
	defer s.mu.Unlock()
	ls, ok := s.locks[string(st[:])]
	if !ok && stateidSeq(st) == 0 {
		ls = s.byOther(st[4:], clientid, key)
		if ls == nil {
			return nil, nfs4.ErrBadStateid
		}
	}
	if ls == nil {
		return nil, nfs4.ErrBadStateid
	}
	if ls.fileKey != key || ls.clientID != clientid {
		return nil, nfs4.ErrBadStateid
	}
	return ls, nfs4.ErrOK
}

// byOther resolves a lock state by its other field alone.
func (s *lockStore) byOther(other []byte, clientid uint64, key string) *lockState {
	for _, ls := range s.locks {
		if string(ls.stateid[4:]) == string(other) && ls.clientID == clientid && ls.fileKey == key {
			return ls
		}
	}
	return nil
}

// test serves LOCKT: it scans the locks of the current file for a conflict
// with the requested range, ignoring the locks of the probing owner, and
// reports the holder of the first conflicting lock.
func (s *lockStore) test(fh nfsfs.Handle, lockClientid uint64, lockOwner []byte, write bool, offset, length uint64) (nfs4.LockDenied, uint32) {
	key := fileKey(fh)
	s.mu.Lock()
	defer s.mu.Unlock()
	probe := lockRange{offset: offset, length: length, write: write}
	for _, other := range s.byFile[key] {
		if other.clientID == lockClientid && string(other.owner) == string(lockOwner) {
			continue
		}
		for _, r := range other.ranges {
			if !rangesConflict(probe, r) {
				continue
			}
			denied := nfs4.LockDenied{
				Offset:   r.offset,
				Length:   r.length,
				LockType: nfs4.LockTypeRead,
				ClientID: other.clientID,
				Owner:    other.owner,
			}
			if r.write {
				denied.LockType = nfs4.LockTypeWrite
			}
			return denied, nfs4.ErrDenied
		}
	}
	return nfs4.LockDenied{}, nfs4.ErrOK
}

// releaseOwner retires every lock state of one lock owner, which
// RELEASE_LOCKOWNER requires. Locks the owner still holds go with it.
func (s *lockStore) releaseOwner(clientid uint64, owner []byte) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for oid, ls := range s.byOwn {
		if ls.clientID != clientid || string(ls.owner) != string(owner) {
			continue
		}
		delete(s.locks, string(ls.stateid[:]))
		delete(s.byOwn, oid)
		list := s.byFile[ls.fileKey]
		for i, cand := range list {
			if cand == ls {
				s.byFile[ls.fileKey] = append(list[:i], list[i+1:]...)
				break
			}
		}
	}
}

// freeStateid retires the lock state the stateid names, which
// FREE_STATEID requires after every lock of the owner is gone. Only the
// client the lock state belongs to may retire it.
func (s *lockStore) freeStateid(st nfs4.Stateid, clientid uint64) uint32 {
	s.mu.Lock()
	defer s.mu.Unlock()
	ls, ok := s.locks[string(st[:])]
	if !ok {
		return nfs4.ErrBadStateid
	}
	if ls.clientID != clientid {
		return nfs4.ErrBadStateid
	}
	if len(ls.ranges) > 0 {
		return nfs4.ErrLocksHeld
	}
	delete(s.locks, string(ls.stateid[:]))
	oid := ownerKey(ls.fileKey, ls.clientID, ls.owner)
	if s.byOwn[oid] == ls {
		delete(s.byOwn, oid)
	}
	list := s.byFile[ls.fileKey]
	for i, cand := range list {
		if cand == ls {
			s.byFile[ls.fileKey] = append(list[:i], list[i+1:]...)
			break
		}
	}
	return nfs4.ErrOK
}

// hasStateid reports whether the lock stateid is live.
func (s *lockStore) hasStateid(st nfs4.Stateid) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, ok := s.locks[string(st[:])]
	return ok
}

// A lockWaiter records one denied LOCK, so the server can notify the
// owner when a conflicting range is released: the CB_NOTIFY_LOCK of RFC
// 8881 section 20.6.
type lockWaiter struct {
	clientID uint64
	owner    []byte
	offset   uint64
	length   uint64
	write    bool
}

// addWaiter records a denied lock request on the file.
func (s *lockStore) addWaiter(fh nfsfs.Handle, clientid uint64, owner []byte, offset, length uint64, write bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.waiters == nil {
		s.waiters = make(map[string][]*lockWaiter)
	}
	s.waiters[fileKey(fh)] = append(s.waiters[fileKey(fh)], &lockWaiter{
		clientID: clientid, owner: owner, offset: offset, length: length, write: write,
	})
}

// takeWaiters returns and forgets every waiter of the file whose range
// overlaps the just released range: their conflict may be gone.
func (s *lockStore) takeWaiters(fh nfsfs.Handle, offset, length uint64) []*lockWaiter {
	s.mu.Lock()
	defer s.mu.Unlock()
	key := fileKey(fh)
	list := s.waiters[key]
	if len(list) == 0 {
		return nil
	}
	delete(s.waiters, key)
	end := offset + length
	if length == 0 {
		end = ^uint64(0)
	}
	var out []*lockWaiter
	for _, w := range list {
		wEnd := w.offset + w.length
		if w.length == 0 {
			wEnd = ^uint64(0)
		}
		if w.offset < end && offset < wEnd {
			out = append(out, w)
		}
	}
	return out
}