// Copyright 2026 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef NET_DISK_CACHE_SQL_SQL_SHARED_CACHE_H_
#define NET_DISK_CACHE_SQL_SQL_SHARED_CACHE_H_

#include <atomic>
#include <vector>

#include "base/containers/id_map.h"
#include "base/containers/queue.h"
#include "base/functional/callback.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/raw_ref.h"
#include "base/memory/ref_counted.h"
#include "base/threading/sequence_bound.h"
#include "base/types/pass_key.h"
#include "net/base/net_export.h"
#include "net/base/network_isolation_key.h"
#include "net/disk_cache/sql/cache_entry_key.h"
#include "net/disk_cache/sql/shared_cache_client_remote.h"
#include "net/disk_cache/sql/sql_persistent_store.h"
#include "net/disk_cache/sql/sql_read_cache_memory_monitor.h"
#include "net/disk_cache/sql/sql_shared_cache_blob_handle.h"
#include "net/disk_cache/sql/sql_shared_cache_isolated_database.h"
#include "net/disk_cache/sql/sql_tracked_sequence_bound.h"
#include "third_party/abseil-cpp/absl/container/flat_hash_set.h"

namespace disk_cache {

class BackendCleanupTracker;
class SqlPersistentStore;
class SqlSharedCacheHandle;

// Represents an isolated SQL disk cache instance shared across requests
// matching a specific NetworkIsolationKey (or string representation).
//
// Managed by `SqlSharedCacheManager`, and reference-counted via
// `SqlSharedCacheHandle`. When all handles (`SqlSharedCacheHandle`) referencing
// this cache are destroyed, `SqlSharedCacheManager` cleans up and deletes this
// cache object.
class NET_EXPORT_PRIVATE SqlSharedCache {
 public:
  SqlSharedCache(
      std::string nik_string,
      SqlPersistentStore& store,
      const base::FilePath& directory,
      base::RepeatingCallback<void(SqlSharedCache&)> on_unreferenced_callback,
      scoped_refptr<base::SequencedTaskRunner> db_task_runner,
      scoped_refptr<SqlReadCacheMemoryMonitor> read_cache_memory_monitor,
      scoped_refptr<BackendCleanupTracker> cleanup_tracker);
  ~SqlSharedCache();

  SqlSharedCache(const SqlSharedCache&) = delete;
  SqlSharedCache& operator=(const SqlSharedCache&) = delete;

  // Asynchronously cleans up resources and notifies `callback` when complete.
  void Cleanup(base::OnceClosure callback);

  // Initializes the underlying isolated database instance for this shared cache
  // associated with `shared_cache_db_id`.
  void InitIsolatedDatabase(SqlSharedCacheDbId shared_cache_db_id,
                            base::OnceCallback<void(bool)> callback);

  // Creates a new reference-counted `SqlSharedCacheHandle` targeting this
  // cache.
  scoped_refptr<SqlSharedCacheHandle> CreateHandle();

  // Returns true if there are any active handles referencing this cache.
  bool IsReferenced() const { return handle_count_ != 0; }
  size_t handle_count_for_testing() const { return handle_count_; }

  // Increments/decrements the count of active `SqlSharedCacheHandle` instances.
  // Restricted via `base::PassKey` to `SqlSharedCacheHandle`.
  void IncrementHandleCount(base::PassKey<SqlSharedCacheHandle>);
  void DecrementHandleCount(base::PassKey<SqlSharedCacheHandle>);

  // Returns the string key derived from NetworkIsolationKey identifying this
  // cache.
  const std::string& nik_string() const { return nik_string_; }

  // Returns the database ID associated with this shared cache, if initialized.
  std::optional<SqlSharedCacheDbId> shared_cache_db_id() const {
    return shared_cache_db_id_;
  }

  // Returns the task runner used for DB operations.
  scoped_refptr<base::SequencedTaskRunner> db_task_runner() const {
    return db_task_runner_;
  }

  SqlTrackedSequenceBound<SqlSharedCacheIsolatedDatabase>&
  isolated_database_for_testing() {
    return isolated_database_;
  }

