/**
 ******************************************************************************
 * Xenia : Xbox 360 Emulator Research Project                                 *
 ******************************************************************************
 * Copyright 2013 Ben Vanik. All rights reserved.                             *
 * Released under the BSD license - see LICENSE in the root for more details. *
 ******************************************************************************
 */

#include "xenia/debug/debugger.h"

#include <gflags/gflags.h>
#include <mstcpip.h>
#include <winsock2.h>
#include <ws2tcpip.h>

#include <mutex>

#include "xenia/base/logging.h"
#include "xenia/base/string.h"
#include "xenia/base/threading.h"
#include "xenia/cpu/function.h"
#include "xenia/cpu/processor.h"
#include "xenia/emulator.h"

// Autogenerated Flatbuffers files:
#include "xenia/debug/proto/breakpoints_generated.h"
#include "xenia/debug/proto/common_generated.h"
#include "xenia/debug/proto/control_generated.h"
#include "xenia/debug/proto/messages_generated.h"
#include "xenia/debug/proto/modules_generated.h"

DEFINE_string(debug_session_path, "", "Debug output path.");
DEFINE_int32(debug_port, 19000, "Port the debugger listens on.");
DEFINE_bool(wait_for_debugger, false,
            "Waits for the debugger to attach before starting the game.");
DEFINE_bool(exit_with_debugger, true, "Exit whe the debugger disconnects.");

