// SPDX-FileCopyrightText: Copyright (c) Kitware Inc.
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkMPIMToNSocketConnection.h"

#include "vtkClientSocket.h"
#include "vtkMPIMToNSocketConnectionPortInformation.h"
#include "vtkMultiProcessController.h"
#include "vtkObjectFactory.h"
#include "vtkProcessModule.h"
#include "vtkRemotingCoreConfiguration.h"
#include "vtkServerSocket.h"
#include "vtkSocketCommunicator.h"

#include <cassert>
#include <iostream>
#include <string>
#include <vector>

vtkStandardNewMacro(vtkMPIMToNSocketConnection);

vtkCxxSetObjectMacro(vtkMPIMToNSocketConnection, Controller, vtkMultiProcessController);
vtkCxxSetObjectMacro(vtkMPIMToNSocketConnection, SocketCommunicator, vtkSocketCommunicator);
class vtkMPIMToNSocketConnectionInternals
{
public:
  struct NodeInformation
  {
    int PortNumber;
    std::string HostName;
  };
  std::vector<NodeInformation> ServerInformation;

  std::string SelfHostName;
};

vtkMPIMToNSocketConnection::vtkMPIMToNSocketConnection()
{
  this->Socket = 0;
  this->PortNumber = 0;
  this->Internals = new vtkMPIMToNSocketConnectionInternals;
  this->Controller = nullptr;
  this->SetController(vtkMultiProcessController::GetGlobalController());
  this->SocketCommunicator = nullptr;
  this->NumberOfConnections = -1;
  this->ServerSocket = nullptr;
  this->IsWaiting = false;
}

