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

#include "services/network/shared_dictionary/shared_dictionary_writer_on_disk.h"

#include <limits>

#include "base/functional/callback_helpers.h"
#include "base/metrics/histogram_functions.h"
#include "base/numerics/checked_math.h"
#include "services/network/shared_dictionary/shared_dictionary_constants.h"
#include "services/network/shared_dictionary/shared_dictionary_disk_cache.h"
#include "services/network/shared_dictionary/shared_dictionary_storage_result.h"

namespace network {

SharedDictionaryWriterOnDisk::SharedDictionaryWriterOnDisk(
    const base::UnguessableToken& token,
    FinishCallback callback,
    base::WeakPtr<SharedDictionaryDiskCache> disk_cache)
    : token_(token), callback_(std::move(callback)), disk_cache_(disk_cache) {}

SharedDictionaryWriterOnDisk::~SharedDictionaryWriterOnDisk() {
  if (callback_) {
    OnFailed(Result::kErrorAborted);
  }
}

void SharedDictionaryWriterOnDisk::Initialize() {
  DCHECK_EQ(State::kBeforeInitialize, state_);
  state_ = State::kInitializing;
  DCHECK(disk_cache_);
  // Binding `this` to keep `this` alive until callback will be called.
  auto split_callback = base::SplitOnceCallback(
      base::BindOnce(&SharedDictionaryWriterOnDisk::OnEntry, this));
  disk_cache::EntryResult result = disk_cache_->OpenOrCreateEntry(
      token_.ToString(), /*create=*/true, std::move(split_callback.first));
  if (result.net_error() != net::ERR_IO_PENDING) {
    std::move(split_callback.second).Run(std::move(result));
  }
}

void SharedDictionaryWriterOnDisk::Append(base::span<const uint8_t> data) {
  DCHECK_GT(data.size(), 0u);
  if (state_ == State::kFailed) {
    return;
  }
  base::CheckedNumeric<size_t> checked_total_size = total_size_;
  checked_total_size += data.size();
  if (checked_total_size.ValueOrDefault(std::numeric_limits<size_t>::max()) >
      shared_dictionary::GetDictionarySizeLimit()) {
    OnFailed(Result::kErrorSizeExceedsLimit);
    return;
  }

  total_size_ = checked_total_size.ValueOrDie();
  hash_.Update(data);
  switch (state_) {
    case State::kBeforeInitialize:
      NOTREACHED();
    case State::kInitializing:
      pending_write_buffers_.push_back(
          base::MakeRefCounted<net::VectorIOBuffer>(data));
      break;
    case State::kInitialized: {
      DCHECK(entry_);
      WriteData(base::MakeRefCounted<net::VectorIOBuffer>(data));
    } break;
    case State::kFailed:
      NOTREACHED();
  }
}

void SharedDictionaryWriterOnDisk::Finish() {
  if (state_ == State::kFailed) {
    return;
  }

  finish_called_ = true;
  MaybeFinish();
}

void SharedDictionaryWriterOnDisk::OnEntry(disk_cache::EntryResult result) {
  if (state_ != State::kInitializing) {
    DCHECK_EQ(State::kFailed, state_);
    return;
  }
  if (result.net_error() != net::OK) {
    OnFailed(Result::kErrorCreateEntryFailed);
    return;
  }
  state_ = State::kInitialized;
  entry_.reset(result.ReleaseEntry());

  if (pending_write_buffers_.empty()) {
    MaybeFinish();
    return;
  }

  while (!pending_write_buffers_.empty()) {
    scoped_refptr<net::IOBuffer> buffer = *pending_write_buffers_.begin();
    pending_write_buffers_.pop_front();
    WriteData(std::move(buffer));
  }
}

void SharedDictionaryWriterOnDisk::WriteData(
    scoped_refptr<net::IOBuffer> buffer) {
  DCHECK_NE(State::kBeforeInitialize, state_);
  DCHECK_NE(State::kInitializing, state_);
  if (state_ != State::kInitialized) {
    DCHECK_EQ(State::kFailed, state_);
    return;
  }
  offset_ += buffer->size();
  ++writing_count_;
  // Binding `this` to keep `this` alive until callback will be called.
  auto split_callback = base::SplitOnceCallback(base::BindOnce(
      &SharedDictionaryWriterOnDisk::OnWrittenData, this, buffer->size()));
  // Stores the dictionary binary in the second stream of disk cache entry
  // (index = 1) which was designed to store the HTTP response body of HTTP
  // Cache.
  int result = entry_->WriteData(
      /*index=*/1, /*offset=*/offset_ - buffer->size(), buffer.get(),
      buffer->size(), std::move(split_callback.first), false);
  if (result != net::ERR_IO_PENDING) {
    std::move(split_callback.second).Run(std::move(result));
  }
}

void SharedDictionaryWriterOnDisk::OnWrittenData(int expected_result,
                                                 int result) {
  DCHECK_NE(State::kBeforeInitialize, state_);
  DCHECK_NE(State::kInitializing, state_);
  if (state_ != State::kInitialized) {
    DCHECK_EQ(State::kFailed, state_);
    return;
  }
  if (result != expected_result) {
    OnFailed(Result::kErrorWriteDataFailed);
    return;
  }
  written_size_ += result;
  --writing_count_;
  MaybeFinish();
}

void SharedDictionaryWriterOnDisk::OnFailed(Result result) {
  DCHECK_NE(State::kFailed, state_);
  state_ = State::kFailed;

  SharedDictionaryStorageResult uma_result;
  switch (result) {
    case Result::kErrorCreateEntryFailed:
      uma_result = SharedDictionaryStorageResult::kErrorCreateEntryFailed;
      break;
    case Result::kErrorWriteDataFailed:
      uma_result = SharedDictionaryStorageResult::kErrorWriteDataFailed;
      break;
    case Result::kErrorSizeExceedsLimit:
      uma_result = SharedDictionaryStorageResult::kErrorSizeExceedsLimit;
      break;
    case Result::kErrorSizeZero:
      uma_result = SharedDictionaryStorageResult::kErrorSizeZero;
      break;
    case Result::kErrorAborted:
      uma_result = SharedDictionaryStorageResult::kErrorAborted;
      break;
    case Result::kSuccess:
      NOTREACHED();
  }
  base::UmaHistogramEnumeration("Net.SharedDictionaryOnDisk.StorageResult",
                                uma_result);

  pending_write_buffers_.clear();
  if (entry_) {
    entry_->Doom();
    entry_.reset();
  }
  std::move(callback_).Run(result, 0u, net::SHA256HashValue());
}

void SharedDictionaryWriterOnDisk::MaybeFinish() {
  if ((writing_count_ != 0) || !finish_called_ || !callback_ ||
      !pending_write_buffers_.empty() || (state_ != State::kInitialized)) {
    return;
  }

  if (total_size_ == 0) {
    OnFailed(Result::kErrorSizeZero);
    return;
  }

  entry_.reset();
  DCHECK_EQ(written_size_, total_size_);
  net::SHA256HashValue sha256;
  hash_.Finish(sha256);
  std::move(callback_).Run(Result::kSuccess, total_size_, sha256);
}

}  // namespace network