//! POSIX record locks: fcntl(F_GETLK / F_SETLK / F_SETLKW), stage 16.
//!
//! A program can lock a byte range of a file: a *read* (shared) lock, which
//! many processes may hold at once, or a *write* (exclusive) lock, which
//! nobody else may overlap. They're "advisory": read() and write() don't
//! look at them, only other fcntl calls do. Databases (SQLite) and mail
//! programs use them so two processes don't change one file at once.
//!
//! The rules, as in Linux:
//!   - locks belong to a *process* (not to an fd): locking a range you
//!     already hold replaces your old lock there (splitting it if needed);
//!     your own locks never conflict with each other;
//!   - closing *any* fd of the file drops all your locks on it, and so does
//!     exiting;
//!   - F_SETLK fails with EAGAIN on a conflict, F_SETLKW waits (a signal
//!     interrupts it), unless waiting would deadlock: EDEADLK.
//!
//! Files are told apart by a key: the path hash stat calls st_ino, or for
//! files without a path (pipes, sockets) the open file's address -- the
//! same number fstat reports as st_ino (`key_of`).
//! Ranges are [start, end) with end = u64::MAX for "to the end of the file
//! and beyond" (l_len = 0).

use crate::errno::Errno;
use crate::thread::WaitQueue;
use alloc::{collections::BTreeMap, vec::Vec};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Lock {
    pub pid: u64,
    pub start: u64,
    pub end: u64,
    pub write: bool,
}

impl Lock {
    fn overlaps(&self, start: u64, end: u64) -> bool {
        self.start < end && start < self.end
    }
}

/// More than this many locks in total: ENOLCK (keeps a runaway program
/// from eating the kernel heap).
const MAX_LOCKS: usize = 1024;

struct Table {
    /// file key -> its locks.
    files: BTreeMap<u64, Vec<Lock>>,
    /// pid -> what it waits for in F_SETLKW (for deadlock detection).
    waiting: BTreeMap<u64, (u64, Lock)>,
}

/// Locked with interrupts off (it's looked at in a wait condition).
static TABLE: Mutex<Table> = Mutex::new(Table { files: BTreeMap::new(), waiting: BTreeMap::new() });
/// F_SETLKW sleepers: woken whenever a lock goes away.
static UNLOCKED: WaitQueue = WaitQueue::new();

fn conflict_in(t: &Table, key: u64, want: &Lock) -> Option<Lock> {
    t.files.get(&key)?.iter().find(|l| l.pid != want.pid && l.overlaps(want.start, want.end) && (l.write || want.write)).copied()
}

/// The key of an open file (see the top).
pub fn key_of(file: &dyn crate::fs::file::File) -> u64 {
    match file.path() {
        Some(path) => crate::fs::file::ino_hash(path.as_bytes()),
        None => file as *const dyn crate::fs::file::File as *const u8 as u64,
    }
}

/// F_GETLK: the first lock that would stop `want`, if any.
pub fn conflict(key: u64, want: Lock) -> Option<Lock> {
    without_interrupts(|| conflict_in(&TABLE.lock(), key, &want))
}

/// Replace `pid`'s locks on [start, end) of file `key` with `new` (None =
/// unlock that range). Doesn't check for conflicts.
fn replace(t: &mut Table, key: u64, pid: u64, start: u64, end: u64, new: Option<bool>) -> Result<(), Errno> {
    let total: usize = t.files.values().map(|v| v.len()).sum();
    let locks = t.files.entry(key).or_default();
    let mut kept = Vec::with_capacity(locks.len() + 2);
    for l in locks.drain(..) {
        if l.pid != pid || !l.overlaps(start, end) {
            kept.push(l);
            continue;
        }
        // Keep the parts of our old lock outside the new range.
        if l.start < start {
            kept.push(Lock { end: start, ..l });
        }
        if l.end > end {
            kept.push(Lock { start: end, ..l });
        }
    }
    if let Some(write) = new {
        if total + 2 > MAX_LOCKS {
            *locks = kept;
            return Err(Errno::ENOLCK);
        }
        kept.push(Lock { pid, start, end, write });
    }
    *locks = kept;
    if locks.is_empty() {
        t.files.remove(&key);
    }
    Ok(())
}

/// F_SETLK: take (Some(write)) or drop (None) a lock. EAGAIN on a conflict.
pub fn set(key: u64, pid: u64, start: u64, end: u64, new: Option<bool>) -> Result<(), Errno> {
    without_interrupts(|| {
        let mut t = TABLE.lock();
        if let Some(write) = new {
            if conflict_in(&t, key, &Lock { pid, start, end, write }).is_some() {
                return Err(Errno::EAGAIN);
            }
        }
        replace(&mut t, key, pid, start, end, new)
    })?;
    UNLOCKED.wake_all(); // (something may have been unlocked)
    Ok(())
}

/// Would `pid` waiting for `want` close a circle of waiters (A waits for
/// B who waits for ... A)?
fn would_deadlock(t: &Table, key: u64, want: &Lock) -> bool {
    let mut next = conflict_in(t, key, want);
    for _ in 0..64 {
        let Some(holder) = next else { return false };
        if holder.pid == want.pid {
            return true;
        }
        next = t.waiting.get(&holder.pid).and_then(|(k, w)| conflict_in(t, *k, w));
    }
    false
}

