// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "connections/implementation/bwu_manager.h"

#include <algorithm>
#include <cassert>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/functional/bind_front.h"
#include "absl/strings/str_cat.h"
#include "absl/time/time.h"
#include "connections/implementation/analytics/connection_attempt_metadata_params.h"
#include "connections/implementation/awdl_bwu_handler.h"
#include "connections/implementation/bluetooth_bwu_handler.h"
#include "connections/implementation/bwu_handler.h"
#include "connections/implementation/client_proxy.h"
#include "connections/implementation/endpoint_channel.h"
#include "connections/implementation/endpoint_channel_manager.h"
#include "connections/implementation/endpoint_manager.h"
#include "connections/implementation/flags/nearby_connections_feature_flags.h"
#include "connections/implementation/mediums/mediums.h"
#include "connections/implementation/offline_frames.h"
#include "connections/implementation/service_id_constants.h"
#include "internal/flags/nearby_flags.h"
#ifdef NO_WEBRTC
#include "connections/implementation/webrtc_bwu_handler_stub.h"
#else
#include "connections/implementation/webrtc_bwu_handler.h"
#endif
#include "connections/implementation/wifi_direct_bwu_handler.h"
#include "connections/implementation/wifi_hotspot_bwu_handler.h"
#include "connections/implementation/wifi_lan_bwu_handler.h"
#include "connections/medium_selector.h"
#include "internal/platform/byte_array.h"
#include "internal/platform/cancelable_alarm.h"
#include "internal/platform/count_down_latch.h"
#include "internal/platform/expected.h"
#include "internal/platform/feature_flags.h"
#include "internal/platform/implementation/system_clock.h"
#include "internal/platform/logging.h"
#include "internal/platform/runnable.h"
#include "proto/connections_enums.pb.h"

namespace nearby {
namespace connections {

namespace {
using ::location::nearby::connections::BandwidthUpgradeNegotiationFrame;
using ::location::nearby::connections::MediumRole;
using ::location::nearby::connections::OfflineFrame;
using ::location::nearby::connections::OsInfo;
using ::location::nearby::connections::V1Frame;
using ::location::nearby::proto::connections::BandwidthUpgradeErrorStage;
using ::location::nearby::proto::connections::BandwidthUpgradeResult;
using ::location::nearby::proto::connections::ConnectionAttemptResult;
using ::location::nearby::proto::connections::ConnectionAttemptType;
using ::location::nearby::proto::connections::DisconnectionReason;
using ::location::nearby::proto::connections::OperationResultCode;
}  // namespace

BwuManager::BwuManager(
    Mediums& mediums, EndpointManager& endpoint_manager,
    EndpointChannelManager& channel_manager,
    absl::flat_hash_map<Medium, std::unique_ptr<BwuHandler>> handlers,
    Config config)
    : config_(config),
      mediums_(&mediums),
      endpoint_manager_(&endpoint_manager),
      channel_manager_(&channel_manager) {
  if (config_.bandwidth_upgrade_retry_delay == absl::ZeroDuration()) {
    if (FeatureFlags::GetInstance().GetFlags().use_exp_backoff_in_bwu_retry) {
      config_.bandwidth_upgrade_retry_delay =
          FeatureFlags::GetInstance()
              .GetFlags()
              .bwu_retry_exp_backoff_initial_delay;
    } else {
      config_.bandwidth_upgrade_retry_delay = absl::Seconds(5);
    }
  }
  if (config_.bandwidth_upgrade_retry_max_delay == absl::ZeroDuration()) {
    if (FeatureFlags::GetInstance().GetFlags().use_exp_backoff_in_bwu_retry) {
      config_.bandwidth_upgrade_retry_max_delay =
          FeatureFlags::GetInstance()
              .GetFlags()
              .bwu_retry_exp_backoff_maximum_delay;
    } else {
      config_.bandwidth_upgrade_retry_max_delay = absl::Seconds(10);
    }
  }
  if (config_.allow_upgrade_to.All(false)) {
    config_.allow_upgrade_to.web_rtc = true;
    config_.allow_upgrade_to.wifi_direct = true;
    config_.allow_upgrade_to.wifi_lan = true;
    config_.allow_upgrade_to.wifi_hotspot = true;
    if (NearbyFlags::GetInstance().GetBoolFlag(
            config_package_nearby::nearby_connections_feature::kEnableAwdl)) {
      config_.allow_upgrade_to.awdl = true;
    }
  }
  if (!handlers.empty()) {
    handlers_ = std::move(handlers);
  } else {
    InitBwuHandlers();
  }

  // Register the offline frame processor.
  endpoint_manager_->RegisterFrameProcessor(
      V1Frame::BANDWIDTH_UPGRADE_NEGOTIATION, this);
}

BwuManager::~BwuManager() {
  LOG(INFO) << "BwuManager going down";
  Shutdown();
}

void BwuManager::InitBwuHandlers() {
  // Register the supported concrete BwuMedium implementations.
  if (config_.allow_upgrade_to.awdl) {
    handlers_.emplace(
        Medium::AWDL,
        std::make_unique<AwdlBwuHandler>(
            *mediums_,
            absl::bind_front(&BwuManager::OnIncomingConnection, this)));
  }
  if (config_.allow_upgrade_to.wifi_hotspot) {
    handlers_.emplace(
        Medium::WIFI_HOTSPOT,
        std::make_unique<WifiHotspotBwuHandler>(
            *mediums_,
            absl::bind_front(&BwuManager::OnIncomingConnection, this)));
  }
  if (config_.allow_upgrade_to.wifi_direct) {
    handlers_.emplace(
        Medium::WIFI_DIRECT,
        std::make_unique<WifiDirectBwuHandler>(
            *mediums_,
            absl::bind_front(&BwuManager::OnIncomingConnection, this)));
  }
  if (config_.allow_upgrade_to.wifi_lan) {
    handlers_.emplace(
        Medium::WIFI_LAN,
        std::make_unique<WifiLanBwuHandler>(
            *mediums_,
            absl::bind_front(&BwuManager::OnIncomingConnection, this)));
  }
  if (config_.allow_upgrade_to.web_rtc) {
    handlers_.emplace(
        Medium::WEB_RTC,
        std::make_unique<WebrtcBwuHandler>(
            *mediums_,
            absl::bind_front(&BwuManager::OnIncomingConnection, this)));
  }
  if (config_.allow_upgrade_to.bluetooth) {
    handlers_.emplace(
        Medium::BLUETOOTH,
        std::make_unique<BluetoothBwuHandler>(
            *mediums_,
            absl::bind_front(&BwuManager::OnIncomingConnection, this)));
  }
}

void BwuManager::Shutdown() {
  LOG(INFO) << "Initiating shutdown of BwuManager.";

  endpoint_manager_->UnregisterFrameProcessor(
      V1Frame::BANDWIDTH_UPGRADE_NEGOTIATION, this);

  // Stop all the ongoing Runnables (as gracefully as possible).
  ShutdownExecutors();

  // After worker threads are down we became exclusive owners of data and
  // may access it from current thread.
  for (auto& item : previous_endpoint_channels_) {
    EndpointChannel* channel = item.second.get();
    if (!channel) continue;
    channel->Close(DisconnectionReason::SHUTDOWN);
  }

  CancelAllRetryUpgradeAlarms();
  medium_ = Medium::UNKNOWN_MEDIUM;
  endpoint_id_to_bwu_medium_.clear();
  for (auto& medium_handler_pair : handlers_) {
    assert(medium_handler_pair.second);
    medium_handler_pair.second->RevertInitiatorState();
  }
  handlers_.clear();

  LOG(INFO) << "BwuManager has shut down.";
}

void BwuManager::MakeSingleThreadedForTesting() {
  is_single_threaded_for_testing_ = true;
}

void BwuManager::InvokeOnIncomingConnectionForTesting(
    ClientProxy* client,
    std::unique_ptr<BwuHandler::IncomingSocketConnection> mutable_connection) {
  OnIncomingConnection(client, std::move(mutable_connection));
}

void BwuManager::ShutdownExecutors() {
  alarm_executor_.Shutdown();
  serial_executor_.Shutdown();
}

void BwuManager::InitiateBwuForEndpoint(ClientProxy* client,
                                        const std::string& endpoint_id,
                                        Medium new_medium) {
  // Select the best medium if one is not provided.
  Medium proposed_medium =
      new_medium == Medium::UNKNOWN_MEDIUM
          ? ChooseBestUpgradeMedium(
                endpoint_id,
                client->GetUpgradeMediums(endpoint_id).GetMediums(true))
          : new_medium;

  RunOnBwuManagerThread("bwu-init", [this, client, endpoint_id,
                                     proposed_medium]() {
    LOG(INFO) << "InitiateBwuForEndpoint for endpoint " << endpoint_id
              << " with medium "
              << location::nearby::proto::connections::Medium_Name(
                     proposed_medium);

    if (channel_manager_->isWifiLanConnected() &&
        (proposed_medium == Medium::WIFI_HOTSPOT)) {
      LOG(INFO)
          << "Some endpoint is using WIFI_LAN and proposed upgrade medium is "
             "WIFI_HOTSPOT. Don't do the BWU because STA connecting to "
             "WIFI_HOTSPOT will destroy WIFI_LAN which will lead BWU fail and "
             "other endpoint connection fail";
      return;
    }

    SetBwuMediumForEndpoint(endpoint_id, proposed_medium);
    BwuHandler* handler = GetHandlerForMedium(proposed_medium);
    if (!handler) {
      LOG(ERROR)
          << "BwuManager cannot initiate bandwidth upgrade for endpoint "
          << endpoint_id
          << " because the current BandwidthUpgradeMedium cannot be deduced.";
      return;
    }

    if (in_progress_upgrades_.contains(endpoint_id)) {
      LOG(INFO)
          << "BwuManager is ignoring bandwidth upgrade for endpoint "
          << endpoint_id
          << " because we're already upgrading bandwidth for that endpoint.";

      return;
    }

    CancelRetryUpgradeAlarm(endpoint_id);

    auto channel = channel_manager_->GetChannelForEndpoint(endpoint_id);
    Medium channel_medium =
        channel ? channel->GetMedium() : Medium::UNKNOWN_MEDIUM;

    // Ignore requests where the medium we're upgrading to is the medium we're
    // already connected over. This can happen now that Bluetooth is both an
    // advertising medium and a potential bandwidth upgrade, and will continue
    // to be possible as we add other new advertising mediums like mDNS (WiFi
    // LAN). Very specifically, this happens now when a device uses
    // P2P_CLUSTER, connects over Bluetooth, and is not connected to LAN.
    // Bluetooth is the best medium, and we attempt to upgrade from Bluetooth
    // to Bluetooth.
    if (proposed_medium == channel_medium) {
      LOG(INFO) << "BwuManager ignoring the upgrade for endpoint "
                << endpoint_id
                << " because it is already connected over medium "
                << location::nearby::proto::connections::Medium_Name(
                       proposed_medium);
      return;
    }

    client->GetAnalyticsRecorder().OnBandwidthUpgradeStarted(
        endpoint_id, channel_medium, proposed_medium,
        location::nearby::proto::connections::INCOMING,
        client->GetConnectionToken(endpoint_id));

    if (channel == nullptr) {
      LOG(INFO)
          << "BwuManager couldn't complete the upgrade for endpoint "
          << endpoint_id
          << " because it couldn't find an existing EndpointChannel for it.";
      client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
          endpoint_id, BandwidthUpgradeResult::CHANNEL_ERROR,
          BandwidthUpgradeErrorStage::NETWORK_AVAILABLE,
          OperationResultCode::NEARBY_GENERIC_OLD_ENDPOINT_CHANNEL_NULL);
      return;
    }

    if (is_dynamic_role_switch_enabled_ &&
        client->GetMediumRole(endpoint_id).has_value()) {
      MediumRole medium_role = client->GetMediumRole(endpoint_id).value();
      if (NeedToSwitchRole(client, endpoint_id, proposed_medium, medium_role)) {
        if (!channel
                 ->Write(parser::ForBwuPathRequest(
                     client->GetUpgradeMediums(endpoint_id).GetMediums(true),
                     medium_role))
                 .Ok()) {
          LOG(ERROR) << "BwuManager couldn't complete the upgrade for endpoint "
                     << endpoint_id << " to medium "
                     << location::nearby::proto::connections::Medium_Name(
                            proposed_medium)
                     << " because it failed to write the "
                        "BWU_NEGOTIATION.UPGRADE_PATH_REQUEST OfflineFrame.";

          client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
              endpoint_id, BandwidthUpgradeResult::RESULT_IO_ERROR,
              BandwidthUpgradeErrorStage::NETWORK_AVAILABLE,
              OperationResultCode::
                  CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);
          return;
        }
        LOG(INFO) << "BwuManager successfully wrote the "
                     "BANDWIDTH_UPGRADE_NEGOTIATION.UPGRADE_PATH_REQUEST "
                     "OfflineFrame while upgrading endpoint "
                  << endpoint_id << " to medium "
                  << location::nearby::proto::connections::Medium_Name(
                         proposed_medium);
        return;
      }
    }

    std::string service_id = channel->GetServiceId();
    ByteArray bytes = handler->InitializeUpgradedMediumForEndpoint(
        client, service_id, endpoint_id);

    // Because we grab the endpointChannel first thing, it is possible the
    // endpointChannel is stale by the time we attempt to write over it.
    if (bytes.Empty()) {
      LOG(ERROR) << "BwuManager couldn't complete the upgrade for endpoint "
                 << endpoint_id << " to medium "
                 << location::nearby::proto::connections::Medium_Name(
                        proposed_medium)
                 << " because it failed to initialize the "
                    "BWU_NEGOTIATION.UPGRADE_PATH_AVAILABLE OfflineFrame.";
      UpgradePathInfo info;
      info.set_medium(parser::MediumToUpgradePathInfoMedium(proposed_medium));

      ProcessUpgradeFailureEvent(
          client, endpoint_id, info, BandwidthUpgradeResult::MEDIUM_ERROR,
          /*record_analytic=*/false,
          OperationResultCode::CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);

      client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
          endpoint_id, BandwidthUpgradeResult::MEDIUM_ERROR,
          BandwidthUpgradeErrorStage::NETWORK_AVAILABLE,
          OperationResultCode::CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);
      return;
    }
    if (!channel->Write(bytes).Ok()) {
      LOG(ERROR) << "BwuManager couldn't complete the upgrade for endpoint "
                 << endpoint_id << " to medium "
                 << location::nearby::proto::connections::Medium_Name(
                        proposed_medium)
                 << " because it failed to write the "
                    "BWU_NEGOTIATION.UPGRADE_PATH_AVAILABLE OfflineFrame.";
      UpgradePathInfo info;
      info.set_medium(parser::MediumToUpgradePathInfoMedium(proposed_medium));

      ProcessUpgradeFailureEvent(
          client, endpoint_id, info, BandwidthUpgradeResult::RESULT_IO_ERROR,
          /*record_analytic=*/false,
          OperationResultCode::CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);

      client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
          endpoint_id, BandwidthUpgradeResult::RESULT_IO_ERROR,
          BandwidthUpgradeErrorStage::NETWORK_AVAILABLE,
          OperationResultCode::CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);
      return;
    }