  // Copies multiple shared-cache eligible entries into the shared cache. Must
  // only be called when `entries` is non-empty, no copy operation is currently
  // in progress, and `InitIsolatedDatabase` has completed (`shared_cache_db_id`
  // and `isolated_database_` are set).
  void CopyEntries(
      base::queue<SqlPersistentStore::SharedCacheEligibleEntry> entries,
      scoped_refptr<base::RefCountedData<std::atomic_bool>> abort_flag,
      base::OnceCallback<void(
          base::queue<SqlPersistentStore::SharedCacheEligibleEntry>)> callback,
      base::RepeatingCallback<void(const CacheEntryKey&)>
          on_entry_copied_callback = {});

  // Deletes entries specified by `shared_cache_row_ids` from the isolated
  // database.
  void DeleteEntries(
      const std::vector<SqlSharedCacheRowId>& shared_cache_row_ids,
      base::OnceCallback<
          void(base::expected<void, SqlSharedCacheIsolatedDatabase::Error>)>
          callback);

  // Asynchronously reads entry body data from the shared cache isolated
  // database into `buffer` starting at `offset`. `body_size` is the total body
  // size used to validate the read range.
  void Read(const CacheEntryKey& entry_key,
            SqlSharedCacheRowId shared_cache_row_id,
            int body_size,
            int64_t offset,
            scoped_refptr<net::IOBuffer> buffer,
            SqlPersistentStore::ReadResultOrErrorCallback callback);

  // Asynchronously retrieves a `SqlSharedCacheBlobHandle` for a shared cache
  // entry.
  void GetBlobHandle(
      const CacheEntryKey& entry_key,
      SqlSharedCacheRowId shared_cache_row_id,
      int body_size,
      base::OnceCallback<void(
          base::expected<scoped_refptr<SqlSharedCacheBlobHandle>,
                         SqlSharedCacheIsolatedDatabase::Error>)> callback);

  // Registers a remote client to receive database connection handles.
  void RegisterClient(std::unique_ptr<SharedCacheClientRemote> client);

 private:
  using ClientId = int32_t;
  using ClientsMap =
      base::IDMap<std::unique_ptr<SharedCacheClientRemote>, ClientId>;

  // Entry Copying Call Flow Overview:
  //
  // CopyEntries()
  //       |
  //       v          (empty/abort)
  // CopyNextEntry() --------------> FinishCopy()
  //  ^    |    ^
  //  |    |    +------------------------------------------+
  //  |    |                                               |
  //  |    v                            (error)            |
  //  |  OnEntryOpenedForSharedCache() ----------------> OnCopyEntryFailed()
  //  |    | (body == 0)   | (body > 0)                    ^ ^ ^     ^ ^
  //  |    |               v                      (error)  | | |     | |
  //  |    |          OnEntryDataReadForInsert() ----------+ | |     | |
  //  |    |               | (OK)                            | |     | |
  //  |    v               v           (error)               | |     | |
  //  |  OnIsolatedDatabaseInserted() -----------------------+ |     | |
  //  |                    | (OK)                              |     | |
  //  |                    v                                   |     | |
  //  |               ReadNextChunk()                          |     | |
  //  |               |(done)  ^    | (has data)               |     | |
  //  |               |        |    v                 (error)  |     | |
  //  |               |        |   OnEntryDataRead() ----------+     | |
  //  |               |        |    | (OK)                           | |
  //  |               |    (OK)|    v                         (error)| |
  //  |               |        +- OnIsolatedDatabaseWritten() -------+ |
  //  |               |                                                |
  //  |               v                (error)                         |
  //  |      MoveBlobsToSharedCache() ---------------------------------+
  //  |               |  (OK)
  //  |               v
  //  +----- OnCopyEntryComplete()
  void CopyNextEntry();
  void OnEntryOpenedForSharedCache(
      SqlPersistentStore::SharedCacheEligibleEntry entry,
      SqlPersistentStore::EntryInfoOrError result);
  void OnEntryDataReadForInsert(
      SqlPersistentStore::SharedCacheEligibleEntry entry,
      SqlPersistentStore::ResId res_id,
      scoped_refptr<net::PickledIOBuffer> headers,
      int64_t body_end,
      scoped_refptr<net::IOBuffer> buffer,
      base::expected<SqlPersistentStore::ReadResult, SqlPersistentStore::Error>
          result);
  void OnIsolatedDatabaseInserted(
      CacheEntryKey key,
      SqlPersistentStore::ResId res_id,
      int64_t body_end,
      int64_t offset,
      base::expected<SqlSharedCacheRowId, SqlSharedCacheIsolatedDatabase::Error>
          result);
  void ReadNextChunk(CacheEntryKey key,
                     SqlPersistentStore::ResId res_id,
                     int64_t body_end,
                     int64_t offset,
                     SqlSharedCacheRowId shared_cache_row_id);
  void OnEntryDataRead(CacheEntryKey key,
                       SqlPersistentStore::ResId res_id,
                       int64_t body_end,
                       int64_t offset,
                       SqlSharedCacheRowId shared_cache_row_id,
                       scoped_refptr<net::IOBuffer> buffer,
                       base::expected<SqlPersistentStore::ReadResult,
                                      SqlPersistentStore::Error> result);
  void OnIsolatedDatabaseWritten(
      CacheEntryKey key,
      SqlPersistentStore::ResId res_id,
      int64_t body_end,
      int64_t next_offset,
      SqlSharedCacheRowId shared_cache_row_id,
      base::expected<void, SqlSharedCacheIsolatedDatabase::Error> result);
  void MoveBlobsToSharedCache(CacheEntryKey key,
                              SqlPersistentStore::ResId res_id,
                              SqlSharedCacheRowId shared_cache_row_id);
  void OnCopyEntryComplete(const CacheEntryKey& key);
  void OnCopyEntryFailed();
  void FinishCopy();

