use std::io;
use std::os::fd::{FromRawFd, OwnedFd, RawFd};
use std::os::unix::fs::{MetadataExt, PermissionsExt};
use std::path::Path;

const MAX_WINSIZE: u16 = 10_000;

/// Create a directory hierarchy with mode 0700, validating ownership of existing components.
/// Trusted system roots (`/`, `/tmp`, `/run`, `$XDG_RUNTIME_DIR`) are accepted without
/// ownership checks. All other existing directories must be owned by the current user
/// and must not be symlinks.
///
/// Uses create-then-validate instead of check-then-create to avoid TOCTOU races.
pub fn secure_create_dir_all(path: &Path) -> io::Result<()> {
    if is_trusted_root(path) {
        return Ok(());
    }

    if let Some(parent) = path.parent() {
        secure_create_dir_all(parent)?;
    }

    // Try to create; if it already exists, validate ownership/type.
    // The umask (0o077 set by daemon) ensures 0o700 mode on creation.
    match std::fs::create_dir(path) {
        Ok(()) => {
            // Explicitly set permissions in case umask wasn't set (e.g. client-side calls).
            std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
            Ok(())
        }
        Err(e) if e.kind() == io::ErrorKind::AlreadyExists => validate_dir(path),
        Err(e) => Err(e),
    }
}

/// Bind a `UnixListener` with TOCTOU-safe stale socket handling and 0600 permissions.
///
/// On `AddrInUse`, probes the existing socket: if it responds to a connect, returns
/// an error (socket is alive). Otherwise, removes the stale socket and retries.
///
/// The whole probe/remove/bind sequence is serialized against concurrent callers via
/// an exclusive flock on a companion `<path>.bindlock` file. Without the lock, two
/// racers that both observe a dead socket can each `remove_file` + `bind`: the second
/// unlinks the first racer's just-bound live socket, orphaning its listener.
pub fn bind_unix_listener(path: &Path) -> io::Result<tokio::net::UnixListener> {
    let lock_path = bind_lock_path(path);
    let _guard = acquire_bind_lock(&lock_path)?;
    match tokio::net::UnixListener::bind(path) {
        Ok(listener) => {
            set_socket_permissions(path)?;
            Ok(listener)
        }
        Err(e) if e.kind() == io::ErrorKind::AddrInUse => {
            match std::os::unix::net::UnixStream::connect(path) {
                Ok(_) => Err(io::Error::new(
                    io::ErrorKind::AddrInUse,
                    format!("{} is already in use by a running process", path.display()),
                )),
                Err(_) => {
                    std::fs::remove_file(path)?;
                    let listener = tokio::net::UnixListener::bind(path)?;
                    set_socket_permissions(path)?;
                    Ok(listener)
                }
            }
        }
        Err(e) => Err(e),
    }
}

fn bind_lock_path(socket_path: &Path) -> std::path::PathBuf {
    let mut os = socket_path.as_os_str().to_owned();
    os.push(".bindlock");
    std::path::PathBuf::from(os)
}

/// Open (creating if needed) the bind-lock file with 0600 and acquire an exclusive
/// flock, retrying briefly if another process is racing us through bind_unix_listener.
fn acquire_bind_lock(lock_path: &Path) -> io::Result<BindLockGuard> {
    use std::fs::OpenOptions;
    use std::os::unix::fs::OpenOptionsExt;
    let file =
        OpenOptions::new().create(true).truncate(false).write(true).mode(0o600).open(lock_path)?;
    // Up to ~2s total: one racer's probe/remove/bind is a handful of syscalls.
    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2);
    loop {
        match nix::fcntl::Flock::lock(
            file.try_clone()?,
            nix::fcntl::FlockArg::LockExclusiveNonblock,
        ) {
            Ok(flock) => return Ok(BindLockGuard { _flock: flock }),
            Err((_, nix::errno::Errno::EWOULDBLOCK)) if std::time::Instant::now() < deadline => {
                std::thread::sleep(std::time::Duration::from_millis(25));
                continue;
            }
            Err((_, e)) => {
                return Err(io::Error::other(format!(
                    "could not acquire bind lock {}: {e}",
                    lock_path.display()
                )));
            }
        }
    }
}

struct BindLockGuard {
    _flock: nix::fcntl::Flock<std::fs::File>,
}