    LOG(INFO) << "BwuManager successfully wrote the "
                 "BWU_NEGOTIATION.UPGRADE_PATH_AVAILABLE OfflineFrame while "
                 "upgrading endpoint "
              << endpoint_id << " to medium "
              << location::nearby::proto::connections::Medium_Name(
                     proposed_medium);
    in_progress_upgrades_.emplace(endpoint_id, client);
  });
}

void BwuManager::OnIncomingFrame(OfflineFrame& frame,
                                 const std::string& endpoint_id,
                                 ClientProxy* client, Medium medium,
                                 PacketMetaData& packet_meta_data) {
  V1Frame::FrameType frame_type = parser::GetFrameType(frame);
  if (frame_type != V1Frame::BANDWIDTH_UPGRADE_NEGOTIATION) return;

  auto bwu_frame = frame.v1().bandwidth_upgrade_negotiation();
  LOG(INFO) << "OnIncomingFrame: bwu_frame="
            << BandwidthUpgradeNegotiationFrame::EventType_Name(
                   bwu_frame.event_type())
            << ", endpoint_id=" << endpoint_id << ", medium="
            << location::nearby::proto::connections::Medium_Name(medium);
  if (FeatureFlags::GetInstance().GetFlags().enable_async_bandwidth_upgrade) {
    RunOnBwuManagerThread(
        "bwu-on-incoming-frame", [this, client, endpoint_id, bwu_frame]() {
          OnBwuNegotiationFrame(client, bwu_frame, endpoint_id);
        });
  } else {
    CountDownLatch latch(1);
    RunOnBwuManagerThread("bwu-on-incoming-frame", [this, client, endpoint_id,
                                                    bwu_frame, &latch]() {
      OnBwuNegotiationFrame(client, bwu_frame, endpoint_id);
      latch.CountDown();
    });
    latch.Await();
  }
}

void BwuManager::OnEndpointDisconnect(ClientProxy* client,
                                      const std::string& service_id,
                                      const std::string& endpoint_id,
                                      CountDownLatch barrier,
                                      DisconnectionReason reason) {
  LOG(INFO) << "BwuManager has processed endpoint disconnection for endpoint "
            << endpoint_id << " with reason "
            << DisconnectionReason_Name(reason);
  RunOnBwuManagerThread("bwu-on-endpoint-disconnect", [this, client, service_id,
                                                       endpoint_id,
                                                       barrier]() mutable {
    Medium medium = GetBwuMediumForEndpoint(endpoint_id);
    BwuHandler* handler = GetHandlerForMedium(medium);
    if (handler) {
      handler->OnEndpointDisconnect(client, endpoint_id);
    }

    auto item = previous_endpoint_channels_.extract(endpoint_id);
    if (!item.empty()) {
      auto old_channel = item.mapped();
      if (old_channel != nullptr) {
        old_channel->Close(DisconnectionReason::SHUTDOWN);
      }
    }
    in_progress_upgrades_.erase(endpoint_id);
    retry_delays_.erase(endpoint_id);
    CancelRetryUpgradeAlarm(endpoint_id);
    successfully_upgraded_endpoints_.erase(endpoint_id);

    // Note(nohle): I'm skeptical of the "<= 1", which seems like it should be
    // "== 0". Luckily, we will enable the flag by default, and it won't matter.
    if (FeatureFlags::GetInstance().GetFlags().support_multiple_bwu_mediums ||
        channel_manager_->GetConnectedEndpointsCount() <= 1) {
      RevertBwuMediumForEndpoint(service_id, endpoint_id);
    }
    barrier.CountDown();
  });
}

