/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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
 *
 *     http://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 <cstring>
#include <cstdlib> // for abort
#include <ts/ts.h> // for debug

// debug messages viewable by setting 'proxy.config.diags.debug.tags'
// in 'records.yaml'

// debug messages during one-time initialization
static const char DEBUG_TAG_INIT[] = "TSContSchedule_test.init";
static const char DEBUG_TAG_SCHD[] = "TSContSchedule_test.schedule";
static const char DEBUG_TAG_HDL[]  = "TSContSchedule_test.handler";
static const char DEBUG_TAG_CHK[]  = "TSContSchedule_test.check";

static DbgCtl dbg_ctl_init{DEBUG_TAG_INIT};
static DbgCtl dbg_ctl_schd{DEBUG_TAG_SCHD};
static DbgCtl dbg_ctl_hdl{DEBUG_TAG_HDL};
static DbgCtl dbg_ctl_chk{DEBUG_TAG_CHK};

// plugin registration info
static char plugin_name[]   = "TSContSchedule_test";
static char vendor_name[]   = "apache";
static char support_email[] = "duke8253@apache.org";

static int test_flag = 0;

static TSEventThread thread_1 = nullptr;
static TSEventThread thread_2 = nullptr;

static TSCont contp_1 = nullptr;
static TSCont contp_2 = nullptr;

static int TSContScheduleOnPool_handler_1(TSCont contp, TSEvent event, void *edata);
static int TSContScheduleOnPool_handler_2(TSCont contp, TSEvent event, void *edata);
static int TSContScheduleOnThread_handler_1(TSCont contp, TSEvent event, void *edata);
static int TSContScheduleOnThread_handler_2(TSCont contp, TSEvent event, void *edata);
static int TSContThreadAffinity_handler(TSCont contp, TSEvent event, void *edata);

static int
TSContScheduleOnPool_handler_1(TSCont /* contp ATS_UNUSED */, TSEvent /* event ATS_UNUSED */, void * /* edata ATS_UNUSED */)
{
  // This runs on ET_NET threads.
  Dbg(dbg_ctl_hdl, "TSContScheduleOnPool handler 1 thread [%p]", TSThreadSelf());
  if (thread_1 == nullptr) {
    // First time here, record thread id.
    thread_1 = TSEventThreadSelf();
  } else {
    // Second time here, we should be on a different thread since affinity was cleared.
    if (thread_1 != TSEventThreadSelf()) {
      Dbg(dbg_ctl_chk, "pass [should not be the same thread]");
    } else {
      Dbg(dbg_ctl_chk, "fail [on the same thread]");
    }
  }
  return 0;
}

static int
TSContScheduleOnPool_handler_2(TSCont /* contp ATS_UNUSED */, TSEvent /* event ATS_UNUSED */, void * /* edata ATS_UNUSED */)
{
  // This runs on ET_TASK threads.
  Dbg(dbg_ctl_hdl, "TSContScheduleOnPool handler 2 thread [%p]", TSThreadSelf());
  if (thread_2 == nullptr) {
    // First time here, record thread id.
    thread_2 = TSEventThreadSelf();
  } else {
    if (thread_2 == TSEventThreadSelf()) {
      // Second time there, we should be on the same thread even though affinity was cleared,
      // reason being plugin is running on ET_TASK threads, and we were scheduled on ET_TASK
      // threads as well, so the thread the plugin is on is used and set to affinity.
      Dbg(dbg_ctl_chk, "pass [should be the same thread]");
    } else {
      Dbg(dbg_ctl_chk, "fail [not the same thread]");
    }
  }
  return 0;
}

void
TSContScheduleOnPool_test()
{
  contp_1 = TSContCreate(TSContScheduleOnPool_handler_1, TSMutexCreate());
  contp_2 = TSContCreate(TSContScheduleOnPool_handler_2, TSMutexCreate());

  if (contp_1 == nullptr || contp_2 == nullptr) {
    Dbg(dbg_ctl_schd, "[%s] could not create continuation", plugin_name);
    abort();
  } else {
    Dbg(dbg_ctl_schd, "[%s] scheduling continuation", plugin_name);

    TSContScheduleOnPool(contp_1, 0, TS_THREAD_POOL_NET);
    TSContThreadAffinityClear(contp_1);
    TSContScheduleOnPool(contp_1, 100, TS_THREAD_POOL_NET);

    TSContScheduleOnPool(contp_2, 200, TS_THREAD_POOL_TASK);
    TSContThreadAffinityClear(contp_2);
    TSContScheduleOnPool(contp_2, 300, TS_THREAD_POOL_TASK);
  }
}

static int
TSContScheduleOnThread_handler_1(TSCont /* contp ATS_UNUSED */, TSEvent /* event ATS_UNUSED */, void * /* edata ATS_UNUSED */)
{
  // Mostly same as TSContScheduleOnPool_handler_1, no need to set affinity
  // since we are scheduling directly on to a thread.
  Dbg(dbg_ctl_hdl, "TSContScheduleOnThread handler 1 thread [%p]", TSThreadSelf());
  if (thread_1 == nullptr) {
    thread_1 = TSEventThreadSelf();

    Dbg(dbg_ctl_hdl, "[%s] scheduling continuation", plugin_name);
    TSContScheduleOnThread(contp_2, 0, thread_1);
    TSContScheduleOnThread(contp_2, 0, thread_1);
  } else if (thread_2 == nullptr) {
    Dbg(dbg_ctl_chk, "fail [schedule delay not applied]");
  } else {
    if (thread_2 != TSEventThreadSelf()) {
      Dbg(dbg_ctl_chk, "pass [should not be the same thread]");
    } else {
      Dbg(dbg_ctl_chk, "fail [on the same thread]");
    }
  }
  return 0;
}