vtkMPIMToNSocketConnection::~vtkMPIMToNSocketConnection()
{
  if (this->ServerSocket)
  {
    this->ServerSocket->Delete();
    this->ServerSocket = nullptr;
  }
  if (this->SocketCommunicator)
  {
    this->SocketCommunicator->CloseConnection();
    this->SocketCommunicator->Delete();
  }
  this->SetController(nullptr);
  delete this->Internals;
  this->Internals = nullptr;
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);

  os << indent << "SelfHostName: " << this->Internals->SelfHostName.c_str() << endl;

  os << indent << "NumberOfConnections: (" << this->NumberOfConnections << ")\n";
  os << indent << "Controller: (" << this->Controller << ")\n";
  os << indent << "Socket: (" << this->Socket << ")\n";
  os << indent << "SocketCommunicator: (" << this->SocketCommunicator << ")\n";
  vtkIndent i2 = indent.GetNextIndent();
  for (unsigned int i = 0; i < this->Internals->ServerInformation.size(); ++i)
  {
    os << i2 << "Server Information for Process: " << i << ":\n";
    vtkIndent i3 = i2.GetNextIndent();
    os << i3 << "PortNumber: " << this->Internals->ServerInformation[i].PortNumber << "\n";
    os << i3 << "HostName: " << this->Internals->ServerInformation[i].HostName.c_str() << "\n";
  }
  os << indent << "PortNumber: " << this->PortNumber << endl;
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::Initialize(int waiting_process_type)
{
  auto pm = vtkProcessModule::GetProcessModule();
  this->SetController(pm->GetGlobalController());
  this->Internals->SelfHostName = vtkRemotingCoreConfiguration::GetInstance()->GetHostName();
  this->IsWaiting = (waiting_process_type == pm->GetProcessType());
  if (this->IsWaiting)
  {
    this->SetupWaitForConnection();
  }
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::ConnectMtoN()
{
  std::cout << "ConnectMtoN" << endl;
  if (this->IsWaiting)
  {
    this->WaitForConnection();
  }
  else
  {
    this->Connect();
  }
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::SetupWaitForConnection()
{
  if (this->SocketCommunicator)
  {
    vtkErrorMacro("SetupWaitForConnection called more than once");
    return;
  }
  unsigned int myId = this->Controller->GetLocalProcessId();
  if (myId >= (unsigned int)this->NumberOfConnections)
  {
    return;
  }
  this->SocketCommunicator = vtkSocketCommunicator::New();
  // open a socket on a random port
  vtkDebugMacro(<< "open with port " << this->PortNumber);
  this->ServerSocket = vtkServerSocket::New();
  this->ServerSocket->CreateServer(this->PortNumber);

  // find out the random port picked
  this->PortNumber = this->ServerSocket->GetServerPort();
  if (this->NumberOfConnections == -1)
  {
    this->NumberOfConnections = this->Controller->GetNumberOfProcesses();
  }
  std::cout.flush();
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::WaitForConnection()
{
  unsigned int myId = this->Controller->GetLocalProcessId();
  if (myId >= static_cast<unsigned int>(this->NumberOfConnections))
  {
    return;
  }
  if (!this->SocketCommunicator || !this->ServerSocket)
  {
    vtkErrorMacro("SetupWaitForConnection must be called before WaitForConnection");
    return;
  }
  std::cout << "Waiting for connection:"
            << " rank :" << myId << " host :" << this->Internals->SelfHostName.c_str()
            << " port :" << this->PortNumber << "\n";

  vtkClientSocket* socket = this->ServerSocket->WaitForConnection();
  this->ServerSocket->Delete();
  this->ServerSocket = nullptr;
  if (!socket)
  {
    vtkErrorMacro("Failed to get connection!");
    return;
  }
  this->SocketCommunicator->SetSocket(socket);
  this->SocketCommunicator->ServerSideHandshake();
  socket->Delete();

  int data;
  this->SocketCommunicator->Receive(&data, 1, 1, 1238);
  std::cout << "Received Hello from process " << data << "\n";
  std::cout.flush();
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::Connect()
{
  if (this->SocketCommunicator)
  {
    vtkErrorMacro("Connect called more than once");
    return;
  }
  unsigned int myId = this->Controller->GetLocalProcessId();
  if (myId >= this->Internals->ServerInformation.size())
  {
    return;
  }

  this->SocketCommunicator = vtkSocketCommunicator::New();

  const vtkMPIMToNSocketConnectionInternals::NodeInformation& targetNode =
    this->Internals->ServerInformation[myId];

  std::cout << "Connecting :"
            << " rank :" << myId << " dest-host :" << targetNode.HostName.c_str()
            << " dest-port :" << targetNode.PortNumber << endl;

  this->SocketCommunicator->ConnectTo(
    const_cast<char*>(targetNode.HostName.c_str()), targetNode.PortNumber);

  int id = static_cast<int>(myId);
  this->SocketCommunicator->Send(&id, 1, 1, 1238);
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::SetNumberOfConnections(int c)
{
  this->NumberOfConnections = c;
  this->Internals->ServerInformation.resize(this->NumberOfConnections);
  this->Modified();
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::SetPortInformation(
  unsigned int processNumber, int port, const char* host)
{
  if (processNumber >= this->Internals->ServerInformation.size())
  {
    vtkErrorMacro("Attempt to set port information for process larger than number of processes.\n"
      << "Max process id " << this->Internals->ServerInformation.size() << " attempted "
      << processNumber << "\n");
    return;
  }
  this->Internals->ServerInformation[processNumber].PortNumber = port;
  if (host)
  {
    this->Internals->ServerInformation[processNumber].HostName = host;
  }
}

//------------------------------------------------------------------------------
void vtkMPIMToNSocketConnection::GetPortInformation(vtkMPIMToNSocketConnectionPortInformation* info)
{
  // This method must be called only on the "Waiting" process.
  if (this->NumberOfConnections == -1)
  {
    return;
  }

  // Pass the current processes information.
  info->SetConnectionInformation(
    this->Controller->GetLocalProcessId(), this->PortNumber, this->Internals->SelfHostName.c_str());
}