void BwuManager::RevertBwuMediumForEndpoint(const std::string& service_id,
                                            const std::string& endpoint_id) {
  Medium medium = GetBwuMediumForEndpoint(endpoint_id);

  // If |support_multiple_bwu_mediums| is disabled, we take a less fine-grained
  // approach and revert the handler for _all_ endpoints.
  if (!FeatureFlags::GetInstance().GetFlags().support_multiple_bwu_mediums) {
    LOG(INFO) << "Reverting medium "
              << location::nearby::proto::connections::Medium_Name(medium)
              << " for all endpoints for service " << service_id;
    medium_ = Medium::UNKNOWN_MEDIUM;
    BwuHandler* handler = GetHandlerForMedium(medium);
    if (!handler) return;

    handler->RevertInitiatorState();
    return;
  }

  LOG(INFO) << "Reverting medium "
            << location::nearby::proto::connections::Medium_Name(medium)
            << " for service ID " << service_id << " and endpoint "
            << endpoint_id;
  endpoint_id_to_bwu_medium_.erase(endpoint_id);

  BwuHandler* handler = GetHandlerForMedium(medium);
  if (!handler) {
    LOG(INFO) << "No BWU handler can be found for "
              << location::nearby::proto::connections::Medium_Name(medium);
    return;
  }
  // If |service_id| isn't of the INITIATOR-upgrade format--for example, if this
  // is called by the RESPONDER--there is no need to call RevertInitiatorState
  // unless the BWU Medium is Hotspot. The client needs to disconnect from
  // Hotspot, then it can restore the previous AP connection right away.
  if (!IsInitiatorUpgradeServiceId(service_id)) {
    if (medium == Medium::WIFI_HOTSPOT || medium == Medium::WIFI_DIRECT) {
      handler->RevertResponderState(service_id);
    }
    return;
  }

  handler->RevertInitiatorState(service_id, endpoint_id);
}

bool BwuManager::IsUpgradeOngoing(const std::string& endpoint_id) {
  CountDownLatch latch(1);
  RunOnBwuManagerThread("is_upgrade_ongoing",
                        [&latch]() { latch.CountDown(); });
  latch.Await();
  return in_progress_upgrades_.contains(endpoint_id);
}

Medium BwuManager::GetBwuMediumForEndpoint(
    const std::string& endpoint_id) const {
  if (!FeatureFlags::GetInstance().GetFlags().support_multiple_bwu_mediums) {
    return medium_;
  }

  auto it = endpoint_id_to_bwu_medium_.find(endpoint_id);
  return it == endpoint_id_to_bwu_medium_.end() ? Medium::UNKNOWN_MEDIUM
                                                : it->second;
}
void BwuManager::SetBwuMediumForEndpoint(const std::string& endpoint_id,
                                         Medium medium) {
  if (!FeatureFlags::GetInstance().GetFlags().support_multiple_bwu_mediums) {
    medium_ = medium;
    return;
  }

  endpoint_id_to_bwu_medium_[endpoint_id] = medium;
}

BwuHandler* BwuManager::GetHandlerForMedium(Medium medium) const {
  if (medium == Medium::UNKNOWN_MEDIUM) return nullptr;

  auto it = handlers_.find(medium);
  if (it == handlers_.end()) return nullptr;

  return it->second.get();
}

void BwuManager::OnBwuNegotiationFrame(ClientProxy* client,
                                       const BwuNegotiationFrame frame,
                                       const std::string& endpoint_id) {
  LOG(INFO) << "OnBwuNegotiationFrame: processing incoming "
            << BwuNegotiationFrame::EventType_Name(frame.event_type())
            << " frame for endpoint " << endpoint_id;

  if (!client->IsConnectedToEndpoint(endpoint_id)) {
    LOG(WARNING)
        << "BwuManager skips the process BANDWIDTH_UPGRADE_NEGOTIATION before "
           "PCP connected, "
        << frame.event_type();

    // For the case discover side not yet get the local accept from client, but
    // advertise side already get, discover side should inform advertise side
    // the upgrade failed, so the advertise side could have chance to initialize
    // another upgrade flow again.
    if (frame.event_type() == BwuNegotiationFrame::UPGRADE_PATH_AVAILABLE) {
      RunUpgradeFailedProtocol(client, endpoint_id, frame.upgrade_path_info());
    }
    return;
  }

  switch (frame.event_type()) {
    case BwuNegotiationFrame::UPGRADE_PATH_AVAILABLE:
      ProcessBwuPathAvailableEvent(client, endpoint_id,
                                   frame.upgrade_path_info());
      break;
    case BwuNegotiationFrame::UPGRADE_FAILURE:
      ProcessUpgradeFailureEvent(
          client, endpoint_id, frame.upgrade_path_info(),
          BandwidthUpgradeResult::REMOTE_CONNECTION_ERROR,
          /* record_analytic= */ true,
          OperationResultCode::NEARBY_GENERIC_REMOTE_UPGRADE_FAILURE);
      break;
    case BwuNegotiationFrame::LAST_WRITE_TO_PRIOR_CHANNEL:
      ProcessLastWriteToPriorChannelEvent(client, endpoint_id);
      break;
    case BwuNegotiationFrame::SAFE_TO_CLOSE_PRIOR_CHANNEL:
      ProcessSafeToClosePriorChannelEvent(client, endpoint_id);
      break;
    default:
      LOG(WARNING)
          << "BwuManager can't process unknown incoming OfflineFrame of type "
          << BandwidthUpgradeNegotiationFrame::EventType_Name(
                 frame.event_type())
          << ", ignoring it.";
      break;
  }
}

void BwuManager::OnIncomingConnection(
    ClientProxy* client,
    std::unique_ptr<BwuHandler::IncomingSocketConnection> mutable_connection) {
  LOG(INFO) << "BwuManager process incoming connection";
  std::shared_ptr<BwuHandler::IncomingSocketConnection> connection(
      mutable_connection.release());
  RunOnBwuManagerThread("bwu-on-incoming-connection", [this, client,
                                                       connection]() {
    absl::Time connection_attempt_start_time = SystemClock::ElapsedRealtime();
    EndpointChannel* channel = connection->channel.get();
    if (channel == nullptr) {
      LOG(ERROR)
          << "BwuManager failed to create new EndpointChannel for incoming "
             "socket.";
      connection->socket->Close();
      AttemptToRecordBandwidthUpgradeErrorForUnknownEndpoint(
          location::nearby::proto::connections::MEDIUM_ERROR,
          BandwidthUpgradeErrorStage::SOCKET_CREATION,
          OperationResultCode::NEARBY_GENERIC_NEW_ENDPOINT_CHANNEL_NULL);
      return;
    }

    VLOG(1) << "BwuManager successfully created new EndpointChannel for "
               "incoming socket";

    ClientIntroduction introduction;
    if (!ReadClientIntroductionFrame(channel, introduction)) {
      // This was never a fully EstablishedConnection, no need to provide a
      // closure reason.
      channel->Close();
      LOG(ERROR) << "BwuManager failed to read "
                    "BWU_NEGOTIATION.CLIENT_INTRODUCTION OfflineFrame from "
                    "newly-created EndpointChannel "
                 << channel->GetName()
                 << ", so the EndpointChannel was discarded.";
      return;
    }

    VLOG(1) << "BwuManager successfully received "
               "BWU_NEGOTIATION.CLIENT_INTRODUCTION "
               "OfflineFrame on EndpointChannel "
            << channel->GetName();

    if (!WriteClientIntroductionAckFrame(channel)) {
      // This was never a fully EstablishedConnection, no need to provide a
      // closure reason.
      LOG(ERROR) << "BwuManager failed to write"
                    "BWU_NEGOTIATION.CLIENT_INTRODUCTION_ACK "
                    "OfflineFrame on EndpointChannel "
                 << channel->GetName();
      channel->Close();
      return;
    }

    VLOG(1) << "BwuManager successfully wrote "
               "BWU_NEGOTIATION.CLIENT_INTRODUCTION_ACK "
               "OfflineFrame on EndpointChannel "
            << channel->GetName();

    const std::string& endpoint_id = introduction.endpoint_id();
    ClientProxy* mapped_client;
    const auto item = in_progress_upgrades_.find(endpoint_id);
    if (item == in_progress_upgrades_.end()) return;
    mapped_client = item->second;
    CancelRetryUpgradeAlarm(endpoint_id);
    if (mapped_client == nullptr) {
      // This was never a fully EstablishedConnection, no need to provide a
      // closure reason.
      channel->Close();
      return;
    }

    CHECK(client == mapped_client);

    // The ConnectionAttempt has now succeeded, so record it as such.
    std::unique_ptr<ConnectionAttemptMetadataParams>
        connections_attempt_metadata_params;
    if (channel != nullptr) {
      connections_attempt_metadata_params =
          client->GetAnalyticsRecorder().BuildConnectionAttemptMetadataParams(
              channel->GetTechnology(), channel->GetBand(),
              channel->GetFrequency(), channel->GetTryCount());
      connections_attempt_metadata_params->operation_result_code =
          OperationResultCode::DETAIL_SUCCESS;
    }
    client->GetAnalyticsRecorder().OnIncomingConnectionAttempt(
        location::nearby::proto::connections::UPGRADE, channel->GetMedium(),
        location::nearby::proto::connections::RESULT_SUCCESS,
        SystemClock::ElapsedRealtime() - connection_attempt_start_time,
        client->GetConnectionToken(endpoint_id),
        connections_attempt_metadata_params.get());

    // Use the introductory client information sent over to run the upgrade
    // protocol.
    RunUpgradeProtocol(mapped_client, endpoint_id,
                       std::move(connection->channel),
                       !introduction.supports_disabling_encryption());
  });
}

