#pragma once

#include <CoreFoundation/CoreFoundation.h>
#include <DiskArbitration/DiskArbitration.h>
#include <IOKit/IOKitLib.h>
#include <sys/mount.h>
#include <utility> // for std::move() in the move ctor/assignment below

// RAII (Resource Acquisition Is Initialization) utilities for macOS APIs.
// These wrappers ensure proper cleanup of system resources even in the
// presence of exceptions, preventing memory leaks and resource exhaustion.

namespace FSMeta {

// Generic RAII wrapper for resources that need free().
// This is used for C-style allocations that must be freed with free().
// Common usage: buffers returned by system APIs like getmntinfo_r_np().
template <typename T> class ResourceRAII {
private:
  T *resource_;

public:
  ResourceRAII() noexcept : resource_(nullptr) {}
  ~ResourceRAII() noexcept {
    if (resource_) {
      free(resource_);
    }
  }

  T **ptr() { return &resource_; }
  T *get() { return resource_; }

  // Add move operations for better resource management
  ResourceRAII(ResourceRAII &&other) noexcept : resource_(other.resource_) {
    other.resource_ = nullptr;
  }

  ResourceRAII &operator=(ResourceRAII &&other) noexcept {
    if (this != &other) {
      if (resource_)
        free(resource_);
      resource_ = other.resource_;
      other.resource_ = nullptr;
    }
    return *this;
  }

  // Prevent copying
  ResourceRAII(const ResourceRAII &) = delete;
  ResourceRAII &operator=(const ResourceRAII &) = delete;
};

// Specialized RAII wrapper for mount buffer from getmntinfo_r_np().
// getmntinfo_r_np() allocates a buffer that the caller must free.
// This wrapper ensures the buffer is freed even if exceptions occur.
class MountBufferRAII {
private:
  struct statfs *buffer_;

public:
  MountBufferRAII() noexcept : buffer_(nullptr) {}
  ~MountBufferRAII() noexcept {
    if (buffer_) {
      free(buffer_);
    }
  }

  struct statfs **ptr() { return &buffer_; }
  struct statfs *get() { return buffer_; }

  // Add move operations for better resource management
  MountBufferRAII(MountBufferRAII &&other) noexcept : buffer_(other.buffer_) {
    other.buffer_ = nullptr;
  }

  MountBufferRAII &operator=(MountBufferRAII &&other) noexcept {
    if (this != &other) {
      if (buffer_)
        free(buffer_);
      buffer_ = other.buffer_;
      other.buffer_ = nullptr;
    }
    return *this;
  }

  // Prevent copying
  MountBufferRAII(const MountBufferRAII &) = delete;
  MountBufferRAII &operator=(const MountBufferRAII &) = delete;
};

// CoreFoundation RAII wrapper following the Create/Copy/Get rule.
// Any CF object obtained via Create or Copy functions must be released.
// This wrapper automatically calls CFRelease() in the destructor,
// preventing memory leaks from Core Foundation objects.
template <typename T> class CFReleaser {
private:
  T ref_;

public:
  explicit CFReleaser(T ref = nullptr) noexcept : ref_(ref) {}
  ~CFReleaser() noexcept { reset(); }

  void reset(T ref = nullptr) noexcept {
    if (ref_) {
      CFRelease(ref_);
    }
    ref_ = ref;
  }

  operator T() const noexcept { return ref_; }
  T get() const noexcept { return ref_; }
  bool isValid() const noexcept { return ref_ != nullptr; }

  // Prevent copying
  CFReleaser(const CFReleaser &) = delete;
  CFReleaser &operator=(const CFReleaser &) = delete;

  // Allow moving
  CFReleaser(CFReleaser &&other) noexcept : ref_(other.ref_) {
    other.ref_ = nullptr;
  }
  CFReleaser &operator=(CFReleaser &&other) noexcept {
    if (this != &other) {
      reset();
      ref_ = other.ref_;
      other.ref_ = nullptr;
    }
    return *this;
  }
};

// RAII wrapper for IOKit io_object_t handles (io_service_t,
// io_registry_entry_t). IOKit objects must be released with IOObjectRelease(),
// not CFRelease().
class IOObjectGuard {
private:
  io_object_t obj_;

public:
  explicit IOObjectGuard(io_object_t obj = 0) noexcept : obj_(obj) {}
  ~IOObjectGuard() noexcept {
    if (obj_) {
      IOObjectRelease(obj_);
    }
  }

  io_object_t get() const noexcept { return obj_; }
  bool isValid() const noexcept { return obj_ != 0; }

  // Prevent copying
  IOObjectGuard(const IOObjectGuard &) = delete;
  IOObjectGuard &operator=(const IOObjectGuard &) = delete;

  // Allow moving
  IOObjectGuard(IOObjectGuard &&other) noexcept : obj_(other.obj_) {
    other.obj_ = 0;
  }
  IOObjectGuard &operator=(IOObjectGuard &&other) noexcept {
    if (this != &other) {
      if (obj_) {
        IOObjectRelease(obj_);
      }
      obj_ = other.obj_;
      other.obj_ = 0;
    }
    return *this;
  }
};

// Specialized RAII wrapper for DASession that handles dispatch queue lifecycle.
// DASessionSetDispatchQueue must be called with NULL before the session is
// released. This wrapper ensures proper cleanup order: unschedule then release.
class DASessionRAII {
private:
  CFReleaser<DASessionRef> session_;
  bool is_scheduled_;

public:
  explicit DASessionRAII(DASessionRef session = nullptr) noexcept
      : session_(session), is_scheduled_(false) {}

  ~DASessionRAII() noexcept { unschedule(); }

  // Schedule the session on a dispatch queue
  void scheduleOnQueue(dispatch_queue_t queue) {
    if (session_.isValid() && queue != nullptr) {
      DASessionSetDispatchQueue(session_.get(), queue);
      is_scheduled_ = true;
    }
  }

  // Unschedule the session (must be called before session is released)
  void unschedule() noexcept {
    if (is_scheduled_ && session_.isValid()) {
      DASessionSetDispatchQueue(session_.get(), nullptr);
      is_scheduled_ = false;
    }
  }

  DASessionRef get() const noexcept { return session_.get(); }
  bool isValid() const noexcept { return session_.isValid(); }

  // Prevent copying
  DASessionRAII(const DASessionRAII &) = delete;
  DASessionRAII &operator=(const DASessionRAII &) = delete;

  // Allow moving
  DASessionRAII(DASessionRAII &&other) noexcept
      : session_(std::move(other.session_)),
        is_scheduled_(other.is_scheduled_) {
    other.is_scheduled_ = false;
  }

  DASessionRAII &operator=(DASessionRAII &&other) noexcept {
    if (this != &other) {
      unschedule();
      session_ = std::move(other.session_);
      is_scheduled_ = other.is_scheduled_;
      other.is_scheduled_ = false;
    }
    return *this;
  }
};

} // namespace FSMeta
