/*
 *  psinfo.c
 *  
 *  This is a distopian mess of ifdefs.  You have been warned.
 *
 *  Created by Charles Mercadal on Thu Oct 16 2003.
 *
 */


/* include common.h before other headers to make sure we have our defines */
#include "common.h"

#include <stdio.h>

#if defined(USE_MACH) || defined(USE_KVM)	/* both mach and kvm use these */
#include <fcntl.h>
#include <kvm.h>
#endif

#if defined(USE_KVM)
#include <math.h>
#include <nlist.h>
#endif

#if defined(USE_PROCFS)
#define _GNU_SOURCE
#endif
#include <stdlib.h>

#if defined(USE_PROCFS)
#include <string.h>

#include <asm/param.h>
#endif

#if defined(USE_MACH)
#include <mach/mach_error.h>
#include <mach/mach_init.h>
#include <mach/mach_traps.h>
#include <mach/mach_port.h>
#include <mach/task.h>
#include <mach/thread_act.h>
#include <mach/thread_info.h>
#endif

#include <sys/param.h>

#if defined(USE_MACH) || defined(USE_KVM)
#include <sys/fcntl.h>
#include <sys/proc.h>
#include <sys/types.h>
#include <sys/sysctl.h>
#include <sys/user.h>
#endif




#include "psinfo.h"


bool Valid, UseGeneric;
static char *getProcessNameByProcessID(int);
static float getProcessCPUPercentByProcessID(int);

#if defined(USE_MACH) || defined(USE_KVM)
static const char *errorstr = "Couldn't read /dev/mem; you should run pscpug with -g or suid";
static const char *toomany = "too many procs found in kvm_getprocs";
#endif


char *getProcessNameByProcessID(int pid)
{
    char *name;
    
    
    name = NULL;
    
    if (UseGeneric)		/* generic mode specified on the command line */
        goto generic;
    
    #if defined(USE_MACH) || defined(USE_PROCFS)  || defined(USE_KVM)
    #if defined(USE_MACH) || defined(USE_KVM)
    {
        int found;
        kvm_t *kd;
        struct kinfo_proc *givenproc;
    
        
        /* get kvm_t descriptor on /dev/mem */
        if ((kd = kvm_open(NULL, NULL, NULL, O_RDONLY, NULL)) == NULL) {
            perror(errorstr);
            goto generic;
        }
    
        /* 
         * get kinfo_proc structure from kernel for our pid.
         * Since pids are had better be unique, we should only get one.
         */
        givenproc = kvm_getprocs(kd, KERN_PROC_PID, pid, &found);
    
        switch (found) {
            case 0:
                /* no process found for pid */
                Valid = false;
                break;
            case 1:
                /* found; copy process name and return it */
                
                #if defined (__FreeBSD__)
                name = xstrdup(givenproc->ki_comm);
                # else
                name = xstrdup(givenproc->kp_proc.p_comm);
                #endif
                
                Valid = true;
                break;
            default:
                /* we got more than one process?! */
                perror(toomany);
                exit(1);
                /* NOTREACHED */
        }
        
        kvm_close(kd);
        
        return name;
    }
    #elif defined(USE_PROCFS)
    {
        char *proccmdlinefilepath;
        const char *proccmdlineformat = "/proc/%d/cmdline";
        FILE *proccmdline;
        
        (void)asprintf(&proccmdlinefilepath, proccmdlineformat, pid);
        
        /* open /proc/PID/cmdline to read name utility was executed with */
        if ((proccmdline = fopen(proccmdlinefilepath, "r")) == NULL) {
            Valid = false;	/* not found */
        } else {
            ssize_t linelen;
            /*
             * I'm just going to freakin' assume GNU since and use getline
			 * since if I'm using procfs this machine is almost certainly a 
			 * GNU userland/libraries machine.
             */
            getline(&name, &linelen, proccmdline);
            
            fclose(proccmdline);
        }
        
        free(proccmdlinefilepath);
        return name;
    }
    #endif
    #endif /* USE_MACH or USE_PROCFS or USE_KVM */
    
    generic:
    {
        char *cmdbuf;
        #if defined(BSD)
        const char *psnamecmd = "ps -p %d | awk \'/%d/ {print $5}\'";
        #else
        const char *psnamecmd = "ps -p %d | awk \'/%d/ {print $4}\'";
        #endif
        FILE *psoutput;
        
        
        (void)asprintf(&cmdbuf, psnamecmd, pid, pid);
        
        if ((psoutput = popen(cmdbuf, "r")) == NULL) {
            perror(NULL);
            exit(1);
        }
        
        name = (char *)xmalloc(1024 * sizeof(char));
        fgets(name, 1024, psoutput);
        
        if (feof(psoutput))	/* no input */
            Valid = false;
        else
            Valid = true;
        
        pclose(psoutput);
        free(cmdbuf);
        
        return name;
    }

}