void BwuManager::RunOnBwuManagerThread(const std::string& name,
                                       Runnable runnable) {
  if (is_single_threaded_for_testing_) {
    runnable();
    return;
  }

  serial_executor_.Execute(name, std::move(runnable));
}

void BwuManager::RunUpgradeProtocol(
    ClientProxy* client, const std::string& endpoint_id,
    std::unique_ptr<EndpointChannel> new_channel, bool enable_encryption) {
  LOG(INFO) << "RunUpgradeProtocol new channel @" << new_channel.get()
            << " name: " << new_channel->GetName() << ", medium: "
            << location::nearby::proto::connections::Medium_Name(
                   new_channel->GetMedium());
  // First, register this new EndpointChannel as *the* EndpointChannel to use
  // for this endpoint here onwards. NOTE: We pause this new EndpointChannel
  // until we've completely drained the old EndpointChannel to avoid out of
  // order reads on the other side. This is a consequence of using the same
  // UKEY2 context for both the previous and new EndpointChannels. UKEY2 uses
  // sequence numbers for writes and reads, and simultaneously sending Payloads
  // on the new channel and control messages on the old channel cause the other
  // side to read messages out of sequence
  new_channel->Pause();
  auto old_channel = channel_manager_->GetChannelForEndpoint(endpoint_id);
  if (!old_channel) {
    LOG(INFO)
        << "BwuManager didn't find a previous EndpointChannel for "
        << endpoint_id
        << " when registering the new EndpointChannel, short-circuiting the "
           "upgrade protocol.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::CHANNEL_ERROR,
        BandwidthUpgradeErrorStage::PRIOR_ENDPOINT_CHANNEL,
        OperationResultCode::NEARBY_GENERIC_OLD_ENDPOINT_CHANNEL_NULL);
    return;
  }
  channel_manager_->ReplaceChannelForEndpoint(
      client, endpoint_id, std::move(new_channel), enable_encryption);

  // Next, initiate a clean shutdown for the previous EndpointChannel used for
  // this endpoint by telling the remote device that it will not receive any
  // more writes over that EndpointChannel.
  if (!old_channel->Write(parser::ForBwuLastWrite()).Ok()) {
    LOG(ERROR) << "BwuManager failed to write "
                  "BWU_NEGOTIATION.LAST_WRITE_TO_PRIOR_CHANNEL OfflineFrame to "
                  "endpoint "
               << endpoint_id << ", short-circuiting the upgrade protocol.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::RESULT_IO_ERROR,
        BandwidthUpgradeErrorStage::LAST_WRITE_TO_PRIOR_CHANNEL,
        OperationResultCode::CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);
    return;
  }
  VLOG(1) << "BwuManager successfully wrote "
             "BWU_NEGOTIATION.LAST_WRITE_TO_PRIOR_CHANNEL "
             "OfflineFrame while upgrading endpoint "
          << endpoint_id;

  // The remainder of this clean shutdown for the previous EndpointChannel will
  // continue when we receive a corresponding
  // BANDWIDTH_UPGRADE_NEGOTIATION.LAST_WRITE_TO_PRIOR_CHANNEL OfflineFrame from
  // the remote device, so for now, just store that previous EndpointChannel.
  previous_endpoint_channels_.emplace(endpoint_id, old_channel);

  // If we already read LAST_WRITE on the old endpoint channel, then we can
  // safely close it now.
  auto item = successfully_upgraded_endpoints_.extract(endpoint_id);
  if (!item.empty()) {
    ProcessLastWriteToPriorChannelEvent(client, endpoint_id);
  }
}

// Outgoing BWU session.
void BwuManager::ProcessBwuPathAvailableEvent(
    ClientProxy* client, const std::string& endpoint_id,
    const UpgradePathInfo& upgrade_path_info) {
  Medium upgrade_medium =
      parser::UpgradePathInfoMediumToMedium(upgrade_path_info.medium());
  LOG(INFO) << "ProcessBwuPathAvailableEvent for endpoint " << endpoint_id
            << " medium "
            << location::nearby::proto::connections::Medium_Name(
                   upgrade_medium);

  if (channel_manager_->isWifiLanConnected() &&
      ((upgrade_medium == Medium::WIFI_HOTSPOT) ||
       ((upgrade_medium == Medium::WIFI_DIRECT) &&
        (client->GetLocalOsInfo().type() ==
         location::nearby::connections::OsInfo::WINDOWS)))) {
    LOG(INFO)
        << "Some endpoint is using WIFI_LAN and proposed upgrade medium is "
        << location::nearby::proto::connections::Medium_Name(upgrade_medium)
        << ". Don't do the BWU because this will destroy WIFI_LAN which will "
           "lead BWU fail and other endpoint connection fail";
    RunUpgradeFailedProtocol(client, endpoint_id, upgrade_path_info);
    return;
  }

  bool abort_bwu = false;
  if (client->IsIncomingConnection(endpoint_id)) {
    if (!is_dynamic_role_switch_enabled_) {
      abort_bwu = true;
    } else {
      auto medium_role = client->GetMediumRole(endpoint_id);
      if (medium_role.has_value() &&
          !NeedToSwitchRole(client, endpoint_id, upgrade_medium,
                            medium_role.value())) {
        abort_bwu = true;
      }
    }
  }
  if (abort_bwu) {
    LOG(INFO)
        << "ProcessBandwidthUpgradePathAvailableEvent ignored by Advertiser";
    return;
  }

  if (in_progress_upgrades_.contains(endpoint_id)) {
    LOG(ERROR)
        << "BwuManager received a duplicate bandwidth upgrade for endpoint "
        << endpoint_id
        << ". We're out of sync with the remote device and cannot recover; "
           "closing all channels.";

    auto item = previous_endpoint_channels_.extract(endpoint_id);
    if (!item.empty()) {
      std::shared_ptr<EndpointChannel> previous_endpoint_channel =
          item.mapped();
      if (previous_endpoint_channel) {
        previous_endpoint_channel->Close(DisconnectionReason::UNFINISHED);
      }
    }
    std::shared_ptr<EndpointChannel> new_channel =
        channel_manager_->GetChannelForEndpoint(endpoint_id);
    if (new_channel) {
      // The upgraded channel never finished upgrading, and therefore is still
      // paused.
      new_channel->Resume();
      new_channel->Close(DisconnectionReason::UNFINISHED);
    }

    return;
  }
  Medium current_bwu_medium = GetBwuMediumForEndpoint(endpoint_id);
  if (current_bwu_medium == Medium::UNKNOWN_MEDIUM) {
    SetBwuMediumForEndpoint(endpoint_id, upgrade_medium);
  }
  // Check for the correct medium so we don't process an incorrect OfflineFrame.
  if (upgrade_medium != GetBwuMediumForEndpoint(endpoint_id)) {
    LOG(INFO) << "Medium not matching";
    RunUpgradeFailedProtocol(client, endpoint_id, upgrade_path_info);
    return;
  }

  auto current_channel = channel_manager_->GetChannelForEndpoint(endpoint_id);
  Medium current_medium =
      current_channel ? current_channel->GetMedium() : Medium::UNKNOWN_MEDIUM;
  client->GetAnalyticsRecorder().OnBandwidthUpgradeStarted(
      endpoint_id, current_medium, upgrade_medium,
      location::nearby::proto::connections::OUTGOING,
      client->GetConnectionToken(endpoint_id));

  absl::Time connection_attempt_start_time = SystemClock::ElapsedRealtime();
  ErrorOr<std::unique_ptr<EndpointChannel>> result =
      ProcessBwuPathAvailableEventInternal(client, endpoint_id,
                                           upgrade_path_info);
  std::unique_ptr<EndpointChannel> channel =
      result.has_value() ? std::move(result.value()) : nullptr;
  ConnectionAttemptResult connection_attempt_result;
  if (channel != nullptr) {
    connection_attempt_result =
        location::nearby::proto::connections::RESULT_SUCCESS;
  } else if (client->GetCancellationFlag(endpoint_id)->Cancelled()) {
    connection_attempt_result =
        location::nearby::proto::connections::RESULT_CANCELLED;
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::RESULT_REMOTE_ERROR,
        BandwidthUpgradeErrorStage::UPGRADE_CANCEL,
        OperationResultCode::CLIENT_CANCELLATION_UPGRADE_CANCELED_BY_REMOTE);
  } else {
    connection_attempt_result =
        location::nearby::proto::connections::RESULT_ERROR;
  }

  if (channel != nullptr) {
    std::unique_ptr<ConnectionAttemptMetadataParams>
        connections_attempt_metadata_params =
            client->GetAnalyticsRecorder().BuildConnectionAttemptMetadataParams(
                channel->GetTechnology(), channel->GetBand(),
                channel->GetFrequency(), channel->GetTryCount());
    connections_attempt_metadata_params->operation_result_code =
        OperationResultCode::DETAIL_SUCCESS;
    client->GetAnalyticsRecorder().OnOutgoingConnectionAttempt(
        endpoint_id, ConnectionAttemptType::UPGRADE, upgrade_medium,
        connection_attempt_result,
        SystemClock::ElapsedRealtime() - connection_attempt_start_time,
        client->GetConnectionToken(endpoint_id),
        connections_attempt_metadata_params.get());
  } else {
    auto connections_attempt_metadata_params =
        std::make_unique<ConnectionAttemptMetadataParams>();
    connections_attempt_metadata_params->operation_result_code =
        result.error().operation_result_code().value();
    client->GetAnalyticsRecorder().OnOutgoingConnectionAttempt(
        endpoint_id, ConnectionAttemptType::UPGRADE, upgrade_medium,
        connection_attempt_result,
        SystemClock::ElapsedRealtime() - connection_attempt_start_time,
        client->GetConnectionToken(endpoint_id),
        connections_attempt_metadata_params.get());
  }

  if (channel == nullptr) {
    LOG(INFO) << "Failed to get new channel.";
    RunUpgradeFailedProtocol(client, endpoint_id, upgrade_path_info);
    return;
  }

  in_progress_upgrades_.emplace(endpoint_id, client);
  RunUpgradeProtocol(client, endpoint_id, std::move(channel),
                     !upgrade_path_info.supports_disabling_encryption());
}

