/*
 * pstree: display a process heirarchy.
 */
#include "config.h"

#include <stdio.h>
#include <pwd.h>
#include <sys/types.h>
#include <string.h>
#include <unistd.h>
#if HAVE_LIBGEN_H
# include <libgen.h>
#endif

#include "ptree.h"

#if !HAVE_BASENAME
#include <string.h>

char *
basename(char *p)
{
    char *ret = strrchr(p, '/');

    return ret ? ret+1 : p;
}
#endif


extern int  printcard(char *,...);
extern void ejectcard();
extern void cardwidth();

int showargs = 0;	/* -a:  show the entire command line */
int showargpath = 0;	/* -A:  but show full pathname of argv[0] */
int compress = 1;	/* !-c: compact duplicate subtrees */
int clipping = 1;	/* !-l: clip output to screenwidth */
int sortme   = 1;	/* !-n: sort output */
int showpid  = 0;	/* -p:  show process ids */
int showuser = 0;	/* -u:  show username transitions */

Proc * sibsort(Proc *);

/* set the udata field in the Proc*
 */
static void
udata(Proc *t)
{
    char *p = t->process;
    char *q;

    if ( showargs && S(t->cmdline) && T(t->cmdline)[0] )
	p = T(t->cmdline);

    t->udata = (void*) (showargpath || !(q = strrchr(p, '/')) ? p : 1+q);
}


/* compare two ->child trees
 */
cmpchild(Proc *a, Proc *b)
{
    int rc;

    if ( a && b ) {
	if ( a->children != b->children )
	    return a->children - b->children;

	if ( !a->sorted ) {
	    a->child = sibsort(a->child);
	    a->sorted = 1;
	}
	if ( !b->sorted ) {
	    b->child = sibsort(b->child);
	    b->sorted = 1;
	}

	a = a->child;
	b = b->child;

	while ( a && b ) {
	    if ( rc = cmp(a,b) )
		return rc;
	    a = a->sib;
	    b = b->sib;
	}
    }
    if ( a == b) return 0;
    else if ( a ) return 1;
    else return -1;
}


/* do a more macosish sort;  try to not pay attention to case when sorting.
 */
int
cmp(Proc *a, Proc *b)
{
    int rc;

    if ( !a->udata )
	udata(a);

    if ( !b->udata )
	udata(b);

    rc = strcasecmp((char*)a->udata, (char*)b->udata);

    if ( rc == 0 )
	rc = strcmp((char*)a->udata, (char*)b->udata);

    if ( rc == 0 )
	rc = cmpchild(a, b);

    if ( rc == 0 )
	return a->pid - b->pid;

    return rc;
}


/* sort the sibling list by ripping into two substrings, sorting
 * each substring, then stitching them together with a merge sort.
 */
Proc *
sibsort(Proc *p)
{
    Proc *d, *tail, *left = 0, *right = 0;
    int even = 0;

    if ( !(p && p->sib) ) return p;

    /* split into two lists */
    while (p) {
	d = p;
	p = p->sib;

	if (even) {
	    d->sib = left;
	    left = d;
	}
	else {
	    d->sib = right;
	    right = d;
	}
	even = !even;
    }

    /* sort them */
    if ( left ) left = sibsort(left);
    if ( right ) right = sibsort(right);

    /* merge them together */
    for ( p = tail = 0; left && right;  ) {
	if (cmp(left,right) < 0) {
	    d = left;
	    left = left->sib;
	}
	else {
	    d = right;
	    right = right->sib;
	}
	if ( p )
	    tail->sib = d;
	else
	    p = d;
	tail = d;
    }
    tail->sib = left ? left : right;

    return p;
}


/* fancy output printing:  po() and pc() are for
 * comma-delimited ()'ed strings.
 */
static int _paren = 0;

/* when first called, po() prints '(', then every other time
 * it's called it prints ',' until pc() is called to close
 * the parentheses.
 */
int
po()
{
    if (_paren)
	return putcard(',');
    _paren = 1;
    return putcard('(');
}


int
pc()
{
    if (_paren) {
	_paren = 0;
	return putcard(')');
    }
    return 0;
}


static int _bracket = 0;

void
bo()
{
    _bracket++;
}

void
bc()
{
    if (_bracket) --_bracket;
}

void
eol()
{
    int i;

    for (i=0; i < _bracket; i++)
	putcard(']');
    ejectcard();
}


/* to keep track of downward branches, we stuff (column,downward arrow)
 * a tabstack and have dle() properly expand them into spaces, '|', and
 * '`'s
 */
typedef struct tabstack {
    int column;
    int active;
} TabStack;

STRING(TabStack) stack;

int tsp = 0;


/* shove a branch onto the tabstack.
 */
void
push(int column, int c)
{
    if ( tsp >= S(stack) )
	EXPAND(stack);

    T(stack)[tsp].column = column;
    T(stack)[tsp].active = c;
    ++tsp;
}


/* set the downward arrow at tos.
 */
void
active(char c)
{
    if (tsp) T(stack)[tsp-1].active = c;
}


/* pop() [and discard] tos
 */
void
pop()
{
    if (tsp) --tsp;
}


/* return the column position at tos, or 0 if the stack is empty.
 */
int
peek()
{
    return tsp ? T(stack)[tsp-1].column : 0 ;
}


/* print whitespace and downward branches at the start of a line.
 * '`' is a specialcase downward branch;  it's where a branch turns
 * towards the end of the line, so when dle() prints it it resets it
 * to ' '
 */