  void OnClientDisconnected(ClientId client_id,
                            scoped_refptr<SqlSharedCacheHandle> handle);

  void OnIsolatedDatabaseRead(
      SqlPersistentStore::ReadResultOrErrorCallback callback,
      SqlSharedCacheIsolatedDatabase::ReadResultOrError result);

  void OnPendingFileSetForClient(
      ClientId client_id,
      base::expected<sqlite_vfs::PendingFileSet,
                     SqlSharedCacheIsolatedDatabase::Error> pending_file_set);

  void OnHashesLoaded(
      base::expected<std::vector<uint32_t>,
                     SqlSharedCacheIsolatedDatabase::Error> hashes);

  const std::string nik_string_;
  const raw_ref<SqlPersistentStore> store_;
  const base::FilePath directory_;

  base::RepeatingCallback<void(SqlSharedCache&)> on_unreferenced_callback_;
  int handle_count_ = 0;
  scoped_refptr<base::SequencedTaskRunner> db_task_runner_;
  scoped_refptr<SqlReadCacheMemoryMonitor> read_cache_memory_monitor_;
  scoped_refptr<BackendCleanupTracker> cleanup_tracker_;

  std::optional<SqlSharedCacheDbId> shared_cache_db_id_;

  SqlTrackedSequenceBound<SqlSharedCacheIsolatedDatabase> isolated_database_;

  std::optional<absl::flat_hash_set<uint32_t>> cached_hashes_;
  ClientsMap clients_;

  base::queue<SqlPersistentStore::SharedCacheEligibleEntry>
      pending_copy_entries_;
  scoped_refptr<base::RefCountedData<std::atomic_bool>> copy_abort_flag_;
  base::OnceCallback<void(
      base::queue<SqlPersistentStore::SharedCacheEligibleEntry>)>
      copy_callback_;
  std::optional<SqlSharedCacheRowId> current_copy_row_id_;

  base::RepeatingCallback<void(const CacheEntryKey&)> on_entry_copied_callback_;
  std::optional<uint32_t> current_entry_hash_;
  std::vector<uint32_t> copy_new_hashes_;

  base::WeakPtrFactory<SqlSharedCache> weak_factory_{this};
};

}  // namespace disk_cache

#endif  // NET_DISK_CACHE_SQL_SQL_SHARED_CACHE_H_