ErrorOr<std::unique_ptr<EndpointChannel>>
BwuManager::ProcessBwuPathAvailableEventInternal(
    ClientProxy* client, const std::string& endpoint_id,
    const UpgradePathInfo& upgrade_path_info) {
  Medium medium =
      parser::UpgradePathInfoMediumToMedium(upgrade_path_info.medium());
  if (medium != GetBwuMediumForEndpoint(endpoint_id)) {
    LOG(ERROR) << "ProcessBwuPathAvailableEventInternal failed for endpoint "
               << endpoint_id << " medium "
               << location::nearby::proto::connections::Medium_Name(medium)
               << ". Upgrade medium not yet set for endpoint.";
    return {Error(OperationResultCode::NEARBY_UPGRADE_PATH_ON_WRONG_MEDIUM)};
  }

  BwuHandler* handler = GetHandlerForMedium(medium);
  if (!handler) {
    LOG(ERROR) << "ProcessBwuPathAvailableEventInternal failed for endpoint "
               << endpoint_id << " medium "
               << location::nearby::proto::connections::Medium_Name(medium)
               << ". No handler for medium.";
    return {Error(OperationResultCode::NEARBY_UPGRADE_PATH_ON_WRONG_MEDIUM)};
  }

  LOG(INFO) << "ProcessBwuPathAvailableEventInternal for "
               "endpoint "
            << endpoint_id << " medium "
            << location::nearby::proto::connections::Medium_Name(medium);

  // Get service ID from the old channel. Don't keep the old channel's shared
  // pointer in scope longer than necessary.
  std::string service_id;
  Medium old_medium = Medium::UNKNOWN_MEDIUM;
  bool disable_ble_scanning = false;
  {
    std::shared_ptr<EndpointChannel> old_channel =
        channel_manager_->GetChannelForEndpoint(endpoint_id);
    if (!old_channel) {
      LOG(ERROR) << "ProcessBwuPathAvailableEventInternal failed for endpoint "
                 << endpoint_id << " medium "
                 << location::nearby::proto::connections::Medium_Name(medium)
                 << ". Old endpoint channel is missing.";
      return {
          Error(OperationResultCode::NEARBY_GENERIC_OLD_ENDPOINT_CHANNEL_NULL)};
    }
    service_id = old_channel->GetServiceId();
    old_medium = old_channel->GetMedium();
  }

  bool enable_ble_v2 = NearbyFlags::GetInstance().GetBoolFlag(
      config_package_nearby::nearby_connections_feature::kEnableBleV2);
  if (NearbyFlags::GetInstance().GetBoolFlag(
          config_package_nearby::nearby_connections_feature::
              kEnableStopBleScanningOnWifiUpgrade)) {
    if (client->GetLocalOsInfo().type() ==
            location::nearby::connections::OsInfo::APPLE &&
        old_medium == Medium::BLE && medium == Medium::WIFI_HOTSPOT) {
      disable_ble_scanning = true;
      if (enable_ble_v2) {
        LOG(INFO) << "For Apple OS, if upgrade from BLE_V2 to WIFI_HOTSPOT, "
                     "we need to pause "
                     "BLE_V2 scanning because it can interfere with WIFI "
                     "Hotspot scanning and connection.";
        ble_v2_medium_.PauseMediumScanning();
      }
    }
  }

  ErrorOr<std::unique_ptr<EndpointChannel>> result =
      handler->CreateUpgradedEndpointChannel(client, service_id, endpoint_id,
                                             upgrade_path_info);

  if (NearbyFlags::GetInstance().GetBoolFlag(
          config_package_nearby::nearby_connections_feature::
              kEnableStopBleScanningOnWifiUpgrade)) {
    if (disable_ble_scanning) {
      if (enable_ble_v2) {
        LOG(INFO) << "Resume BLE_V2 scanning.";
        ble_v2_medium_.ResumeMediumScanning();
      }
    }
  }

  if (result.has_error() || !result.has_value()) {
    LOG(ERROR) << "BwuManager failed to create an endpoint "
                  "channel to endpoint"
               << endpoint_id << ", aborting upgrade.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::RESULT_IO_ERROR,
        BandwidthUpgradeErrorStage::SOCKET_CREATION,
        result.error().operation_result_code().value());
    return result;
  }
  std::unique_ptr<EndpointChannel> new_channel = std::move(result.value());

  // Write the requisite BANDWIDTH_UPGRADE_NEGOTIATION.CLIENT_INTRODUCTION as
  // the first OfflineFrame on this new EndpointChannel.
  if (!new_channel
           ->Write(parser::ForBwuIntroduction(
               client->GetLocalEndpointId(),
               upgrade_path_info.supports_disabling_encryption()))
           .Ok()) {
    // This was never a fully EstablishedConnection, no need to provide a
    // closure reason.
    new_channel->Close();

    LOG(ERROR)
        << "BwuManager failed to write BWU_NEGOTIATION.CLIENT_INTRODUCTION "
           "OfflineFrame to newly-created EndpointChannel "
        << new_channel->GetName() << ", aborting upgrade.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::RESULT_IO_ERROR,
        BandwidthUpgradeErrorStage::CLIENT_INTRODUCTION,
        OperationResultCode::
            NEARBY_GENERIC_READ_CLIENT_INTRODUCTION_ACK_FORMAT_ERROR);
    return {
        Error(OperationResultCode::
                  NEARBY_GENERIC_READ_CLIENT_INTRODUCTION_ACK_FORMAT_ERROR)};
  }

  if (upgrade_path_info.supports_client_introduction_ack()) {
    if (!ReadClientIntroductionAckFrame(new_channel.get())) {
      // This was never a fully EstablishedConnection, no need to provide a
      // closure reason.
      new_channel->Close();

      LOG(ERROR) << "BwuManager failed to read "
                    "BWU_NEGOTIATION.CLIENT_INTRODUCTION_ACK "
                    "OfflineFrame to newly-created EndpointChannel "
                 << new_channel->GetName() << ", aborting upgrade.";
      return {Error(
          OperationResultCode::
              NEARBY_GENERIC_READ_CLIENT_INTRODUCTION_ACK_FRAME_TYPE_ERROR)};
    }
  }

  LOG(INFO) << "BwuManager successfully wrote "
               "BWU_NEGOTIATION.CLIENT_INTRODUCTION OfflineFrame to "
               "newly-created EndpointChannel "
            << new_channel->GetName() << " while upgrading endpoint "
            << endpoint_id;

  // Set the AnalyticsRecorder so that the future closure of this
  // EndpointChannel will be recorded.
  return {std::move(new_channel)};
}