void
dle()
{
    int dx, xp, dsp;
    char c;

    for ( xp = dsp = 0; dsp < tsp; dsp++ ) {
	dx = T(stack)[dsp].column - xp;
	printcard("%*s%c", dx, "", T(stack)[dsp].active);
	xp = T(stack)[dsp].column;
	if ( T(stack)[dsp].active == '`' )
	    T(stack)[dsp].active = ' ';
    }
}


/* print a string, expanding unusual characters to \ooo escapes
 */
int
prints(unsigned char *s, int len)
{
    int siz, c;

    for (siz=0; len > 0; --len, ++s) {
	if ( !(c = *s) )
	    siz += putcard(' ');
	else if ( isprint(c) && (c > ' ') )
	    siz += putcard(*s);
	else
	    siz += printcard("\\%03o", c);
    }
    return siz;
}


/* print process information (process name, id, uid translation)
 * and branch prefixes and suffixes.   Returns the # of characters
 * printed, so print() can adjust the indent for subtrees
 */
int
printjob(int first, int count, Proc *p)
{
    int tind = 0;

    if ( showargs || !first ) dle();

    if ( count )
	tind = printcard("-%d*[", 1+count);
    else if ( tsp )
	tind = putcard('-');

    if ( !p->udata )
	udata(p);
    tind += prints((char*)p->udata, strlen((char*)p->udata));

    if ( showpid )
	tind += po() + printcard("%d", p->pid);

    if ( showuser && p->parent && (p->uid != p->parent->uid) ) {
	struct passwd *pw = getpwuid(p->uid);

	tind += po();
	if ( pw )
	    tind += printcard("%s", pw->pw_name);
	else
	    tind += printcard("#%d", p->uid);
    }

    tind += pc();

    if ( showargs ) {
	if ( S(p->cmdline) ) {
	    int i;

	    for ( i=0; T(p->cmdline)[i]; ++i)
		;

	    prints(T(p->cmdline) + i, S(p->cmdline) - i);
	}
	eol();
    }
    else if ( p->child ) {
	putcard('-');
	tind++;
    }
    return tind;
}


/* compare two process trees (for subtree compaction)
 * process trees are identical if
 *     (a) the processes are the same
 *     (b) their ->child trees are identical
 *     (c) [if required] all of their siblings are identical
 */
int
sameas(Proc *a, Proc *b, int walk)
{
    if ( ! (a && b) )
	return (a == b);

    if ( strcmp(a->process, b->process) != 0 )
	return 0;

    if ( !sameas(a->child, b->child, 1) )
	return 0;

    if ( !walk ) return 1;

    do {
	if ( !sameas(a->sib, b->sib, 0) )
	    return 0;
	a = a->sib;
	b = b->sib;
    } while ( a && b );
    return 1;
}


/* print() a subtree, indented by a header.
 */
void
print(int first, int count, Proc *node)
{
    int indent;
    char branch;

    indent = printjob(first,count,node);

    if ( node->child ) {
	if ( sortme && !node->sorted ) {
	    node->child = sibsort(node->child);
	    node->sorted = 1;
	}
	node = node->child;
    }
    else {
	if ( !showargs ) eol();
	return;
    }

    count = 0;
    first = 1;

    do {
	if ( compress && sameas(node, node->sib, 0) )
	    count++;
	else {
	    if ( first ) {
		if ( !showargs )
		    putcard(node->sib ? '+' : '-' );
		branch = node->sib ? '|' : (showargs ? '`' : ' ');
		push(peek() + (showargs ? 2 : indent), branch);
	    }
	    if ( branch != ' ' && !node->sib) active('`');
	    if ( count ) bo();
	    print(first,count,node);
	    if ( count ) bc();
	    count=first=0;
	}
    } while ( node = node->sib );
    pop();
}


void
userjobs(Proc *p, uid_t user)
{
    for ( ; p ; p = p->sib )
	if (p->uid == user)
	    print(1,0,p);
	else
	    userjobs(p->child, user);
}


main(int argc, char **argv)
{
    pid_t curid;
    Proc *init;
    Proc *newest;
    int opt;
    char *e;
    extern char version[];

    argv[0] = basename(argv[0]);

    opterr = 1;
    while ( (opt = getopt(argc, argv, "aAclnpuV")) != EOF )
	switch (opt) {
	case 'A':   showargpath = 1; break;
	case 'a':   showargs = 1; compress = 0; break;
	case 'c':   compress = 0; break;
	case 'l':   clipping = 0; break;
	case 'n':   sortme = compress = 0; break;
	case 'p':   showpid  = 1; compress = 0; break;
	case 'u':   showuser = 1; break;
	case 'V':   printf("%s (ps-etc) %s\n", argv[0], version); exit(0);
	default :   exit(1);
	}

    init = ptree(showargs ? PTREE_ARGS : 0);

    if ( !init ) {
	perror(argv[0]);
	exit(1);
    }

    if ( clipping )
	cardwidth();

    argc -= optind;
    argv += optind;

    if ( argc < 1 ) {
	print(1,0,init);
	exit(0);
    }

    curid = strtol(argv[0], &e, 10);

    if ( *e == 0 ) {
	if ( init = pfind(curid) )
	    print(1,0,init);
    }
    else {
	struct passwd *pwd;

	if ( !(pwd = getpwnam(argv[0])) ) {
	    fprintf(stderr, "No such user name: %s\n", argv[0]);
	    exit(1);
	}
	userjobs(init, pwd->pw_uid);
    }

    exit(0);
}