/// F_SETLKW: like `set`, but wait while somebody else holds a conflicting
/// lock. A signal ends the wait with ERESTARTSYS (EINTR for the program).
pub fn set_wait(key: u64, pid: u64, start: u64, end: u64, write: bool) -> Result<(), Errno> {
    let want = Lock { pid, start, end, write };
    let result = UNLOCKED.wait_until_intr(|| {
        let mut t = TABLE.lock();
        if conflict_in(&t, key, &want).is_none() {
            t.waiting.remove(&pid);
            return Some(replace(&mut t, key, pid, start, end, Some(write)));
        }
        if would_deadlock(&t, key, &want) {
            t.waiting.remove(&pid);
            return Some(Err(Errno::EDEADLK));
        }
        t.waiting.insert(pid, (key, want));
        None
    });
    without_interrupts(|| TABLE.lock().waiting.remove(&pid));
    result?
}

/// `pid` closed an fd of file `key`: all its locks on that file go.
pub fn file_closed(key: u64, pid: u64) {
    let removed = without_interrupts(|| {
        let mut t = TABLE.lock();
        let Some(locks) = t.files.get_mut(&key) else { return false };
        let before = locks.len();
        locks.retain(|l| l.pid != pid);
        let removed = locks.len() != before;
        if locks.is_empty() {
            t.files.remove(&key);
        }
        removed
    });
    if removed {
        UNLOCKED.wake_all();
    }
}

/// Process `pid` exited: drop all its locks.
pub fn process_exited(pid: u64) {
    without_interrupts(|| {
        let mut t = TABLE.lock();
        t.waiting.remove(&pid);
        for locks in t.files.values_mut() {
            locks.retain(|l| l.pid != pid);
        }
        t.files.retain(|_, v| !v.is_empty());
    });
    UNLOCKED.wake_all();
}

/// How many locks exist, record locks and flocks (tests: nothing left behind).
pub fn count() -> usize {
    without_interrupts(|| TABLE.lock().files.values().map(|v| v.len()).sum::<usize>()) + flock_count()
}

// --- flock (stage 17) ------------------------------------------------------
//
// flock(fd, LOCK_SH / LOCK_EX / LOCK_UN, maybe | LOCK_NB) locks a *whole*
// file, BSD style. Unlike the record locks above, a flock belongs to the
// *open file* (what open() made; dup() and fork() share it), not to a
// process: two open()s of the same file in one process do conflict, and the
// lock goes away when the last fd sharing that open file is closed
// (`Drop for InodeFile` calls `flock_release`).
//
// Many holders may share LOCK_SH; LOCK_EX allows nobody else. Asking again
// on an open file that already holds one *converts* it (SH <-> EX).

/// file key -> the open files holding a flock on it: (owner, exclusive).
static FLOCKS: Mutex<BTreeMap<u64, Vec<(usize, bool)>>> = Mutex::new(BTreeMap::new());

fn flock_conflict(t: &BTreeMap<u64, Vec<(usize, bool)>>, key: u64, owner: usize, exclusive: bool) -> bool {
    t.get(&key).is_some_and(|v| v.iter().any(|&(o, ex)| o != owner && (ex || exclusive)))
}

/// Remove `owner`'s flock (any file). True if it had one.
fn flock_drop(t: &mut BTreeMap<u64, Vec<(usize, bool)>>, owner: usize) -> bool {
    let mut found = false;
    for v in t.values_mut() {
        let before = v.len();
        v.retain(|&(o, _)| o != owner);
        found |= v.len() != before;
    }
    t.retain(|_, v| !v.is_empty());
    found
}

/// Take (or convert to) a shared / exclusive flock on file `key` for open
/// file `owner`. `wait`: sleep until it's free (a signal interrupts:
/// ERESTARTSYS); else EAGAIN (= EWOULDBLOCK) at once.
pub fn flock(key: u64, owner: usize, exclusive: bool, wait: bool) -> Result<(), Errno> {
    let try_take = || {
        let mut t = FLOCKS.lock();
        if flock_conflict(&t, key, owner, exclusive) {
            return None;
        }
        let v = t.entry(key).or_default();
        v.retain(|&(o, _)| o != owner);
        v.push((owner, exclusive));
        Some(())
    };
    if !wait {
        return without_interrupts(try_take).ok_or(Errno::EAGAIN);
    }
    // Linux drops the old lock before waiting for a conversion (so two
    // holders upgrading at once can't deadlock); so do we.
    if without_interrupts(|| flock_drop(&mut FLOCKS.lock(), owner)) {
        UNLOCKED.wake_all();
    }
    UNLOCKED.wait_until_intr(try_take)
}

/// LOCK_UN, or the open file went away: drop its flock, if any.
pub fn flock_release(owner: usize) {
    if without_interrupts(|| flock_drop(&mut FLOCKS.lock(), owner)) {
        UNLOCKED.wake_all();
    }
}

/// How many flocks are held (tests: nothing left behind).
pub fn flock_count() -> usize {
    without_interrupts(|| FLOCKS.lock().values().map(|v| v.len()).sum())
}