void BwuManager::RunUpgradeFailedProtocol(
    ClientProxy* client, const std::string& endpoint_id,
    const UpgradePathInfo& upgrade_path_info) {
  LOG(INFO) << "RunUpgradeFailedProtocol for endpoint " << endpoint_id
            << " medium "
            << location::nearby::proto::connections::Medium_Name(
                   parser::UpgradePathInfoMediumToMedium(
                       upgrade_path_info.medium()));
  // We attempted to connect to the new medium that the remote device has set up
  // for us but we failed. We need to let the remote device know so that they
  // can pick another medium for us to try.
  std::shared_ptr<EndpointChannel> channel =
      channel_manager_->GetChannelForEndpoint(endpoint_id);
  if (!channel) {
    LOG(ERROR)
        << "BwuManager didn't find a previous EndpointChannel for "
        << endpoint_id
        << " when sending an upgrade failure frame, short-circuiting the "
           "upgrade protocol.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::CHANNEL_ERROR,
        BandwidthUpgradeErrorStage::NETWORK_AVAILABLE,
        OperationResultCode::NEARBY_GENERIC_OLD_ENDPOINT_CHANNEL_NULL);
    return;
  }

  // Report UPGRADE_FAILURE to the remote device.
  if (!channel->Write(parser::ForBwuFailure(upgrade_path_info)).Ok()) {
    channel->Close(DisconnectionReason::IO_ERROR);

    LOG(ERROR) << "BwuManager failed to write BWU_NEGOTIATION.UPGRADE_FAILURE "
                  "OfflineFrame to endpoint "
               << endpoint_id << ", short-circuiting the upgrade protocol.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::RESULT_IO_ERROR,
        BandwidthUpgradeErrorStage::NETWORK_AVAILABLE,
        OperationResultCode::CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);
    return;
  }

  // And lastly, clean up our medium since we failed to utilize it anyways.
  if (GetBwuMediumForEndpoint(endpoint_id) != Medium::UNKNOWN_MEDIUM) {
    RevertBwuMediumForEndpoint(channel->GetServiceId(), endpoint_id);
  }
  in_progress_upgrades_.erase(endpoint_id);
  LOG(INFO) << "BwuManager has informed endpoint " << endpoint_id
            << " that the bandwidth upgrade failed.";
}

bool BwuManager::ReadClientIntroductionFrame(EndpointChannel* channel,
                                             ClientIntroduction& introduction) {
  LOG(INFO) << "ReadClientIntroductionFrame with channel name: "
            << channel->GetName() << ", medium: "
            << location::nearby::proto::connections::Medium_Name(
                   channel->GetMedium());
  CancelableAlarm timeout_alarm(
      "BwuManager::ReadClientIntroductionFrame",
      [channel]() {
        LOG(ERROR) << "In BwuManager, failed to read the "
                      "ClientIntroductionFrame after "
                   << absl::FormatDuration(kReadClientIntroductionFrameTimeout)
                   << ". Timing out and closing EndpointChannel "
                   << channel->GetType();
        channel->Close();
      },
      kReadClientIntroductionFrameTimeout, &alarm_executor_);
  auto data = channel->Read();
  timeout_alarm.Cancel();
  if (!data.ok()) return false;
  auto transfer(parser::FromBytes(data.result()));
  if (!transfer.ok()) {
    LOG(ERROR) << "In ReadClientIntroductionFrame, attempted to read a "
                  "ClientIntroductionFrame from EndpointChannel "
               << channel->GetType()
               << " but was unable to obtain any OfflineFrame.";
    return false;
  }
  OfflineFrame frame = transfer.result();
  if (!frame.has_v1() || !frame.v1().has_bandwidth_upgrade_negotiation()) {
    LOG(ERROR) << "In ReadClientIntroductionFrame, expected a "
                  "BANDWIDTH_UPGRADE_NEGOTIATION v1 OfflineFrame but got a "
               << parser::GetFrameType(frame) << " frame instead.";
    return false;
  }
  if (frame.v1().bandwidth_upgrade_negotiation().event_type() !=
      BandwidthUpgradeNegotiationFrame::CLIENT_INTRODUCTION) {
    LOG(ERROR)
        << "In ReadClientIntroductionFrame, expected a CLIENT_INTRODUCTION "
           "v1 OfflineFrame but got a BANDWIDTH_UPGRADE_NEGOTIATION frame "
           "with eventType "
        << frame.v1().bandwidth_upgrade_negotiation().event_type()
        << " instead.";
    return false;
  }
  const auto& frame_intro =
      frame.v1().bandwidth_upgrade_negotiation().client_introduction();
  introduction = frame_intro;
  return true;
}

bool BwuManager::ReadClientIntroductionAckFrame(EndpointChannel* channel) {
  LOG(INFO) << "ReadClientIntroductionAckFrame with channel name: "
            << channel->GetName() << ", medium: "
            << location::nearby::proto::connections::Medium_Name(
                   channel->GetMedium());
  CancelableAlarm timeout_alarm(
      "BwuManager::ReadClientIntroductionAckFrame",
      [channel]() {
        LOG(ERROR)
            << "In BwuManager, failed to read the ClientIntroductionAckFrame "
               "after "
            << absl::FormatDuration(kReadClientIntroductionFrameTimeout)
            << ". Timing out and closing EndpointChannel "
            << channel->GetType();
        channel->Close();
      },
      kReadClientIntroductionFrameTimeout, &alarm_executor_);
  auto data = channel->Read();
  timeout_alarm.Cancel();
  if (!data.ok()) return false;
  auto transfer(parser::FromBytes(data.result()));
  if (!transfer.ok()) return false;
  OfflineFrame frame = transfer.result();
  if (!frame.has_v1() || !frame.v1().has_bandwidth_upgrade_negotiation())
    return false;
  if (frame.v1().bandwidth_upgrade_negotiation().event_type() !=
      BandwidthUpgradeNegotiationFrame::CLIENT_INTRODUCTION_ACK)
    return false;
  return true;
}

bool BwuManager::WriteClientIntroductionAckFrame(EndpointChannel* channel) {
  LOG(INFO) << "WriteClientIntroductionAckFrame channel name: "
            << channel->GetName() << ", medium: "
            << location::nearby::proto::connections::Medium_Name(
                   channel->GetMedium());
  return channel->Write(parser::ForBwuIntroductionAck()).Ok();
}

void BwuManager::ProcessLastWriteToPriorChannelEvent(
    ClientProxy* client, const std::string& endpoint_id) {
  // By this point in the upgrade protocol, there is the guarantee that both
  // involved endpoints have registered a new EndpointChannel with the
  // EndpointChannelManager as the official channel for communication; given
  // the way communication is structured in the EndpointManager, this means
  // that all new writes are happening over that new EndpointChannel, but
  // reads are still happening over this prior EndpointChannel (to avoid data
  // loss). But now that we've received this definitive final write over that
  // prior EndpointChannel, we can let the remote device that they can safely
  // close their end of this now-dormant EndpointChannel.
  EndpointChannel* previous_endpoint_channel =
      previous_endpoint_channels_[endpoint_id].get();
  if (!previous_endpoint_channel) {
    LOG(ERROR)
        << "BwuManager received a BWU_NEGOTIATION.LAST_WRITE_TO_PRIOR_CHANNEL "
           "OfflineFrame for unknown endpoint "
        << endpoint_id << ", can't complete the upgrade protocol.";
    successfully_upgraded_endpoints_.emplace(endpoint_id);
    return;
  }

  LOG(INFO) << "ProcessLastWriteToPriorChannelEvent: service_id="
            << previous_endpoint_channel->GetServiceId()
            << ", endpoint_id=" << endpoint_id << ", medium="
            << location::nearby::proto::connections::Medium_Name(
                   previous_endpoint_channel->GetMedium());

  if (!previous_endpoint_channel->Write(parser::ForBwuSafeToClose()).Ok()) {
    previous_endpoint_channel->Close(DisconnectionReason::IO_ERROR);
    // Remove this prior EndpointChannel from previous_endpoint_channels to
    // avoid leaks.
    previous_endpoint_channels_.erase(endpoint_id);

    LOG(ERROR) << "BwuManager failed to write "
                  "BWU_NEGOTIATION.SAFE_TO_CLOSE_PRIOR_CHANNEL "
                  "OfflineFrame to endpoint "
               << endpoint_id << ", short-circuiting the upgrade protocol.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::RESULT_IO_ERROR,
        BandwidthUpgradeErrorStage::SAFE_TO_CLOSE_PRIOR_CHANNEL,
        OperationResultCode::CONNECTIVITY_GENERIC_WRITING_CHANNEL_IO_ERROR);
    return;
  }
  VLOG(1) << "BwuManager successfully wrote "
             "BWU_NEGOTIATION.SAFE_TO_CLOSE_PRIOR_CHANNEL "
             "OfflineFrame while trying to upgrade endpoint "
          << endpoint_id;

  // The upgrade protocol's clean shutdown of the prior EndpointChannel will
  // conclude when we receive a corresponding
  // BANDWIDTH_UPGRADE_NEGOTIATION.SAFE_TO_CLOSE_PRIOR_CHANNEL OfflineFrame
  // from the remote device.
}

