/**
 * Copyright 2015 Google Inc. All Rights Reserved.
 *
 * 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
 *
 *      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.
 */

// Ensure that Python.h is included before any other header.
#include "common.h"

#include "bytecode_breakpoint.h"

#include "bytecode_manipulator.h"
#include "python_util.h"
#include "linetable/linetable.h"
#include "codeobject/codeobject.h"

namespace devtools {


namespace cdbg {

// Each method in python has a tuple with all the constants instructions use.
// Breakpoint patching appends more constants. If the index of new constant
// exceed 0xFFFF, breakpoint patching would need to use extended instructions,
// which is not supported. We therefore limit to methods with up to 0xF000
// instructions that leaves us with up to 0x0FFF breakpoints.
static const int kMaxCodeObjectConsts = 0xF000;

BytecodeBreakpoint::BytecodeBreakpoint()
    : cookie_counter_(1000000) {
}


BytecodeBreakpoint::~BytecodeBreakpoint() {
  Detach();
}


void BytecodeBreakpoint::Detach() {
  for (auto it = patches_.begin(); it != patches_.end(); ++it) {
    it->second->breakpoints.clear();
    PatchCodeObject(it->second, nullptr);

    // TODO(vlif): assert zombie_refs.empty() after garbage collection
    // for zombie refs is implemented.

    delete it->second;
  }

  patches_.clear();

  for (auto it = cookie_map_.begin(); it != cookie_map_.end(); ++it) {
    delete it->second;
  }

  cookie_map_.clear();
}


int BytecodeBreakpoint::SetBreakpoint(
    PyCodeObject* code_object,
    int line,
    PyObject* hit_callback,
    std::function<void()> error_callback,
    PyObject** o_new_co_code)
{
    CodeObjectBreakpoints* code_object_breakpoints =
      PreparePatchCodeObject(ScopedPyCodeObject::NewReference(code_object));

    if (code_object_breakpoints == nullptr) {
      error_callback();
      return -1;  // Not a valid cookie, but "ClearBreakpoint" wouldn't mind.
    }

  // Find the offset of the instruction at "line". We use original line
  // table in case "code_object" is already patched with another breakpoint.
  int offset = linetable::GetLineOffset(line, &code_object_breakpoints->original_lnotab, code_object);
    if (-1 == offset) {
        error_callback();
        return -1;
    }

  // Assign cookie to this breakpoint and Register it.
  const int cookie = cookie_counter_++;

  std::unique_ptr<Breakpoint> breakpoint(new Breakpoint);
  breakpoint->code_object = ScopedPyCodeObject::NewReference(code_object);
  breakpoint->line = line;
  breakpoint->offset = offset;
  breakpoint->hit_callable = hit_callback;
  breakpoint->error_callback = error_callback;
  breakpoint->cookie = cookie;

  code_object_breakpoints->breakpoints.insert(
      std::make_pair(breakpoint->offset, breakpoint.get()));

  DEBUG("Number of breakpoints after insertion to db is %zu\n", code_object_breakpoints->breakpoints.size());
  DCHECK(cookie_map_[cookie] == nullptr);
  cookie_map_[cookie] = breakpoint.release();

  PatchCodeObject(code_object_breakpoints, o_new_co_code);

  return cookie;
}


void BytecodeBreakpoint::ClearBreakpoint(int cookie, PyObject** o_new_co_code) {
  auto it_breakpoint = cookie_map_.find(cookie);
  if (it_breakpoint == cookie_map_.end()) {
    return;  // No breakpoint with this cookie.
  }

  // Previously, hit_callable d id not refer to the user's callback directly --
  // but to a callback generated by this extension, which the extension would use to
  // disable the breakpoint here - so that even if the existing bytecode is used after disabling
  // (e.g. in the case of a long-running function),
  // the breakpoint would immediately stop being hit.
  // This has been replaced with functionality implemented in google_bdb._callback in order to allow
  // the modified bytecode to be pickled (for Spark support).

  auto it_code = patches_.find(it_breakpoint->second->code_object);
  if (it_code != patches_.end()) {
    CodeObjectBreakpoints* code = it_code->second;

    auto it = code->breakpoints.begin();
    int erase_count = 0;
    while (it != code->breakpoints.end()) {
      if (it->second == it_breakpoint->second) {
        code->breakpoints.erase(it);
        ++erase_count;
        it = code->breakpoints.begin();
      } else {
        ++it;
      }
    }

    DCHECK_EQ(1, erase_count);

    PatchCodeObject(code, o_new_co_code);
    DEBUG("Returned from PatchedCodeObject\n");

    if (code->breakpoints.empty() && code->zombie_refs.empty()) {
      delete it_code->second;
      patches_.erase(it_code);
    }
  } else {
    DCHECK(false) << "Missing code object";
  }

  delete it_breakpoint->second;
  cookie_map_.erase(it_breakpoint);
}


BytecodeBreakpoint::CodeObjectBreakpoints*
BytecodeBreakpoint::PreparePatchCodeObject(
    const ScopedPyCodeObject& code_object) {

    if (code_object.is_null() || !PyCode_Check(code_object.get())) {
    DEBUG("Bad code_object argument");
    return nullptr;
  }

    auto it = patches_.find(code_object);
  if (it != patches_.end()) {
      DEBUG("Found patches entry\n");
      DEBUG("co_consts size is %zd\n", PyTuple_GET_SIZE(it->second->original_consts.get()));
      DEBUG("co_consts is %s\n", PrintTuple(it->second->original_consts.get()).c_str());
      return it->second;  // Already loaded.
  }

  std::unique_ptr<CodeObjectBreakpoints> data(new CodeObjectBreakpoints);
  data->code_object = code_object;
  data->original_stacksize = code_object.get()->co_stacksize;

  data->original_consts =
      ScopedPyObject::NewReference(code_object.get()->co_consts);
  if ((data->original_consts == nullptr) ||
      !PyTuple_CheckExact(data->original_consts.get())) {
      DEBUG("Code object has null or corrupted constants tuple");
    return nullptr;
  }
    DEBUG("co_consts size is %zd\n", PyTuple_GET_SIZE(data->original_consts.get()));
    DEBUG("co_consts is %s\n", PrintTuple(data->original_consts.get()).c_str());

    if (PyTuple_GET_SIZE(data->original_consts.get()) >= kMaxCodeObjectConsts) {
        DEBUG("Code object with more than %d constants not supported\n", kMaxCodeObjectConsts);

      return nullptr;
  }

    data->original_code = ScopedPyObject::NewReference(codeobject::GetCoCode(code_object.get()));
  if ((data->original_code == nullptr) ||
      !PyBytes_CheckExact(data->original_code.get())) {
     DEBUG("Code object has no code");
      return nullptr;  // Probably a built-in method or uninitialized code object.
  }

  data->original_lnotab = ScopedPyObject::NewReference(linetable::GetLinetable(code_object.get()));

#if PY_VERSION_HEX >= 0x030B0000
    DEBUG("BytecodeBreakpoint::PreparePatchCodeObject: Changing original exception table\n");
    data->original_exception_table =
            ScopedPyObject::NewReference(code_object.get()->co_exceptiontable);
#endif

  patches_[code_object] = data.get();
    return data.release();
}


void BytecodeBreakpoint::PatchCodeObject(CodeObjectBreakpoints* code, PyObject** o_new_co_code) {
  PyCodeObject* code_object = code->code_object.get();

  if (code->breakpoints.empty()) {
      DEBUG("No breakpoints\n");
    code->zombie_refs.push_back(ScopedPyObject(code_object->co_consts));
    code_object->co_consts = code->original_consts.get();
    Py_INCREF(code_object->co_consts);


    code_object->co_stacksize = code->original_stacksize;

    PyObject *co_code = codeobject::GetCoCode(code_object);
    code->zombie_refs.push_back(ScopedPyObject(co_code));

#if PY_VERSION_HEX >= 0x030B0000
    if(o_new_co_code != nullptr) {
        *o_new_co_code =  code->original_code.get();
        Py_INCREF(*o_new_co_code);
    } else {
        DDEBUG("o_new_co_code is null!!!\n");
    }
#else
    DEBUG("Entering SetCoCode with code->original_code as param\n");
    codeobject::SetCoCode(code_object, code->original_code);
    DEBUG("Returned from SetCoCode with code->original_code as param\n");
#endif
    DEBUG("Code object %s reverted to %p from patched %p\n", CodeObjectDebugString(code_object).c_str(),
          co_code, code->zombie_refs.back().get());
      Py_INCREF(co_code);

    PyObject *line_table = linetable::GetLinetable(code_object);
    if (line_table != nullptr) {
      code->zombie_refs.push_back(ScopedPyObject(line_table));
    }
    line_table = code->original_lnotab.get();
#if PY_VERSION_HEX >= 0x030A0000
    code_object->co_linetable = line_table;
    Py_INCREF(code_object->co_linetable);
#else
    code_object->co_lnotab = line_table;
    Py_INCREF(code_object->co_lnotab);
#endif

#if PY_VERSION_HEX >= 0x030B0000
    DEBUG("Reverting exception table\n");
    // TODO: Why aren't there checks here like for line_table above?
    code_object->co_exceptiontable = code->original_exception_table.get();
    Py_INCREF(code_object->co_exceptiontable);
#endif

    return;
  }

  std::vector<uint8> bytecode = PyBytesToByteArray(code->original_code.get());

  bool has_lnotab = false;
  std::vector<uint8> lnotab;
  if (!code->original_lnotab.is_null() &&
      PyBytes_CheckExact(code->original_lnotab.get())) {
    has_lnotab = true;
    DDEBUG("has_lnotab = true\n");
    lnotab = PyBytesToByteArray(code->original_lnotab.get());
  }

  bool has_exception_table = false;
  std::vector<uint8> exception_table;
  if (!code->original_exception_table.is_null() &&
          PyBytes_CheckExact(code->original_exception_table.get())) {
      has_exception_table = true;
      DDEBUG("has_exception_table = true\n");
      exception_table = PyBytesToByteArray(code->original_exception_table.get());
  }

  BytecodeManipulator bytecode_manipulator(
      std::move(bytecode),
      has_lnotab,
      std::move(lnotab),
      has_exception_table,
      std::move(exception_table),
      code_object);

  // Add callbacks to code object constants and patch the bytecode.
  std::vector<PyObject*> callbacks;
  DDEBUG("code->breakpoints.size() = %zu\n", code->breakpoints.size());
  callbacks.reserve(code->breakpoints.size());

  std::vector<std::function<void()>> errors;

  int const_index = PyTuple_GET_SIZE(code->original_consts.get());
  for (auto it_entry = code->breakpoints.begin();
       it_entry != code->breakpoints.end();
       ++it_entry, ++const_index)
  {
    int offset = it_entry->first;
    bool offset_found = true;
    const Breakpoint& breakpoint = *it_entry->second;
    DCHECK_EQ(offset, breakpoint.offset);
    DEBUG("Add a new breakpoint (%d) at offset: %d, for line; %d\n", breakpoint.cookie, offset, breakpoint.line);

    callbacks.push_back(breakpoint.hit_callable);

#if PY_VERSION_HEX >= 0x03060000
      // In Python 3, since we allow upgrading of instructions to use
    // EXTENDED_ARG, the offsets for lines originally calculated might not be
    // accurate, so we need to recalculate them each insertion.
    offset_found = false;
    if (bytecode_manipulator.has_lnotab())
    {
      ScopedPyObject lnotab(PyBytes_FromStringAndSize(
          reinterpret_cast<const char*>(bytecode_manipulator.lnotab().data()),
          bytecode_manipulator.lnotab().size()));

      offset = linetable::GetLineOffset(breakpoint.line, &lnotab, code_object);
      lnotab.release();
      offset_found = (-1 != offset);
      }
#endif // PY_VERSION_HEX >= 0x03060000

    DDDEBUG("Const_index = %d\n", const_index);
    if (!offset_found ||
        !bytecode_manipulator.InjectMethodCall(offset, const_index)) {
        DEBUG("Failed to insert bytecode for breakpoint %d at line %d\n", breakpoint.cookie, breakpoint.line);
      errors.push_back(breakpoint.error_callback);
    }
  }

  // Create the constants tuple, the new bytecode string and line table.
  code->zombie_refs.push_back(ScopedPyObject(code_object->co_consts));
  ScopedPyObject consts = AppendTuple(code->original_consts.get(), callbacks);
  code_object->co_consts = consts.release();

  code_object->co_stacksize = code->original_stacksize + 1;

  code->zombie_refs.push_back(ScopedPyObject(codeobject::GetCoCode(code_object)));
  DDEBUG("Return 2 from GetCoCode\n");
  ScopedPyObject bytecode_string(PyBytes_FromStringAndSize(
      reinterpret_cast<const char*>(bytecode_manipulator.bytecode().data()),
      bytecode_manipulator.bytecode().size()));
  DCHECK(!bytecode_string.is_null());

#if PY_VERSION_HEX >= 0x030B0000
    if(o_new_co_code != nullptr) {
        *o_new_co_code = bytecode_string.release();
    } else {
        DDEBUG("o_new_co_code is null!\n");
    }
#else
        codeobject::SetCoCode(code_object, bytecode_string);
        bytecode_string.release();
#endif

    DDEBUG("co_consts: %s -> %s\n",
           PrintTuple(code->original_consts.get()).c_str(),
           PrintTuple(code_object->co_consts).c_str());
#if PY_VERSION_HEX >= 0x030B0000
    DDEBUG("co_exceptiontable: %p -> %p\n", code->original_exception_table.get(), code_object->co_exceptiontable);
#endif
    DDEBUG("co_stacksize: %d -> %d\n", code->original_stacksize, code_object->co_stacksize);
    DDEBUG("co_linetable: %p -> %p\n", code->original_lnotab.get(), linetable::GetLinetable(code_object));
    DDEBUG("co_code: %p -> %p\n", code->original_code.get(), codeobject::GetCoCode(code_object));

  if (has_lnotab) {
      code->zombie_refs.push_back(ScopedPyObject(linetable::GetLinetable(code_object)));
      ScopedPyObject lnotab_string(PyBytes_FromStringAndSize(
            reinterpret_cast<const char*>(bytecode_manipulator.lnotab().data()),
            bytecode_manipulator.lnotab().size()));
      DCHECK(!lnotab_string.is_null());
      linetable::UpdateLinetable(lnotab_string, code_object);
  }

  if (has_exception_table) {
#if PY_VERSION_HEX >= 0x030B0000
      code->zombie_refs.push_back(ScopedPyObject(code_object->co_exceptiontable));
      ScopedPyObject exception_table_string(PyBytes_FromStringAndSize(
              reinterpret_cast<const char *>(bytecode_manipulator.exception_table().data()),
              bytecode_manipulator.exception_table().size()));
      DCHECK(!exception_table_string.is_null());
      code_object->co_exceptiontable = exception_table_string.release();
#endif
  }

  // Invoke error callback after everything else is done. The callback may
  // decide to remove the breakpoint, which will change "code".
  for (auto it = errors.begin(); it != errors.end(); ++it) {
    (*it)();
  }
}



}  // namespace cdbg
}  // namespace devtools
