// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <fuchsia/hardware/test/c/banjo.h>
#include <lib/ddk/binding_driver.h>
#include <lib/ddk/device.h>
#include <lib/ddk/driver.h>
#include <lib/fidl/cpp/wire/vector_view.h>
#include <lib/zircon-internal/thread_annotations.h>
#include <lib/zx/channel.h>
#include <lib/zx/vmo.h>
#include <stdio.h>
#include <threads.h>
#include <zircon/assert.h>
#include <zircon/errors.h>
#include <zircon/process.h>
#include <zircon/status.h>
#include <zircon/types.h>

#include <memory>
#include <optional>
#include <type_traits>
#include <variant>

#include <ddktl/device.h>
#include <fbl/array.h>
#include <fbl/auto_lock.h>
#include <fbl/mutex.h>
#include <fbl/vector.h>

#include "ddktl/suspend-txn.h"
#include "fidl.h"

namespace mock_device {

class MockDevice;
using MockDeviceType =
    ddk::Device<MockDevice, ddk::GetProtocolable, ddk::Initializable, ddk::Unbindable,
                ddk::Suspendable, ddk::Resumable, ddk::Rxrpcable>;

class MockDevice : public MockDeviceType {
 public:
  MockDevice(zx_device_t* device, fidl::ClientEnd<device_mock::MockDevice> controller);
  static zx_status_t Create(zx_device_t* parent,
                            fidl::ClientEnd<device_mock::MockDevice> controller,
                            std::unique_ptr<MockDevice>* out);

  // Device protocol implementation.
  void DdkRelease();
  zx_status_t DdkGetProtocol(uint32_t proto_id, void* out);
  void DdkInit(ddk::InitTxn txn) { txn.Reply(ZX_OK); }
  void DdkUnbind(ddk::UnbindTxn txn);
  void DdkSuspend(ddk::SuspendTxn txn);
  void DdkResume(ddk::ResumeTxn txn);
  zx_status_t DdkRxrpc(zx_handle_t channel);

  // Generate an invocation record for a hook RPC
  device_mock::wire::HookInvocation ConstructHookInvocation();
  static device_mock::wire::HookInvocation ConstructHookInvocation(uint64_t device_id);

  // Create a new thread that will serve a MockDeviceThread interface over |server_end|
  void CreateThread(fidl::ServerEnd<device_mock::MockDeviceThread> server_end);

 private:
  // Retrieve the current thread's process and thread koids
  static void GetThreadKoids(zx_koid_t* process, zx_koid_t* thread);

  static int ThreadFunc(void* arg);
  struct ThreadFuncArg {
    // channel that the thread will use to serve a MockDeviceThread
    // interface over
    fidl::ServerEnd<device_mock::MockDeviceThread> server_end;
    // Device this thread is executing for
    MockDevice* dev;
  };

  fbl::Mutex lock_;
  // List of threads spawned by Actions
  fbl::Vector<thrd_t> threads_ TA_GUARDED(lock_);

  // Our half of the controller channel.  We will send requests for input on
  // it.
  fidl::WireSyncClient<device_mock::MockDevice> controller_;
};

struct ProcessActionsContext {
  // A channel that is either borrowing (zx::unowned_channel) or
  // owned (ServerEnd).
  // In the borrowing case, the channel must outlive this variant.
  using ChannelVariants = std::variant<fidl::UnownedClientEnd<device_mock::MockDevice>,
                                       fidl::ServerEnd<device_mock::MockDeviceThread>>;

  // Constructs a process action context. Note that |channel_variants| must
  // outlive this context. Typically, the |channel_variants| is first created
  // on the caller's stack, after which |ProcessActionContext| is constructed
  // referencing it.
  ProcessActionsContext(ChannelVariants& channel_variants, bool has_hook_status,
                        MockDevice* mock_device, zx_device_t* device)
      : channel_variants(channel_variants),
        has_hook_status(has_hook_status),
        mock_device(mock_device),
        device(device) {}