void BwuManager::ProcessSafeToClosePriorChannelEvent(
    ClientProxy* client, const std::string& endpoint_id) {
  LOG(INFO) << "ProcessSafeToClosePriorChannelEvent for endpoint "
            << endpoint_id;
  // By this point in the upgrade protocol, there's no more writes happening
  // over the prior EndpointChannel, and the remote device has given us the
  // go-ahead to close this EndpointChannel [1], so we can safely close it
  // (and depend on the EndpointManager querying the EndpointChannelManager to
  // start reading from the new EndpointChannel).
  //
  // [1] Which also implies that they've received our
  // BANDWIDTH_UPGRADE_NEGOTIATION.LAST_WRITE_TO_PRIOR_CHANNEL OfflineFrame),
  // so there can be no data loss, regardless of whether the EndpointChannel
  // allows reads of queued, unread data after the EndpointChannel has been
  // closed from the other end (as is the case with conventional TCP sockets)
  // or not (as is the case with Android's Bluetooth sockets, where closing
  // instantly throws an IOException on the remote device).
  auto item = previous_endpoint_channels_.extract(endpoint_id);
  auto& previous_endpoint_channel = item.mapped();
  if (previous_endpoint_channel == nullptr) {
    LOG(ERROR)
        << "BwuManager received a BWU_NEGOTIATION.SAFE_TO_CLOSE_PRIOR_CHANNEL "
           "OfflineFrame for unknown endpoint "
        << endpoint_id << ", can't complete the upgrade protocol.";
    return;
  }
  LOG(INFO) << "BwuManager successfully received a "
            << "BWU_NEGOTIATION.SAFE_TO_CLOSE_PRIOR_CHANNEL OfflineFrame while "
            << "trying to upgrade endpoint " << endpoint_id;

  // Each encrypted message includes the key to decrypt the next message. The
  // disconnect message is optional and may not be received under normal
  // circumstances so it is necessary to send it unencrypted. This way the
  // serial crypto context does not increment here.
  previous_endpoint_channel->DisableEncryption();
  LOG(INFO) << "[safe-to-disconnect] Sending "
               "DISCONNECTION frame with request 0, ack 0";
  previous_endpoint_channel->Write(
      parser::ForDisconnection(/* request_safe_to_disconnect */ false,
                               /* ack_safe_to_disconnect */ false));

  // Attempt to read the disconnect message from the previous channel. We don't
  // care whether we successfully read it or whether we get an exception here.
  // The idea is just to make sure the other side has had a chance to receive
  // the full SAFE_TO_CLOSE_PRIOR_CHANNEL message before we actually close the
  // channel. See b/172380349 for more context.
  previous_endpoint_channel->Read();
  previous_endpoint_channel->Close(DisconnectionReason::UPGRADED);

  VLOG(1) << "BwuManager cleanly shut down prior "
          << previous_endpoint_channel->GetType()
          << " EndpointChannel to conclude upgrade protocol for endpoint "
          << endpoint_id;

  // Now the upgrade protocol has completed, record analytics for this new
  // upgraded bandwidth connection...
  client->GetAnalyticsRecorder().OnConnectionEstablished(
      endpoint_id, GetBwuMediumForEndpoint(endpoint_id),
      client->GetConnectionToken(endpoint_id));
  // ...and the success of the upgrade itself.
  client->GetAnalyticsRecorder().OnBandwidthUpgradeSuccess(endpoint_id);

  // Now that the old channel has been drained, we can unpause the new channel
  std::shared_ptr<EndpointChannel> channel =
      channel_manager_->GetChannelForEndpoint(endpoint_id);

  if (!channel) {
    LOG(ERROR) << "BwuManager attempted to resume the current "
                  "EndpointChannel with endpoint "
               << endpoint_id << ", but none was found.";
    return;
  }

  channel->Resume();

  // Report the success to the client
  Medium medium = channel->GetMedium();
  if (FeatureFlags::GetInstance()
          .GetFlags()
          .support_web_rtc_non_cellular_medium) {
    if (medium == Medium::WEB_RTC && !mediums_->GetWebRtc().IsUsingCellular()) {
      medium = Medium::WEB_RTC_NON_CELLULAR;
    }
  }
  client->OnBandwidthChanged(endpoint_id, medium);
  in_progress_upgrades_.erase(endpoint_id);
}

void BwuManager::ProcessUpgradeFailureEvent(
    ClientProxy* client, const std::string& endpoint_id,
    const UpgradePathInfo& upgrade_info, BandwidthUpgradeResult result,
    bool record_analytic, OperationResultCode operation_result_code) {
  LOG(INFO) << "ProcessUpgradeFailureEvent for endpoint " << endpoint_id
            << " from medium: "
            << location::nearby::proto::connections::Medium_Name(
                   parser::UpgradePathInfoMediumToMedium(
                       upgrade_info.medium()));
  // The remote device failed to upgrade to the new medium we set up for them.
  // That's alright! We'll just try the next available medium (if there is one).
  in_progress_upgrades_.erase(endpoint_id);

  // The first thing we have to do is to replace our currentBwuMedium with the
  // next best upgrade medium we share with the remote device. The catch is that
  // we can only do this if we only have one connected endpoint. Otherwise,
  // we'll end up disrupting our other connected peers.
  // Note: This is only necessary when |support_multiple_bwu_mediums| is
  // disabled.
  if (!FeatureFlags::GetInstance().GetFlags().support_multiple_bwu_mediums &&
      channel_manager_->GetConnectedEndpointsCount() > 1) {
    // We can't change the currentBwuMedium, so there are no more upgrade
    // attempts for this endpoint. Sorry.
    LOG(ERROR)
        << "BwuManager failed to attempt a new bandwidth upgrade for endpoint "
        << endpoint_id
        << " because we have other connected endpoints and can't try a new "
           "upgrade medium.";
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, BandwidthUpgradeResult::CHANNEL_ERROR,
        BandwidthUpgradeErrorStage::NETWORK_AVAILABLE, operation_result_code);
    return;
  }

  EndpointChannel* channel =
      channel_manager_->GetChannelForEndpoint(endpoint_id).get();
  std::string service_id =
      channel ? channel->GetServiceId() : std::string(kUnknownServiceId);
  // Revert the existing upgrade medium for now.
  if (GetBwuMediumForEndpoint(endpoint_id) != Medium::UNKNOWN_MEDIUM) {
    // This is the BWU initiator, so append "_UPGRADE" to service_id.
    // Otherwise the Revert logic won't call into platform medium layer code.
    std::string upgrade_service_id = WrapInitiatorUpgradeServiceId(service_id);
    RevertBwuMediumForEndpoint(upgrade_service_id, endpoint_id);
  }

  if (record_analytic) {
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, result, BandwidthUpgradeErrorStage::NETWORK_AVAILABLE,
        operation_result_code);
  }

  // Loop through the ordered list of upgrade mediums. One by one, remove the
  // top element until we get to the medium we last attempted to upgrade to. The
  // remainder of the list will contain the mediums we haven't attempted yet.
  Medium last = parser::UpgradePathInfoMediumToMedium(upgrade_info.medium());
  std::vector<Medium> all_possible_mediums =
      client->GetUpgradeMediums(endpoint_id).GetMediums(true);
  std::vector<Medium> untried_mediums(all_possible_mediums);
  for (Medium medium : all_possible_mediums) {
    untried_mediums.erase(untried_mediums.begin());
    if (medium == last) {
      break;
    }
  }

  TryNextBestUpgradeMediums(client, endpoint_id, untried_mediums);
}

void BwuManager::TryNextBestUpgradeMediums(
    ClientProxy* client, const std::string& endpoint_id,
    std::vector<Medium> upgrade_mediums) {
  Medium next_medium = ChooseBestUpgradeMedium(endpoint_id, upgrade_mediums);
  LOG(INFO) << "Try Next Best Medium for endpoint " << endpoint_id
            << " after ChooseBestUpgradeMedium: "
            << location::nearby::proto::connections::Medium_Name(next_medium);

  // If current medium is not WiFi and we have not succeeded with upgrading yet,
  // retry upgrade.
  // TODO(b/228610864): Do we actually want to treat WIFI_LAN differently?
  auto channel = channel_manager_->GetChannelForEndpoint(endpoint_id);
  Medium current_medium =
      channel ? channel->GetMedium() : Medium::UNKNOWN_MEDIUM;
  VLOG(1) << "current_medium: "
          << location::nearby::proto::connections::Medium_Name(current_medium);
  if (current_medium != Medium::WIFI_LAN &&
      (next_medium == current_medium || next_medium == Medium::UNKNOWN_MEDIUM ||
       upgrade_mediums.empty())) {
    RetryUpgradesAfterDelay(client, endpoint_id);
    return;
  }

  // Attempt to set the new upgrade medium.
  if (!GetHandlerForMedium(next_medium)) {
    // As Medium without handler has been stripped out, this shouldn't be hit
    LOG(INFO)
        << "BwuManager failed to attempt a new bandwidth upgrade for endpoint "
        << endpoint_id
        << " because we couldn't set a new bandwidth upgrade medium.";
    return;
  }
  SetBwuMediumForEndpoint(endpoint_id, next_medium);

  // Now that we've successfully picked a new upgrade medium to try,
  // re-initiate the bandwidth upgrade.
  LOG(INFO) << "BwuManager is attempting to upgrade endpoint " << endpoint_id
            << " again with a new bandwidth upgrade medium.";
  InitiateBwuForEndpoint(client, endpoint_id, next_medium);
}