/// Remove the companion `.bindlock` of a socket that is being cleaned up.
///
/// Every `bind_unix_listener` call creates a `<path>.bindlock`; without this,
/// they accumulate one per socket ever bound. Callers that remove a socket
/// (client fwd cleanup, server session-socket cleanup, daemon shutdown) call
/// this right after removing the socket itself.
///
/// The file is only removed when no concurrent binder holds its flock -- a
/// held lock means someone is mid-probe/bind on this path, and unlinking the
/// file under them would let a third racer create a fresh lock file and bind
/// unserialized (the race the bindlock exists to prevent). Best-effort: any
/// error leaves the file in place, which is always safe.
pub fn remove_bind_lock_if_unheld(socket_path: &Path) {
    let lock_path = bind_lock_path(socket_path);
    // No create: if there is no bindlock there is nothing to clean.
    let Ok(file) = std::fs::OpenOptions::new().write(true).open(&lock_path) else {
        return;
    };
    if let Ok(_flock) = nix::fcntl::Flock::lock(file, nix::fcntl::FlockArg::LockExclusiveNonblock) {
        // Unlink while holding the lock: a racer that opens the path after
        // this unlink creates a brand-new inode and locks it independently,
        // so there is never a window where the path points at an inode we
        // are about to delete (same ordering rationale as the tunnel lock).
        let _ = std::fs::remove_file(&lock_path);
    }
}

/// Connect to a Unix socket and verify the peer's UID matches ours.
/// All client-side connects to daemon/session sockets MUST go through this.
pub async fn connect_verified(path: &Path) -> io::Result<tokio::net::UnixStream> {
    let stream = tokio::net::UnixStream::connect(path).await?;
    verify_peer_uid(&stream)?;
    Ok(stream)
}

/// Verify that the peer on a Unix stream has the same UID as the current process.
pub fn verify_peer_uid(stream: &tokio::net::UnixStream) -> io::Result<()> {
    let cred = stream.peer_cred()?;
    let my_uid = unsafe { libc::getuid() };
    if cred.uid() != my_uid {
        return Err(io::Error::new(
            io::ErrorKind::PermissionDenied,
            format!("rejecting connection from uid {} (expected {my_uid})", cred.uid()),
        ));
    }
    Ok(())
}

/// `dup(2)` that returns an `OwnedFd` with `FD_CLOEXEC` set, or an error.
pub fn checked_dup(fd: RawFd) -> io::Result<OwnedFd> {
    let new_fd = unsafe { libc::fcntl(fd, libc::F_DUPFD_CLOEXEC, 0) };
    if new_fd == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(unsafe { OwnedFd::from_raw_fd(new_fd) })
}

/// Set `FD_CLOEXEC` on a raw fd so it is not inherited across `exec`.
pub fn set_cloexec(fd: RawFd) -> io::Result<()> {
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
    if flags == -1 {
        return Err(io::Error::last_os_error());
    }
    if unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) } == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

/// Clamp window-size values to a sane range, preventing zero-sized or absurdly large values.
pub fn clamp_winsize(cols: u16, rows: u16) -> (u16, u16) {
    (cols.clamp(1, MAX_WINSIZE), rows.clamp(1, MAX_WINSIZE))
}

fn set_socket_permissions(path: &Path) -> io::Result<()> {
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
}

fn is_trusted_root(path: &Path) -> bool {
    if matches!(path.to_str(), Some("/" | "/tmp" | "/run")) {
        return true;
    }
    env_path_matches("XDG_RUNTIME_DIR", path) || env_path_matches("TMPDIR", path)
}

fn env_path_matches(var: &str, path: &Path) -> bool {
    std::env::var(var).ok().is_some_and(|v| path == Path::new(v.trim_end_matches('/')))
}

