/*
 * Copyright 2000-2001 Sun Microsystems, Inc.  All Rights Reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Sun Microsystems, Inc., 4150 Network Circle, Santa Clara,
 * CA 95054 USA or visit www.sun.com if you need additional information or
 * have any questions.
 *
 */

#include "procList.hpp"
#include "nt4internals.hpp"
#include "isNT4.hpp"
#include "toolHelp.hpp"
#include <assert.h>

using namespace std;
using namespace NT4;

typedef void ProcListImplFunc(ProcEntryList& processes);

void procListImplNT4(ProcEntryList& processes);
void procListImplToolHelp(ProcEntryList& processes);

ProcEntry::ProcEntry(ULONG pid, USHORT nameLength, WCHAR* name) {
  this->pid = pid;
  this->nameLength = nameLength;
  this->name = new WCHAR[nameLength];
  memcpy(this->name, name, nameLength * sizeof(WCHAR));
}

ProcEntry::ProcEntry(ULONG pid, USHORT nameLength, char* name) {
  this->pid = pid;
  this->nameLength = nameLength;
  this->name = new WCHAR[nameLength];
  int j = 0;
  for (int i = 0; i < nameLength; i++) {
    // FIXME: what is the proper promotion from ASCII to UNICODE?
    this->name[i] = name[i] & 0xFF;
  }
}

ProcEntry::ProcEntry(const ProcEntry& arg) {
  name = NULL;
  copyFrom(arg);
}

ProcEntry&
ProcEntry::operator=(const ProcEntry& arg) {
  copyFrom(arg);
  return *this;
}

ProcEntry::~ProcEntry() {
  delete[] name;
}

void
ProcEntry::copyFrom(const ProcEntry& arg) {
  if (name != NULL) {
    delete[] name;
  }
  pid = arg.pid;
  nameLength = arg.nameLength;
  name = new WCHAR[nameLength];
  memcpy(name, arg.name, nameLength * sizeof(WCHAR));
}

ULONG
ProcEntry::getPid() {
  return pid;
}

USHORT
ProcEntry::getNameLength() {
  return nameLength;
}

WCHAR*
ProcEntry::getName() {
  return name;
}

void
procList(ProcEntryList& processes) {
  static ProcListImplFunc* impl = NULL;

  if (impl == NULL) {
    // See which operating system we're on
    impl = (isNT4() ? &procListImplNT4 : &procListImplToolHelp);
  }

  assert(impl != NULL);

  (*impl)(processes);
}

void
procListImplNT4(ProcEntryList& processes) {
  using namespace NT4;

  static ZwQuerySystemInformationFunc* query = NULL;

  if (query == NULL) {
    HMODULE ntDLL = loadNTDLL();
    query =
      (ZwQuerySystemInformationFunc*) GetProcAddress(ntDLL,
                                                     "ZwQuerySystemInformation");
    assert(query != NULL);
  }

  ULONG n = 0x100;
  PSYSTEM_PROCESSES sp = new SYSTEM_PROCESSES[n];
  while ((*query)(SystemProcessesAndThreadsInformation,
                  sp, n * sizeof(SYSTEM_PROCESSES), 0) == STATUS_INFO_LENGTH_MISMATCH) {
    delete[] sp;
    n *= 2;
    sp = new SYSTEM_PROCESSES[n];
  }

  bool done = false;
  for (PSYSTEM_PROCESSES p = sp; !done;
       p = PSYSTEM_PROCESSES(PCHAR(p) + p->NextEntryDelta)) {
    processes.push_back(ProcEntry(p->ProcessId,
                                  p->ProcessName.Length / 2,
                                  p->ProcessName.Buffer));
    done = p->NextEntryDelta == 0;
  }
}

void
procListImplToolHelp(ProcEntryList& processes) {
  using namespace ToolHelp;

  static CreateToolhelp32SnapshotFunc* snapshotFunc = NULL;
  static Process32FirstFunc*           firstFunc    = NULL;
  static Process32NextFunc*            nextFunc     = NULL;

  if (snapshotFunc == NULL) {
    HMODULE dll = loadDLL();

    snapshotFunc =
      (CreateToolhelp32SnapshotFunc*) GetProcAddress(dll,
                                                     "CreateToolhelp32Snapshot");

    firstFunc = (Process32FirstFunc*) GetProcAddress(dll,
                                                     "Process32First");

    nextFunc = (Process32NextFunc*) GetProcAddress(dll,
                                                   "Process32Next");

    assert(snapshotFunc != NULL);
    assert(firstFunc    != NULL);
    assert(nextFunc     != NULL);
  }

  HANDLE snapshot = (*snapshotFunc)(TH32CS_SNAPPROCESS, 0 /* ignored */);
  if (snapshot == (HANDLE) -1) {
    // Error occurred during snapshot
    return;
  }

  // Iterate
  PROCESSENTRY32 proc;
  if ((*firstFunc)(snapshot, &proc)) {
    do {
      // FIXME: could make this uniform to the NT version by cutting
      // off the path name just before the executable name
      processes.push_back(ProcEntry(proc.th32ProcessID,
                                    strlen(proc.szExeFile),
                                    proc.szExeFile));
    } while ((*nextFunc)(snapshot, &proc));
  }

  CloseHandle(snapshot);
}