std::vector<Medium> BwuManager::StripOutUnavailableMediums(
    const std::vector<Medium>& mediums) const {
  std::vector<Medium> available_mediums;
  for (Medium m : mediums) {
    bool available = false;
    if (GetHandlerForMedium(m)) {
      switch (m) {
        case Medium::AWDL:
          available = mediums_->GetAwdl().IsAvailable();
          break;
        case Medium::WIFI_LAN:
          available = mediums_->GetWifiLan().IsAvailable();
          break;
        case Medium::WIFI_DIRECT:
          available = mediums_->GetWifiDirect().IsGOAvailable();
          break;
        case Medium::WIFI_HOTSPOT:
          available = mediums_->GetWifiHotspot().IsAPAvailable();
          break;
        case Medium::WEB_RTC:
          available = mediums_->GetWebRtc().IsAvailable();
          break;
        case Medium::BLUETOOTH:
          available = mediums_->GetBluetoothClassic().IsAvailable();
          break;
        default:
          break;
      }
    }
    if (available) {
      available_mediums.push_back(m);
    }
  }
  return available_mediums;
}

// Returns the optimal medium supported by both devices.
// Each medium in the passed in list is checked for its availability with the
// medium_manager_ to ensure that the chosen upgrade medium is supported and
// available locally before continuing the upgrade. Once we pick a medium, all
// future connections will use it too. eg. If we chose Wifi LAN, we'll attempt
// to upgrade the 2nd, 3rd, etc remote endpoints with Wifi LAN even if they're
// on a different network (or had a better medium). This is a quick and easy
// way to prevent mediums, like Wifi Hotspot, from interfering with active
// connections (although it's suboptimal for bandwidth throughput). When all
// endpoints disconnect, we reset the bandwidth upgrade medium.
Medium BwuManager::ChooseBestUpgradeMedium(
    const std::string& endpoint_id, const std::vector<Medium>& mediums) const {
  auto available_mediums = StripOutUnavailableMediums(mediums);
  Medium current_medium = GetBwuMediumForEndpoint(endpoint_id);
  if (current_medium == Medium::UNKNOWN_MEDIUM) {
    if (!available_mediums.empty()) {
      // Case 1: This is our first time upgrading, and we have at least one
      // supported medium to choose from. Return the first medium in the list,
      // since they are ordered by preference.
      return available_mediums[0];
    }
    // Case 2: This is our first time upgrading, but there are no available
    // upgrade mediums. Fall through to returning UNKNOWN_MEDIUM at the
    // bottom.
    LOG(INFO)
        << "Current upgrade medium is unset, but there are no common supported "
           "upgrade mediums.";
  } else {
    // Case 3: We have already upgraded, and there is a list of supported
    // mediums to check against. Return the current upgrade medium if it's in
    // the supported list.
    if (std::find(available_mediums.begin(), available_mediums.end(),
                  current_medium) != available_mediums.end()) {
      return current_medium;
    }
    // Case 4: We have already upgraded, but the current medium is not
    // supported by the remote endpoint (it's not in the list, or the list is
    // empty). Fall through and return Medium.UNKNOWN_MEDIUM because we cannot
    // continue with the current upgrade medium, and we are not allowed to
    // switch.
    std::string mediums_string;
    for (const auto& medium : available_mediums) {
      absl::StrAppend(&mediums_string,
                      location::nearby::proto::connections::Medium_Name(medium),
                      "; ");
    }
    LOG(INFO) << "Current upgrade medium "
              << location::nearby::proto::connections::Medium_Name(
                     current_medium)
              << " is not supported by the remote endpoint (supported mediums: "
              << mediums_string << ")";
  }

  return Medium::UNKNOWN_MEDIUM;
}

void BwuManager::RetryUpgradesAfterDelay(ClientProxy* client,
                                         const std::string& endpoint_id) {
  absl::Duration delay = CalculateNextRetryDelay(endpoint_id);
  CancelRetryUpgradeAlarm(endpoint_id);
  auto alarm = std::make_unique<CancelableAlarm>(
      "BWU alarm",
      [this, client, endpoint_id]() {
        RunOnBwuManagerThread(
            "bwu-retry-upgrade", [this, client, endpoint_id]() {
              if (!client->IsConnectedToEndpoint(endpoint_id)) {
                return;
              }
              TryNextBestUpgradeMediums(
                  client, endpoint_id,
                  client->GetUpgradeMediums(endpoint_id).GetMediums(true));
            });
      },
      delay, &alarm_executor_);

  retry_upgrade_alarms_.emplace(endpoint_id,
                                std::make_pair(std::move(alarm), delay));
  retry_delays_[endpoint_id] = delay;
  LOG(INFO) << "Retry bandwidth upgrade after " << absl::FormatDuration(delay);
}

void BwuManager::AttemptToRecordBandwidthUpgradeErrorForUnknownEndpoint(
    BandwidthUpgradeResult result, BandwidthUpgradeErrorStage error_stage,
    OperationResultCode operation_result_code) {
  if (in_progress_upgrades_.size() == 1) {
    auto it = in_progress_upgrades_.begin();
    std::string endpoint_id = it->first;
    ClientProxy* client = it->second;
    // Note: Even though we know this is an error, we cannot clear state yet.
    // We've sent the remote device the credentials they need and it's up to
    // them if they want to repeatedly attempt to connect or if they want to
    // give up and have us try a different medium. This isn't a decision we can
    // make for them.
    client->GetAnalyticsRecorder().OnBandwidthUpgradeError(
        endpoint_id, result, error_stage, operation_result_code);
    LOG(INFO) << "BwuManager got error " << BandwidthUpgradeResult_Name(result)
              << " at stage " << BandwidthUpgradeErrorStage_Name(error_stage)
              << " when upgrading endpoint " << endpoint_id;
  }
  // Otherwise, we have no way of knowing which endpoint was trying to connect
  // to us :(
  LOG(INFO) << "BwuManager got error " << BandwidthUpgradeResult_Name(result)
            << " at stage " << BandwidthUpgradeErrorStage_Name(error_stage)
            << ", but we don't know which endpoint was trying to "
               "connect to us, so skipping analytics for his error.";
}

bool BwuManager::NeedToSwitchRole(
    ClientProxy* client, const std::string& endpoint_id, Medium medium,
    const location::nearby::connections::MediumRole& medium_role) {
  if (GetLocalOsInfo(client).type() == OsInfo::APPLE) {
    switch (medium) {
      case Medium::WIFI_HOTSPOT:
        return medium_role.support_wifi_hotspot_host();
      default:
        break;
    }
  }
  return false;
}

const location::nearby::connections::OsInfo& BwuManager::GetLocalOsInfo(
    ClientProxy* client) const {
  return client->GetLocalOsInfo();
}

absl::Duration BwuManager::CalculateNextRetryDelay(
    const std::string& endpoint_id) {
  auto item = retry_delays_.find(endpoint_id);
  auto initial_delay = config_.bandwidth_upgrade_retry_delay;
  if (item == retry_delays_.end()) return initial_delay;
  /*
   * Without use_exp_backoff_in_bwu_retry, bwu retry intervals for the same
   * endpoint_id in seconds will be like: 5, 10, 10, 10...
   * With use_exp_backoff_in_bwu_retry enabled, bwu retry intervals in seconds
   * would be like: 3, 6, 12, 24, ... 300, 300, ...
   */
  auto delay =
      FeatureFlags::GetInstance().GetFlags().use_exp_backoff_in_bwu_retry
          ? item->second * 2
          : item->second + initial_delay;
  return std::min(delay, config_.bandwidth_upgrade_retry_max_delay);
}

void BwuManager::CancelRetryUpgradeAlarm(const std::string& endpoint_id) {
  LOG(INFO) << "CancelRetryUpgradeAlarm for endpoint " << endpoint_id;
  auto item = retry_upgrade_alarms_.extract(endpoint_id);
  if (item.empty()) return;
  auto& pair = item.mapped();
  pair.first->Cancel();
}

void BwuManager::CancelAllRetryUpgradeAlarms() {
  LOG(INFO) << "CancelAllRetryUpgradeAlarms invoked";
  for (auto& item : retry_upgrade_alarms_) {
    const std::string& endpoint_id = item.first;
    CancelableAlarm* cancellable_alarm = item.second.first.get();
    LOG(INFO) << "CancelRetryUpgradeAlarm for endpoint " << endpoint_id;
    cancellable_alarm->Cancel();
  }
  retry_upgrade_alarms_.clear();
  retry_delays_.clear();
}

}  // namespace connections
}  // namespace nearby