fn validate_dir(path: &Path) -> io::Result<()> {
    let meta = std::fs::symlink_metadata(path)?;

    // Root-owned entries are system-managed (e.g. /var -> /private/var on macOS).
    // Just verify the path resolves to a directory.
    if meta.uid() == 0 {
        if !path.is_dir() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!("{} does not resolve to a directory", path.display()),
            ));
        }
        return Ok(());
    }

    if meta.file_type().is_symlink() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("refusing to use symlink at {}", path.display()),
        ));
    }
    if !meta.is_dir() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("{} is not a directory", path.display()),
        ));
    }

    let uid = unsafe { libc::getuid() };
    if meta.uid() != uid {
        return Err(io::Error::new(
            io::ErrorKind::PermissionDenied,
            format!(
                "{} is owned by uid {}, expected uid {uid}; \
                 set $XDG_RUNTIME_DIR or use --ctl-socket",
                path.display(),
                meta.uid()
            ),
        ));
    }

    // StrictModes: if group/other have any access, tighten to 0700. A writable
    // socket directory lets another local user unlink and re-bind ctl.sock.
    if meta.mode() & 0o077 != 0 {
        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn clamp_winsize_zeros_to_minimum() {
        assert_eq!(clamp_winsize(0, 0), (1, 1));
    }

    #[test]
    fn clamp_winsize_normal_passthrough() {
        assert_eq!(clamp_winsize(80, 24), (80, 24));
    }

    #[test]
    fn clamp_winsize_max_boundary() {
        assert_eq!(clamp_winsize(10_000, 10_000), (10_000, 10_000));
    }

    #[test]
    fn clamp_winsize_over_max_clamped() {
        assert_eq!(clamp_winsize(10_001, 10_001), (10_000, 10_000));
    }

    #[test]
    fn clamp_winsize_extreme_values() {
        assert_eq!(clamp_winsize(u16::MAX, u16::MAX), (10_000, 10_000));
    }

    #[test]
    fn clamp_winsize_asymmetric() {
        assert_eq!(clamp_winsize(0, 80), (1, 80));
        assert_eq!(clamp_winsize(20_000, 5), (10_000, 5));
    }

    // --- secure_create_dir_all ---

    #[test]
    fn secure_create_dir_all_fresh_hierarchy() {
        let tmp = tempfile::tempdir().unwrap();
        let deep = tmp.path().join("a").join("b").join("c");
        secure_create_dir_all(&deep).unwrap();
        assert!(deep.is_dir());
        let mode = std::fs::metadata(&deep).unwrap().permissions().mode() & 0o777;
        assert_eq!(mode, 0o700);
    }

    #[test]
    fn secure_create_dir_all_idempotent() {
        let tmp = tempfile::tempdir().unwrap();
        let dir = tmp.path().join("mydir");
        secure_create_dir_all(&dir).unwrap();
        secure_create_dir_all(&dir).unwrap(); // second call succeeds
    }

    #[test]
    fn secure_create_dir_all_tightens_loose_mode() {
        let tmp = tempfile::tempdir().unwrap();
        let dir = tmp.path().join("loose");
        std::fs::create_dir(&dir).unwrap();
        std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o777)).unwrap();
        secure_create_dir_all(&dir).unwrap();
        let mode = std::fs::metadata(&dir).unwrap().permissions().mode() & 0o777;
        assert_eq!(mode, 0o700);
    }

    #[test]
    fn secure_create_dir_all_rejects_symlink() {
        let tmp = tempfile::tempdir().unwrap();
        let real = tmp.path().join("real");
        std::fs::create_dir(&real).unwrap();
        let link = tmp.path().join("link");
        std::os::unix::fs::symlink(&real, &link).unwrap();
        let err = secure_create_dir_all(&link).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn secure_create_dir_all_rejects_regular_file() {
        let tmp = tempfile::tempdir().unwrap();
        let file = tmp.path().join("not_a_dir");
        std::fs::write(&file, b"").unwrap();
        let err = secure_create_dir_all(&file).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn secure_create_dir_all_trusted_roots() {
        // These should succeed without ownership checks
        secure_create_dir_all(Path::new("/")).unwrap();
        secure_create_dir_all(Path::new("/tmp")).unwrap();
    }

    // --- bind_unix_listener ---

    #[tokio::test]
    async fn bind_unix_listener_fresh() {
        let tmp = tempfile::tempdir().unwrap();
        let sock = tmp.path().join("test.sock");
        let _listener = bind_unix_listener(&sock).unwrap();
        let mode = std::fs::metadata(&sock).unwrap().permissions().mode() & 0o777;
        assert_eq!(mode, 0o600);
    }

    #[tokio::test]
    async fn bind_unix_listener_stale_socket() {
        let tmp = tempfile::tempdir().unwrap();
        let sock = tmp.path().join("stale.sock");
        // Create a stale socket file using UnixDatagram (never calls listen()).
        // This avoids a macOS kernel race where connect() briefly succeeds on a
        // just-closed listening socket.
        drop(std::os::unix::net::UnixDatagram::bind(&sock).unwrap());
        assert!(sock.exists());
        // Re-bind should clean up stale socket and succeed
        let _listener = bind_unix_listener(&sock).unwrap();
    }

    #[tokio::test]
    async fn bind_unix_listener_live_socket_errors() {
        let tmp = tempfile::tempdir().unwrap();
        let sock = tmp.path().join("live.sock");
        let _listener = bind_unix_listener(&sock).unwrap();
        // Try to bind again while listener is alive
        let err = bind_unix_listener(&sock).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::AddrInUse);
    }

    #[tokio::test]
    async fn bind_unix_listener_nonexistent_dir() {
        let err = bind_unix_listener(Path::new("/no/such/dir/test.sock")).unwrap_err();
        assert!(err.kind() == io::ErrorKind::NotFound || err.kind() == io::ErrorKind::Other);
    }

    // --- remove_bind_lock_if_unheld ---

    #[tokio::test]
    async fn bind_then_cleanup_removes_bindlock() {
        let tmp = tempfile::tempdir().unwrap();
        let sock = tmp.path().join("test.sock");
        let lock = tmp.path().join("test.sock.bindlock");
        // bind creates the companion lock file
        let listener = bind_unix_listener(&sock).unwrap();
        assert!(lock.exists());
        // cleanup removes it once the listener is gone
        drop(listener);
        let _ = std::fs::remove_file(&sock);
        remove_bind_lock_if_unheld(&sock);
        assert!(!lock.exists());
    }

    #[test]
    fn remove_bind_lock_noop_when_missing() {
        let tmp = tempfile::tempdir().unwrap();
        // Must not create the file or panic when there is nothing to remove.
        let sock = tmp.path().join("never-bound.sock");
        remove_bind_lock_if_unheld(&sock);
        assert!(!tmp.path().join("never-bound.sock.bindlock").exists());
    }

    #[test]
    fn remove_bind_lock_skips_held_lock() {
        let tmp = tempfile::tempdir().unwrap();
        let sock = tmp.path().join("busy.sock");
        let lock_path = bind_lock_path(&sock);
        // Simulate a concurrent binder holding the flock. BSD flock conflicts
        // are per open-file-description, so a second open in the same process
        // genuinely contends.
        let file = std::fs::OpenOptions::new()
            .create(true)
            .truncate(false)
            .write(true)
            .open(&lock_path)
            .unwrap();
        let held =
            nix::fcntl::Flock::lock(file, nix::fcntl::FlockArg::LockExclusiveNonblock).unwrap();
        remove_bind_lock_if_unheld(&sock);
        assert!(lock_path.exists(), "must not yank a lock file a binder holds");
        // Once released, cleanup succeeds.
        drop(held);
        remove_bind_lock_if_unheld(&sock);
        assert!(!lock_path.exists());
    }

    // --- verify_peer_uid ---

    #[tokio::test]
    async fn verify_peer_uid_same_process() {
        let (a, _b) = tokio::net::UnixStream::pair().unwrap();
        verify_peer_uid(&a).unwrap();
    }

    // --- checked_dup ---

    #[test]
    fn checked_dup_stdout() {
        use std::os::fd::AsRawFd;
        let fd = checked_dup(1).unwrap();
        assert!(fd.as_raw_fd() > 2);
    }

    #[test]
    fn checked_dup_invalid_fd() {
        assert!(checked_dup(-1).is_err());
    }

    #[test]
    fn checked_dup_sets_cloexec() {
        use std::os::fd::AsRawFd;
        let fd = checked_dup(1).unwrap();
        let flags = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFD) };
        assert_ne!(flags & libc::FD_CLOEXEC, 0);
    }

    #[test]
    fn set_cloexec_on_plain_dup() {
        let raw = unsafe { libc::dup(1) };
        assert!(raw >= 0);
        let fd = unsafe { OwnedFd::from_raw_fd(raw) };
        use std::os::fd::AsRawFd;
        let before = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFD) };
        assert_eq!(before & libc::FD_CLOEXEC, 0);
        set_cloexec(fd.as_raw_fd()).unwrap();
        let after = unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_GETFD) };
        assert_ne!(after & libc::FD_CLOEXEC, 0);
    }
}