static int
TSContScheduleOnThread_handler_2(TSCont /* contp ATS_UNUSED */, TSEvent /* event ATS_UNUSED */, void * /* edata ATS_UNUSED */)
{
  Dbg(dbg_ctl_hdl, "TSContScheduleOnThread handler 2 thread [%p]", TSThreadSelf());
  if (thread_2 == nullptr) {
    thread_2 = TSEventThreadSelf();
  } else if (thread_2 == TSEventThreadSelf()) {
    Dbg(dbg_ctl_chk, "pass [should be the same thread]");
  } else {
    Dbg(dbg_ctl_chk, "fail [not the same thread]");
  }
  return 0;
}

void
TSContScheduleOnThread_test()
{
  contp_1 = TSContCreate(TSContScheduleOnThread_handler_1, TSMutexCreate());
  contp_2 = TSContCreate(TSContScheduleOnThread_handler_2, TSMutexCreate());

  if (contp_1 == nullptr || contp_2 == nullptr) {
    Dbg(dbg_ctl_schd, "[%s] could not create continuation", plugin_name);
    abort();
  } else {
    Dbg(dbg_ctl_schd, "[%s] scheduling continuation", plugin_name);
    TSContScheduleOnPool(contp_1, 0, TS_THREAD_POOL_NET);
    TSContThreadAffinityClear(contp_1);
    TSContScheduleOnPool(contp_1, 200, TS_THREAD_POOL_NET);
  }
}

static int
TSContThreadAffinity_handler(TSCont contp, TSEvent event, void *edata)
{
  Dbg(dbg_ctl_hdl, "TSContThreadAffinity handler thread [%p]", TSThreadSelf());

  thread_1 = TSEventThreadSelf();

  if (TSContThreadAffinityGet(contp) != nullptr) {
    Dbg(dbg_ctl_chk, "pass [affinity thread is not null]");
    TSContThreadAffinityClear(contp);
    if (TSContThreadAffinityGet(contp) == nullptr) {
      Dbg(dbg_ctl_chk, "pass [affinity thread is cleared]");
      TSContThreadAffinitySet(contp, TSEventThreadSelf());
      if (TSContThreadAffinityGet(contp) == thread_1) {
        Dbg(dbg_ctl_chk, "pass [affinity thread is set]");
      } else {
        Dbg(dbg_ctl_chk, "fail [affinity thread is not set]");
      }
    } else {
      Dbg(dbg_ctl_chk, "fail [affinity thread is not cleared]");
    }
  } else {
    Dbg(dbg_ctl_chk, "fail [affinity thread is null]");
  }

  return 0;
}

void
TSContThreadAffinity_test()
{
  TSCont contp = TSContCreate(TSContThreadAffinity_handler, TSMutexCreate());

  if (contp == nullptr) {
    Dbg(dbg_ctl_schd, "[%s] could not create continuation", plugin_name);
    abort();
  } else {
    Dbg(dbg_ctl_schd, "[%s] scheduling continuation", plugin_name);
    TSContScheduleOnPool(contp, 0, TS_THREAD_POOL_NET);
  }
}

static int
LifecycleHookTracer(TSCont /* contp ATS_UNUSED */, TSEvent event, void * /* edata ATS_UNUSED */)
{
  if (event == TS_EVENT_LIFECYCLE_TASK_THREADS_READY) {
    switch (test_flag) {
    case 1:
      TSContScheduleOnPool_test();
      break;
    case 2:
      TSContScheduleOnThread_test();
      break;
    case 3:
      TSContThreadAffinity_test();
      break;
    default:
      break;
    }
  }
  return 0;
}

void
TSPluginInit(int argc, const char *argv[])
{
  if (argc == 2) {
    int len = strlen(argv[1]);
    if (len == 4 && strncmp(argv[1], "pool", 4) == 0) {
      Dbg(dbg_ctl_init, "initializing plugin for testing TSContScheduleOnPool");
      test_flag = 1;
    } else if (len == 6 && strncmp(argv[1], "thread", 6) == 0) {
      Dbg(dbg_ctl_init, "initializing plugin for testing TSContScheduleOnThread");
      test_flag = 2;
    } else if (len == 8 && strncmp(argv[1], "affinity", 8) == 0) {
      Dbg(dbg_ctl_init, "initializing plugin for testing TSContThreadAffinity");
      test_flag = 3;
    } else {
      goto Lerror;
    }
  } else {
    goto Lerror;
  }

  TSPluginRegistrationInfo info;

  info.plugin_name   = plugin_name;
  info.vendor_name   = vendor_name;
  info.support_email = support_email;

  if (TSPluginRegister(&info) != TS_SUCCESS) {
    Dbg(dbg_ctl_init, "[%s] plugin registration failed", plugin_name);
    abort();
  }

  TSLifecycleHookAdd(TS_LIFECYCLE_TASK_THREADS_READY_HOOK, TSContCreate(LifecycleHookTracer, TSMutexCreate()));

  return;

Lerror:
  Dbg(dbg_ctl_init, "[%s] plugin invalid argument", plugin_name);
  abort();
}