  // IN: The channel that these actions came from.  Used for acknowledging
  // add/remove device requests.
  //
  // When this context is running in a separate thread, the context has the
  // |fidl::ServerEnd<device_mock::MockDeviceThread>| variant.
  // When this context is running in the same thread, the context has the
  // |zx::unowned_channel| variant, and is the client-end of the
  // MockDevice protocol.
  //
  // Note that in either case, the context does not own the underlying channel,
  // since this field is a reference. The channel is usually owned by a caller
  // which created the context.
  ChannelVariants& channel_variants;

  bool has_hook_status = false;
  // OUT: What should be returned by the hook
  zx_status_t hook_status = ZX_ERR_INTERNAL;
  // IN: A buffer that can be written to by actions (nullptr if none)
  void* associated_buf = nullptr;
  size_t associated_buf_count = 0;
  // OUT: Number of bytes written by actions
  size_t associated_buf_actual = 0;
  // IN/OUT: MockDevice to use for associating threads with.  NULL'd out if
  // remove was called.
  MockDevice* mock_device = nullptr;
  // IN/OUT: Device to use for invoking add_device/remove_device.  NULL'd out
  // if remove was called.
  zx_device_t* device = nullptr;

  // IN: The txn used to reply to the unbind hook.
  std::optional<ddk::UnbindTxn> pending_unbind_txn = std::nullopt;
  // IN: The txn used to reply to the suspend hook.
  std::optional<ddk::SuspendTxn> pending_suspend_txn = std::nullopt;
  // IN: The txn used to reply to the resume hook.
  std::optional<ddk::ResumeTxn> pending_resume_txn = std::nullopt;
};

// Helper struct to simplify matching an |std::variant|.
template <class... Ts>
struct matchers : Ts... {
  using Ts::operator()...;
};
template <class... Ts>
matchers(Ts...) -> matchers<Ts...>;

// Execute the actions returned by a hook
zx_status_t ProcessActions(fidl::VectorView<device_mock::wire::Action> actions,
                           ProcessActionsContext* ctx);

MockDevice::MockDevice(zx_device_t* device, fidl::ClientEnd<device_mock::MockDevice> controller)
    : MockDeviceType(device),
      controller_(fidl::WireSyncClient<device_mock::MockDevice>(std::move(controller))) {}

int MockDevice::ThreadFunc(void* raw_arg) {
  auto arg = std::unique_ptr<ThreadFuncArg>(static_cast<ThreadFuncArg*>(raw_arg));
  ProcessActionsContext::ChannelVariants server_end = std::move(arg->server_end);

  while (true) {
    fbl::Array<device_mock::wire::Action> actions;
    zx_status_t status = WaitForPerformActions(
        std::get<fidl::ServerEnd<device_mock::MockDeviceThread>>(server_end).channel(), &actions);
    if (status != ZX_OK) {
      ZX_ASSERT_MSG(status == ZX_ERR_STOP, "MockDevice thread exiting: %s\n",
                    zx_status_get_string(status));
      break;
    }
    ProcessActionsContext ctx(server_end, false, arg->dev, arg->dev->zxdev());
    status = ProcessActions(
        fidl::VectorView<device_mock::wire::Action>::FromExternal(actions.data(), actions.size()),
        &ctx);
    ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
    if (ctx.device == nullptr) {
      // If the device was removed, bail out since we're releasing.
      break;
    }
  }
  return 0;
}

void MockDevice::CreateThread(fidl::ServerEnd<device_mock::MockDeviceThread> server_end) {
  auto arg = std::make_unique<ThreadFuncArg>();
  arg->server_end = std::move(server_end);
  arg->dev = this;

  fbl::AutoLock guard(&lock_);
  thrd_t thrd;
  int ret = thrd_create(&thrd, MockDevice::ThreadFunc, arg.get());
  ZX_ASSERT(ret == thrd_success);
  // The thread now owns this pointer
  [[maybe_unused]] auto ptr = arg.release();

  threads_.push_back(thrd);
}

void MockDevice::GetThreadKoids(zx_koid_t* process, zx_koid_t* thread) {
  struct Koids {
    zx_koid_t process;
    zx_koid_t thread;
  };
  thread_local Koids thread_koids;

  if (thread_koids.process == 0 && thread_koids.thread == 0) {
    zx_info_handle_basic_t info;
    zx_status_t status = zx_object_get_info(zx_thread_self(), ZX_INFO_HANDLE_BASIC, &info,
                                            sizeof(info), nullptr, nullptr);
    ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
    thread_koids.process = info.related_koid;
    thread_koids.thread = info.koid;
  }

  *process = thread_koids.process;
  *thread = thread_koids.thread;
}

device_mock::wire::HookInvocation MockDevice::ConstructHookInvocation(uint64_t device_id) {
  device_mock::wire::HookInvocation invoc;
  GetThreadKoids(&invoc.process_koid, &invoc.thread_koid);
  invoc.device_id = device_id;
  return invoc;
}

device_mock::wire::HookInvocation MockDevice::ConstructHookInvocation() {
  return ConstructHookInvocation(reinterpret_cast<uintptr_t>(zxdev()));
}

void MockDevice::DdkRelease() {
  auto result = controller_->Release(ConstructHookInvocation());
  ZX_ASSERT_MSG(result.ok(), "%s", result.status_string());

  // Launch a thread to do the actual joining and delete, since this could get
  // called from a thread.
  thrd_t thrd;
  int ret = thrd_create(
      &thrd,
      [](void* arg) {
        auto me = static_cast<MockDevice*>(arg);
        {
          fbl::AutoLock guard(&me->lock_);
          for (auto& t : me->threads_) {
            thrd_join(t, nullptr);
          }
        }
        delete me;
        return 0;
      },
      this);
  ZX_ASSERT(ret == thrd_success);
  thrd_detach(thrd);
}

zx_status_t MockDevice::DdkGetProtocol(uint32_t proto_id, void* out) {
  auto result = controller_->GetProtocol(ConstructHookInvocation(), proto_id);
  ZX_ASSERT_MSG(result.ok(), "%s", result.status_string());
  ProcessActionsContext::ChannelVariants channel = controller_.client_end().borrow();
  ProcessActionsContext ctx(channel, true, this, zxdev());
  zx_status_t status = ProcessActions(result.value().actions, &ctx);
  ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
  return ctx.hook_status;
}

void MockDevice::DdkUnbind(ddk::UnbindTxn txn) {
  auto result = controller_->Unbind(ConstructHookInvocation());
  ZX_ASSERT_MSG(result.ok(), "%s", result.status_string());
  ProcessActionsContext::ChannelVariants channel = controller_.client_end().borrow();
  ProcessActionsContext ctx(channel, false, this, zxdev());
  ctx.pending_unbind_txn = std::move(txn);
  zx_status_t status = ProcessActions(result.value().actions, &ctx);
  ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
}

void MockDevice::DdkSuspend(ddk::SuspendTxn txn) {
  auto result = controller_->Suspend(ConstructHookInvocation(), txn.requested_state(),
                                     txn.enable_wake(), txn.suspend_reason());
  ZX_ASSERT_MSG(result.ok(), "%s", result.status_string());
  ProcessActionsContext::ChannelVariants channel = controller_.client_end().borrow();
  ProcessActionsContext ctx(channel, true, this, zxdev());
  ctx.pending_suspend_txn = std::move(txn);
  zx_status_t status = ProcessActions(result.value().actions, &ctx);
  ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
}

void MockDevice::DdkResume(ddk::ResumeTxn txn) {
  auto result = controller_->Resume(ConstructHookInvocation(), txn.requested_state());
  ZX_ASSERT_MSG(result.ok(), "%s", result.status_string());
  ProcessActionsContext::ChannelVariants channel = controller_.client_end().borrow();
  ProcessActionsContext ctx(channel, true, this, zxdev());
  ctx.pending_resume_txn = std::move(txn);
  zx_status_t status = ProcessActions(result.value().actions, &ctx);
  ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
}

zx_status_t MockDevice::DdkRxrpc(zx_handle_t channel) {
  auto result = controller_->Rxrpc(ConstructHookInvocation());
  ZX_ASSERT_MSG(result.ok(), "%s", result.status_string());
  ProcessActionsContext::ChannelVariants channel_variants = controller_.client_end().borrow();
  ProcessActionsContext ctx(channel_variants, true, this, zxdev());
  zx_status_t status = ProcessActions(result.value().actions, &ctx);
  ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
  return ctx.hook_status;
}

zx_status_t MockDevice::Create(zx_device_t* parent,
                               fidl::ClientEnd<device_mock::MockDevice> controller,
                               std::unique_ptr<MockDevice>* out) {
  auto dev = std::make_unique<MockDevice>(parent, std::move(controller));
  *out = std::move(dev);
  return ZX_OK;
}

zx_status_t ProcessActions(fidl::VectorView<device_mock::wire::Action> actions,
                           ProcessActionsContext* ctx) {
  for (size_t i = 0; i < actions.count(); ++i) {
    auto& action = actions[i];
    switch (action.Which()) {
      case device_mock::wire::Action::Tag::kReturnStatus: {
        if (i != actions.count() - 1) {
          printf("MockDevice::ProcessActions: return_status was not the final entry\n");
          return ZX_ERR_INVALID_ARGS;
        }
        if (!ctx->has_hook_status) {
          printf("MockDevice::ProcessActions: return_status present for no-status hook\n");
          return ZX_ERR_INVALID_ARGS;
        }
        ctx->hook_status = action.return_status();
        return ZX_OK;
      }
      case device_mock::wire::Action::Tag::kCreateThread: {
        if (ctx->mock_device == nullptr) {
          printf("MockDevice::CreateThread: asked to create thread without device\n");
          return ZX_ERR_INVALID_ARGS;
        }
        ctx->mock_device->CreateThread(std::move(action.create_thread()));
        break;
      }
      case device_mock::wire::Action::Tag::kAsyncRemoveDevice: {
        if (ctx->mock_device == nullptr) {
          printf("MockDevice::RemoveDevice: asked to remove device but none populated\n");
          return ZX_ERR_INVALID_ARGS;
        }
        ctx->mock_device->DdkAsyncRemove();
        break;
      }
      case device_mock::wire::Action::Tag::kUnbindReply: {
        if (!ctx->pending_unbind_txn) {
          printf("MockDevice::UnbindReply: asked to reply to unbind but no unbind is pending\n");
          return ZX_ERR_INVALID_ARGS;
        }
        ctx->pending_unbind_txn->Reply();
        // Null out the device pointers, since the release hook might get
        // activated.
        ctx->device = nullptr;
        ctx->mock_device = nullptr;
        zx_status_t status =
            std::visit(matchers{
                           [&](fidl::ServerEnd<device_mock::MockDeviceThread>& server_end) {
                             return fidl::WireSendEvent(server_end)
                                 ->UnbindReplyDone(action.unbind_reply().action_id)
                                 .status();
                           },
                           [&](fidl::UnownedClientEnd<device_mock::MockDevice>& client_end) {
                             return fidl::WireCall(client_end)
                                 ->UnbindReplyDone(action.unbind_reply().action_id)
                                 .status();
                           },
                       },
                       ctx->channel_variants);
        ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
        break;
      }

      case device_mock::wire::Action::Tag::kSuspendReply: {
        if (!ctx->pending_suspend_txn) {
          printf("MockDevice::SuspendReply: asked to reply to suspend but no suspend is pending\n");
          return ZX_ERR_INVALID_ARGS;
        }
        ctx->pending_suspend_txn->Reply(ZX_OK, 0);
        zx_status_t status =
            std::visit(matchers{
                           [&](fidl::ServerEnd<device_mock::MockDeviceThread>& server_end) {
                             return fidl::WireSendEvent(server_end)
                                 ->SuspendReplyDone(action.suspend_reply().action_id)
                                 .status();
                           },
                           [&](fidl::UnownedClientEnd<device_mock::MockDevice>& client_end) {
                             return fidl::WireCall(client_end)
                                 ->SuspendReplyDone(action.suspend_reply().action_id)
                                 .status();
                           },
                       },
                       ctx->channel_variants);
        ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
        break;
      }

      case device_mock::wire::Action::Tag::kResumeReply: {
        if (!ctx->pending_resume_txn) {
          printf("MockDevice::ResumeReply: asked to reply to resume but no resume is pending\n");
          return ZX_ERR_INVALID_ARGS;
        }
        ctx->pending_resume_txn->Reply(ZX_OK, 0, 0);
        zx_status_t status =
            std::visit(matchers{
                           [&](fidl::ServerEnd<device_mock::MockDeviceThread>& server_end) {
                             return fidl::WireSendEvent(server_end)
                                 ->ResumeReplyDone(action.resume_reply().action_id)
                                 .status();
                           },
                           [&](fidl::UnownedClientEnd<device_mock::MockDevice>& client_end) {
                             return fidl::WireCall(client_end)
                                 ->ResumeReplyDone(action.resume_reply().action_id)
                                 .status();
                           },
                       },
                       ctx->channel_variants);
        ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
        break;
      }
      case device_mock::wire::Action::Tag::kAddDevice: {
        // TODO(teisenbe): Implement more functionality here
        auto& add_device_action = action.add_device();
        ZX_ASSERT_MSG(!add_device_action.do_bind, "bind not yet supported\n");
        std::unique_ptr<MockDevice> dev;
        zx_status_t status =
            MockDevice::Create(ctx->device, std::move(add_device_action.controller), &dev);
        if (status != ZX_OK) {
          return status;
        }

        char name[ZX_DEVICE_NAME_MAX + 1];
        if (add_device_action.name.size() > sizeof(name) - 1) {
          return ZX_ERR_INVALID_ARGS;
        }
        memcpy(name, add_device_action.name.data(), add_device_action.name.size());
        name[add_device_action.name.size()] = 0;

        cpp20::span<zx_device_prop_t> props(
            reinterpret_cast<zx_device_prop_t*>(add_device_action.properties.data()),
            add_device_action.properties.count());

        status = dev->DdkAdd(ddk::DeviceAddArgs(name).set_props(props));
        if (status == ZX_OK) {
          // Devmgr now owns this
          [[maybe_unused]] auto ptr = dev.release();
        }
        if (add_device_action.expect_status != status) {
          return status;
        }

        status = std::visit(matchers{
                                [&](fidl::ServerEnd<device_mock::MockDeviceThread>& server_end) {
                                  return fidl::WireSendEvent(server_end)
                                      ->AddDeviceDone(add_device_action.action_id)
                                      .status();
                                },
                                [&](fidl::UnownedClientEnd<device_mock::MockDevice>& client_end) {
                                  return fidl::WireCall(client_end)
                                      ->AddDeviceDone(add_device_action.action_id)
                                      .status();
                                },
                            },
                            ctx->channel_variants);
        ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
        break;
      }
    }
  }

  if (!ctx->has_hook_status) {
    return ZX_OK;
  }

  printf("MockDevice::ProcessActions: return_status was not the final entry\n");
  return ZX_ERR_INVALID_ARGS;
}

zx_status_t MockDeviceBind(void* ctx, zx_device_t* parent) {
  // It's expected that this driver is binding against a device created by the
  // fuchsia.device.test interface.  Get the protocol from the device we're
  // binding to so we can wire up the control channel.
  test_protocol_t proto;
  zx_status_t status = device_get_protocol(parent, ZX_PROTOCOL_TEST, &proto);
  ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));

  fidl::ClientEnd<device_mock::MockDevice> client_end;
  test_get_channel(&proto, client_end.channel().reset_and_get_address());

  // Ask the control channel what to do about this bind().
  auto result =
      fidl::WireCall(client_end)
          ->Bind(MockDevice::ConstructHookInvocation(reinterpret_cast<uintptr_t>(parent)));
  ZX_ASSERT_MSG(result.ok(), "%s", result.status_string());

  ProcessActionsContext::ChannelVariants channel_variants = client_end.borrow();
  ProcessActionsContext pac_ctx(channel_variants, true, nullptr, parent);
  status = ProcessActions(result.value().actions, &pac_ctx);
  ZX_ASSERT_MSG(status == ZX_OK, "%s", zx_status_get_string(status));
  return pac_ctx.hook_status;
}

const zx_driver_ops_t kMockDeviceOps = []() {
  zx_driver_ops_t ops = {};
  ops.version = DRIVER_OPS_VERSION;
  ops.bind = MockDeviceBind;
  return ops;
}();

}  // namespace mock_device

ZIRCON_DRIVER(mock_device, mock_device::kMockDeviceOps, "zircon", "0.1");