namespace xe {
namespace debug {

using xe::cpu::ThreadState;

Breakpoint::Breakpoint(Type type, uint32_t address)
    : type_(type), address_(address) {}

Breakpoint::~Breakpoint() = default;

Debugger::Debugger(Emulator* emulator)
    : emulator_(emulator),
      listen_socket_(INVALID_SOCKET),
      client_socket_(INVALID_SOCKET) {
  WSADATA wsa_data;
  WSAStartup(MAKEWORD(2, 2), &wsa_data);
}

Debugger::~Debugger() {
  if (listen_socket_ != INVALID_SOCKET) {
    StopSession();
  }
}

bool Debugger::StartSession() {
  std::wstring session_path = xe::to_wstring(FLAGS_debug_session_path);

  functions_path_ = xe::join_paths(session_path, L"functions");
  functions_file_ =
      ChunkedMappedMemoryWriter::Open(functions_path_, 32 * 1024 * 1024, false);

  functions_trace_path_ = xe::join_paths(session_path, L"functions.trace");
  functions_trace_file_ = ChunkedMappedMemoryWriter::Open(
    functions_trace_path_, 32 * 1024 * 1024, true);

  listen_socket_ = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
  if (listen_socket_ < 1) {
    XELOGE("Failed to create debugger socket");
    return false;
  }
  struct tcp_keepalive alive;
  alive.onoff = TRUE;
  alive.keepalivetime = 7200000;
  alive.keepaliveinterval = 6000;
  DWORD bytes_returned;
  WSAIoctl(listen_socket_, SIO_KEEPALIVE_VALS, &alive, sizeof(alive), nullptr,
           0, &bytes_returned, nullptr, nullptr);
  int opt_value = 1;
  setsockopt(listen_socket_, SOL_SOCKET, SO_REUSEADDR,
             reinterpret_cast<const char*>(&opt_value), sizeof(opt_value));
  opt_value = 1;
  setsockopt(listen_socket_, IPPROTO_TCP, TCP_NODELAY,
             reinterpret_cast<const char*>(&opt_value), sizeof(opt_value));

  sockaddr_in socket_addr = {0};
  socket_addr.sin_family = AF_INET;
  socket_addr.sin_addr.s_addr = htonl(INADDR_ANY);
  socket_addr.sin_port = htons(FLAGS_debug_port);
  if (bind(listen_socket_, reinterpret_cast<sockaddr*>(&socket_addr),
           sizeof(socket_addr)) == SOCKET_ERROR) {
    int e = WSAGetLastError();
    XELOGE("Unable to bind debug socket");
    return false;
  }
  if (listen(listen_socket_, 5) == SOCKET_ERROR) {
    XELOGE("Unable to listen on debug socket");
    return 1;
  }

  return true;
}

void SendResponse(SOCKET client_socket, flatbuffers::FlatBufferBuilder& fbb,
                  uint32_t request_id, proto::ResponseData response_data_type,
                  flatbuffers::Offset<void> response_data_offset) {
  auto response_offset = proto::CreateResponse(
      fbb, request_id, response_data_type, response_data_offset);
  fbb.Finish(response_offset);
  int buffer_length = fbb.GetSize();
  send(client_socket, reinterpret_cast<const char*>(&buffer_length), 4, 0);
  send(client_socket, reinterpret_cast<const char*>(fbb.GetBufferPointer()),
       fbb.GetSize(), 0);
}

void Debugger::PreLaunch() {
  accept_thread_ = std::thread([this]() {
    while (listen_socket_ != INVALID_SOCKET) {
      sockaddr_in6 client_addr;
      int client_count = sizeof(client_addr);
      SOCKET client_socket_id =
          accept(listen_socket_, reinterpret_cast<sockaddr*>(&client_addr),
                 &client_count);
      if (client_socket_id == INVALID_SOCKET) {
        XELOGE("Failed to accept socket");
        continue;
      }

      // Only one debugger at a time.
      if (client_socket_ != INVALID_SOCKET) {
        XELOGW("Ignoring debugger connection as one is already connected");
        closesocket(client_socket_id);
        continue;
      }

      // Setup recv thread.
      client_socket_ = client_socket_id;
      receive_thread_ = std::thread([this]() {
        while (client_socket_ != INVALID_SOCKET) {
          // Read length prefix.
          uint32_t length = 0;
          int r = recv(client_socket_, reinterpret_cast<char*>(&length), 4,
                       MSG_WAITALL);
          if (r != 4) {
            // Failed?
            XELOGE("Failed to recv debug data length - dead connection?");
            if (FLAGS_exit_with_debugger) {
              exit(1);
            }
            break;
          }

          // Read body.
          std::vector<uint8_t> body(length);
          r = recv(client_socket_, reinterpret_cast<char*>(body.data()), length,
                   MSG_WAITALL);
          if (r != length) {
            // Failed?
            XELOGE("Failed to recv debug data body - dead connection?");
            if (FLAGS_exit_with_debugger) {
              exit(1);
            }
            break;
          }

          // Read message contents and dispatch.
          OnMessage(std::move(body));
        }
      });

      // This will WaitForClient if it was waiting.
    }
  });

  if (FLAGS_wait_for_debugger) {
    // Wait for the first client.
    XELOGI("Waiting for debugger because of --wait_for_debugger...");
    accept_fence_.Wait();
    XELOGI("Debugger attached, continuing...");
  }
}

void Debugger::OnMessage(std::vector<uint8_t> buffer) {
  flatbuffers::FlatBufferBuilder fbb;
  proto::ResponseData response_data_type;
  flatbuffers::Offset<void> response_data_offset;

  auto request = flatbuffers::GetRoot<proto::Request>(buffer.data());
  switch (request->request_data_type()) {
    case proto::RequestData_AttachRequest: {
      // Send debug info.
      response_data_type = proto::ResponseData_AttachResponse;
      response_data_offset =
          proto::CreateAttachResponse(
              fbb, fbb.CreateString(
                       xe::to_string(emulator()->memory()->file_name())),
              fbb.CreateString(xe::to_string(functions_path_)),
              fbb.CreateString(xe::to_string(functions_trace_path_))).Union();

      // Allow continuation if we were blocked waiting for a client.
      accept_fence_.Signal();
    } break;

    case proto::RequestData_ListBreakpointsRequest: {
      response_data_type = proto::ResponseData_ListBreakpointsResponse;
      auto response_data = proto::ListBreakpointsResponseBuilder(fbb);
      //
      response_data_offset = response_data.Finish().Union();
    } break;
    case proto::RequestData_AddBreakpointsRequest: {
      response_data_type = proto::ResponseData_RemoveBreakpointsResponse;
      auto response_data = proto::AddBreakpointsResponseBuilder(fbb);
      //
      response_data_offset = response_data.Finish().Union();
    } break;
    case proto::RequestData_UpdateBreakpointsRequest: {
      response_data_type = proto::ResponseData_UpdateBreakpointsResponse;
      auto response_data = proto::UpdateBreakpointsResponseBuilder(fbb);
      //
      response_data_offset = response_data.Finish().Union();
    } break;
    case proto::RequestData_RemoveBreakpointsRequest: {
      response_data_type = proto::ResponseData_AddBreakpointsResponse;
      auto response_data = proto::RemoveBreakpointsResponseBuilder(fbb);
      //
      response_data_offset = response_data.Finish().Union();
    } break;

    case proto::RequestData_ListModulesRequest: {
      response_data_type = proto::ResponseData_ListModulesResponse;
      auto response_data = proto::ListModulesResponseBuilder(fbb);
      //
      response_data_offset = response_data.Finish().Union();
    } break;
    case proto::RequestData_GetModuleRequest: {
      response_data_type = proto::ResponseData_GetModuleResponse;
      auto response_data = proto::GetModuleResponseBuilder(fbb);
      //
      response_data_offset = response_data.Finish().Union();
    } break;

    case proto::RequestData_StopRequest: {
      response_data_type = proto::ResponseData_StopResponse;
      auto response_data = proto::StopResponseBuilder(fbb);
      // TODO(benvanik): gracefully die?
      exit(1);
      response_data_offset = response_data.Finish().Union();
    } break;
    case proto::RequestData_BreakRequest: {
      response_data_type = proto::ResponseData_BreakResponse;
      auto response_data = proto::BreakResponseBuilder(fbb);
      //
      SuspendAllThreads();
      response_data_offset = response_data.Finish().Union();
    } break;
    case proto::RequestData_ContinueRequest: {
      response_data_type = proto::ResponseData_ContinueResponse;
      auto response_data = proto::ContinueResponseBuilder(fbb);
      //
      ResumeAllThreads();
      response_data_offset = response_data.Finish().Union();
    } break;
    case proto::RequestData_StepRequest: {
      response_data_type = proto::ResponseData_StepResponse;
      auto response_data = proto::StepResponseBuilder(fbb);
      //
      response_data_offset = response_data.Finish().Union();
    } break;

    default:
      assert_unhandled_case(request->request_data_type());
      break;
  }

  SendResponse(client_socket_, std::move(fbb), request->id(),
               response_data_type, response_data_offset);
}

void Debugger::StopSession() {
  FlushSession();

  if (client_socket_ != INVALID_SOCKET) {
    shutdown(client_socket_, SD_SEND);
    closesocket(client_socket_);
    client_socket_ = INVALID_SOCKET;
  }

  linger so_linger;
  so_linger.l_onoff = TRUE;
  so_linger.l_linger = 30;
  setsockopt(listen_socket_, SOL_SOCKET, SO_LINGER,
             reinterpret_cast<const char*>(&so_linger), sizeof(so_linger));
  shutdown(listen_socket_, SD_SEND);
  closesocket(listen_socket_);
  listen_socket_ = INVALID_SOCKET;
  if (accept_thread_.joinable()) {
    accept_thread_.join();
  }

  functions_file_.reset();
  functions_trace_file_.reset();
}

void Debugger::FlushSession() {
  if (functions_file_) {
    functions_file_->Flush();
  }
  if (functions_trace_file_) {
    functions_trace_file_->Flush();
  }
}

uint8_t* Debugger::AllocateFunctionData(size_t size) {
  if (!functions_file_) {
    return nullptr;
  }
  return functions_file_->Allocate(size);
}

uint8_t* Debugger::AllocateFunctionTraceData(size_t size) {
  if (!functions_trace_file_) {
    return nullptr;
  }
  return functions_trace_file_->Allocate(size);
}

int Debugger::SuspendAllThreads(uint32_t timeout_ms) {
  std::lock_guard<std::mutex> guard(threads_lock_);

  int result = 0;
  for (auto thread_state : threads_) {
    if (thread_state.second->Suspend(timeout_ms)) {
      result = 1;
    }
  }
  return result;
}

int Debugger::ResumeThread(uint32_t thread_id) {
  std::lock_guard<std::mutex> guard(threads_lock_);

  auto it = threads_.find(thread_id);
  if (it == threads_.end()) {
    return 1;
  }

  // Found thread. Note that it could be deleted as soon as we unlock.
  ThreadState* thread_state = it->second;
  int result = thread_state->Resume();

  return result;
}

int Debugger::ResumeAllThreads(bool force) {
  std::lock_guard<std::mutex> guard(threads_lock_);

  int result = 0;
  for (auto thread_state : threads_) {
    if (thread_state.second->Resume(force)) {
      result = 1;
    }
  }
  return result;
}

void Debugger::ForEachThread(std::function<void(ThreadState*)> callback) {
  std::lock_guard<std::mutex> guard(threads_lock_);

  for (auto thread_state : threads_) {
    callback(thread_state.second);
  }
}

int Debugger::AddBreakpoint(Breakpoint* breakpoint) {
  // Add to breakpoints map.
  {
    std::lock_guard<std::mutex> guard(breakpoints_lock_);
    breakpoints_.insert(
        std::pair<uint32_t, Breakpoint*>(breakpoint->address(), breakpoint));
  }

  // Find all functions that contain the breakpoint address.
  auto fns =
      emulator_->processor()->FindFunctionsWithAddress(breakpoint->address());

  // Add.
  for (auto fn : fns) {
    if (fn->AddBreakpoint(breakpoint)) {
      return 1;
    }
  }

  return 0;
}

int Debugger::RemoveBreakpoint(Breakpoint* breakpoint) {
  // Remove from breakpoint map.
  {
    std::lock_guard<std::mutex> guard(breakpoints_lock_);
    auto range = breakpoints_.equal_range(breakpoint->address());
    if (range.first == range.second) {
      return 1;
    }
    bool found = false;
    for (auto it = range.first; it != range.second; ++it) {
      if (it->second == breakpoint) {
        breakpoints_.erase(it);
        found = true;
        break;
      }
    }
    if (!found) {
      return 1;
    }
  }

  // Find all functions that have the breakpoint set.
  auto fns =
      emulator_->processor()->FindFunctionsWithAddress(breakpoint->address());

  // Remove.
  for (auto fn : fns) {
    fn->RemoveBreakpoint(breakpoint);
  }

  return 0;
}

void Debugger::FindBreakpoints(uint32_t address,
                               std::vector<Breakpoint*>& out_breakpoints) {
  std::lock_guard<std::mutex> guard(breakpoints_lock_);

  out_breakpoints.clear();

  auto range = breakpoints_.equal_range(address);
  if (range.first == range.second) {
    return;
  }

  for (auto it = range.first; it != range.second; ++it) {
    Breakpoint* breakpoint = it->second;
    out_breakpoints.push_back(breakpoint);
  }
}

void Debugger::OnThreadCreated(ThreadState* thread_state) {
  std::lock_guard<std::mutex> guard(threads_lock_);
  threads_[thread_state->thread_id()] = thread_state;
}

void Debugger::OnThreadDestroyed(ThreadState* thread_state) {
  std::lock_guard<std::mutex> guard(threads_lock_);
  auto it = threads_.find(thread_state->thread_id());
  if (it != threads_.end()) {
    threads_.erase(it);
  }
}

void Debugger::OnFunctionDefined(cpu::FunctionInfo* symbol_info,
                                 cpu::Function* function) {
  // Man, I'd love not to take this lock.
  std::vector<Breakpoint*> breakpoints;
  {
    std::lock_guard<std::mutex> guard(breakpoints_lock_);
    for (uint32_t address = symbol_info->address();
         address <= symbol_info->end_address(); address += 4) {
      auto range = breakpoints_.equal_range(address);
      if (range.first == range.second) {
        continue;
      }
      for (auto it = range.first; it != range.second; ++it) {
        Breakpoint* breakpoint = it->second;
        breakpoints.push_back(breakpoint);
      }
    }
  }

  if (breakpoints.size()) {
    // Breakpoints to add!
    for (auto breakpoint : breakpoints) {
      function->AddBreakpoint(breakpoint);
    }
  }
}

void Debugger::OnBreakpointHit(ThreadState* thread_state,
                               Breakpoint* breakpoint) {
  // Suspend all threads immediately.
  SuspendAllThreads();

  // Notify listeners.
  BreakpointHitEvent e(this, thread_state, breakpoint);
  breakpoint_hit(e);

  // Note that we stay suspended.
}

}  // namespace debug
}  // namespace xe