float getProcessCPUPercentByProcessID(int pid)
{
    float cpuval;
    
    
    cpuval = 0;
    
    if (UseGeneric)
        goto generic;
    
    #if defined(USE_MACH) || defined(USE_PROCFS)  || defined(USE_KVM)
    #if defined(USE_MACH)
    {
        register int cnt;
        unsigned int numthreads, thread_info_count;
        kern_return_t errors;
        thread_act_array_t allthreads;
        struct thread_basic_info presentThread;
        mach_port_t task;
        
        
        /* get mach task */
        if (task_for_pid(mach_task_self(), pid, &task) != KERN_SUCCESS) {
            goto generic;
        }
        
        /* get all threads from task */
        if ((errors = task_threads(task, &allthreads, &numthreads)) != KERN_SUCCESS) {
            mach_port_deallocate(mach_task_self(), task);
            mach_error("Call to task_threads() failed", errors);
            exit(1);
            /* NOTREACHED */
        }
        
        for (cnt = 0, cpuval = 0; cnt < numthreads; cnt++) {
            thread_info_count = THREAD_BASIC_INFO_COUNT;
            
            if ((errors = thread_info(allthreads[cnt], THREAD_BASIC_INFO, (thread_info_t)&presentThread, &thread_info_count)) != KERN_SUCCESS)
                continue;
                /*
                 * above continue is because sometimes this fails if a task has removed a thread between 
                 * our call to task_threads and our call to thread_info.  We'll just skip this time around...
                 */
            
            cpuval += (float)(((float)presentThread.cpu_usage * 100) / TH_USAGE_SCALE);
        }
    
        mach_port_deallocate(mach_task_self(), task);
        
        Valid = true;
        return cpuval;
    }
    #elif defined(USE_KVM)
    {
        int found, fscale;
        
        #if defined (__FreeBSD__)
        struct kinfo_proc *givenproc;
        #else
        struct kinfo_proc2 *givenproc;
        struct nlist readnl[] = {{"_fscale"}, {NULL}};
        #endif	/* defined(__FreeBSD__) */
        kvm_t *kd;
        
        
        if ((kd = kvm_open(NULL, NULL, NULL, O_RDONLY, NULL)) == NULL) {
            goto generic;
        }
        
        #if defined(__FreeBSD__)
        {
        	size_t oldlen;
        	
        	
        	oldlen = sizeof(fscale);
        	
       		if (sysctlbyname("kern.fscale", &fscale, &oldlen, NULL, 0) == -1) {
       			perror(NULL);
       			exit(1);
       		}
       	}
        #else
        /* get fscale from kernel */
        if (kvm_nlist(kd, readnl)) {
            perror(NULL);
            exit(1);
        }
        
        if ((kvm_read(kd, readnl[0].n_value, &fscale, sizeof(fscale))) != sizeof(fscale)) {
            perror(NULL);
            exit(1);
        }
        #endif	/* defined (__FreeBSD__) */
        
        
        /* get proc */
		#if defined(__FreeBSD__)
		givenproc = kvm_getprocs(kd, KERN_PROC_PID, pid, &found);
		#else
        givenproc = kvm_getproc2(kd, KERN_PROC_PID, pid, sizeof(struct kinfo_proc2), &found);
        #endif	/* defined(__FreeBSD__) */
        
        switch (found) {
            case 0:
                Valid = false;
                break;
            case 1:
                /* these maths gleaned from NetBSD top and ps */
                #define fxtofl(fixpt) ((double)(fixpt) / fscale)
                
                #if defined(__FreeBSD__)
                cpuval = (100 * fxtofl(givenproc->ki_pctcpu));
                #else
                cpuval = (100 * fxtofl(givenproc->p_pctcpu));
                #endif	/* defined(__FreeBSD__) */

                Valid = true;
                break;
            default:
                perror(toomany);
                exit(1);
                /* NOTREACHED */
        }
        
        kvm_close(kd);
        
        return cpuval;
    }
    #elif defined(USE_PROCFS)
    {
        /*
         * You know what?
         * I tried to find a proper value to read from /proc to get %CPU
         * but I seem to have failed miserably.
         * Since the generic hack works fine for this, I'm going to 
         * forgo this part of the code right now and just let 
         * generic do the heavy lifting.
         */
    }
    #endif
    #endif /* USE_MACH or USE_PROCFS or USE_KVM */
    
    generic:
    {
        char *cmdbuf, *cpu;
        const char *cpupctcmd = "ps uww -p %d | awk \'/%d/ {print $3}\'";
        FILE *psoutput;
        
        (void)asprintf(&cmdbuf, cpupctcmd, pid, pid);
        
        if ((psoutput = popen(cmdbuf, "r")) == NULL) {
            perror(NULL);
            exit(1);
        }
        
        cpu = (char *)xmalloc(1024 * sizeof(char));
        fgets(cpu, 1024, psoutput);
        
        if (feof(psoutput))	/* no input */
            Valid = false;
        else
            Valid = true;
        
        pclose(psoutput);
        
        cpuval = atof(cpu);
        
        free(cpu);
        free(cmdbuf);
    
        return cpuval;
    }
}


void setinfo(struct psinfo *set, const int newpsnum, bool gen)
{

    Valid = true;
    UseGeneric = gen;
    set->psname = getProcessNameByProcessID(newpsnum);
    set->psnum = newpsnum;
    set->cpu = 100;
    
    set->statistics.increments = 0;		/* no runs yet */
    set->statistics.cumulative = set->statistics.maxuse = 0;	/* cumulative and max usage are nothing yet */
    set->statistics.minuse = 100;	/* more than the minimum will be */
    
    set->valid = Valid;
}


void updateinfo(struct psinfo *update)
{
    
    Valid = true;
    update->cpu = getProcessCPUPercentByProcessID(update->psnum);
    update->valid = Valid;
    
    if (Valid) {
        update->statistics.increments++;
        update->statistics.cumulative += update->cpu;
        
        if (update->statistics.maxuse < update->cpu) update->statistics.maxuse = update->cpu;
        if (update->statistics.minuse > update->cpu) update->statistics.minuse = update->cpu;
    }
}


void removeinfo(struct psinfo *rem)
{

    free(rem->psname);
}
