/*
 * distmult - distribute input and run multiple different commands
 *
 * Usage: distmult [-v] [-s<sep>] [-k<str>] args
 *
 * The "args" are split into one or more sequences delimited by
 *  arguments equaling <sep>, which defaults to "--".  Each sequence is
 *  taken as a command, to be run with execvp, with modifications as
 *  described below.  args begins with the first argument not beginning
 *  with a dash, or the first argument matching <sep>, whichever occurs
 *  first.  Note that flags to distmult are not recognized once args
 *  begins.
 *
 * If the first sequence is empty (that is, if args begins with <seq>),
 *  it is ignored.  Other empty sequences are errors.
 *
 * Input lines are read.  As each line is read, distmult starts a
 *  process corresponding to the next sequence, with that line as
 *  argument.  When all sequences have processes running, distmult
 *  blocks until one dies, then reading the next input line and running
 *  another process.  This continues until EOF is reached on stdin (if
 *  the input ends with a partial line, a warning is printed and a
 *  newline is supplied), at which point distmult then waits for all
 *  child processes to terminate.
 *
 * Before each process is run, if -k is given, any occurrences of <str>
 *  in the sequence are replaced by the input line (with its trailing
 *  newline stripped).  Note that this may involve replacing only part
 *  of an argument.  If -k is not given, a single argument is added
 *  after all explicitly present arguments, holding the input line.
 *  Note that if -k is given but the sequence does not contain <str>,
 *  the corresponding input line is ignored.
 *
 * If at some point there are multiple sequences that could be started
 *  for a given input line, distmult picks one that has never before
 *  run, if there are any such; otherwise, the one that least recently
 *  finished.
 *
 * The commands run inherit distmult's own stdout and stderr, and get
 *  /dev/null on their stdin - except that before this is done,
 *  /dev/null is placed on any of the standard descriptors which were
 *  closed at startup.
 *
 * By default distmult prints nothing except in error cases.  With -v,
 *  it prints each sequence as it starts a process for it; the sequence
 *  is printed after the processing for -k, or the appending that takes
 *  place by default if -k is not given.
 *
 * distmult's own exit status is a count of errorful child
 *  terminations, or 1 if some other error caused a premature exit -
 *  but never more than 127 in any case.
 *
 * Copyright status: this program is explicitly in the public domain.
 *  Anyone may use it in any way for any purpose, though I would
 *  appreciate credit where it is due.
 *	--Mouse, mouse@rodents.montreal.qc.ca, 2005-12-20
 */

#include <stdio.h>
#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <stdlib.h>
#include <signal.h>
#include <strings.h>
#include <sys/wait.h>
#include <searchstr.h>

extern const char *__progname;

typedef struct line LINE;

struct line {
  LINE *link;
  char *txt;
  } ;

static const char *devnull = "/dev/null";

static const unsigned char equiv[256]
 = { 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,
     25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,
     47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,
     69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,
     91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,
     109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,
     125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,
     141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,
     157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,
     173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,
     189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,
     205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,
     221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,
     237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,
     253,254,255 };

static int vflag = 0;
static const char *sep = "--";
static const char *kflag = 0;
static unsigned char (*ktbl)[256];
static int klen;
static int dnfd;
static int nseqs;
static char ***seqs;
static int *seqlen;
static pid_t *kids;
static int livekids;
static int *order;
static int nextorder;
static int ateof;
static int nerr;
static LINE *lines;
static LINE **linetail;
static char *part;
static int partall;
static int partlen;

static void usage(void) __attribute__((__noreturn__));
static void usage(void)
{
 fprintf(stderr,"Usage: %s [-v] [-s<sep>] [-k<str>] args\n",__progname);
 exit(1);
}

static void findseqs(int ac, char **av)
{
 int i;
 int l;

 nseqs = 0;
 seqs = 0;
 l = 0;
 for (i=0;i<=ac;i++)
  { if ((i == ac) || !strcmp(av[i],sep))
     { av[i] = 0;
       seqs = realloc(seqs,(nseqs+1)*sizeof(*seqs));
       seqs[nseqs++] = av + i - l;
       l = 0;
     }
    else
     { l ++;
     }
  }
}

static void print_seq(char **seq, FILE *f)
{
 int i;

 for (i=0;seq[i];i++) fprintf(f,"%s%s",i?" ":"",seq[i]);
}

static void kidwait(int flags)
{
 pid_t dead;
 int status;
 int i;

 while (1)
  { dead = wait3(&status,flags,0);
    if (dead <= 0) return;
    for (i=0;i<nseqs;i++)
     { if (dead == kids[i])
	{ if (WIFEXITED(status))
	   { if (WEXITSTATUS(status) != 0)
	      { fprintf(stderr,"%s: ",__progname);
		print_seq(seqs[i],stderr);
		fprintf(stderr,": exit %d\n",WEXITSTATUS(status));
		nerr ++;
	      }
	   }
	  else if (WIFSIGNALED(status))
	   { fprintf(stderr,"%s: ",__progname);
	     print_seq(seqs[i],stderr);
	     fprintf(stderr,": %s%s\n",strsignal(WTERMSIG(status)),WCOREDUMP(status)?" (core dumped)":"");
	     nerr ++;
	   }
	  else if (WIFSTOPPED(status))
	   { fprintf(stderr,"%s: ",__progname);
	     print_seq(seqs[i],stderr);
	     fprintf(stderr,": stopped (%s)\n",strsignal(WSTOPSIG(status)));
	     fprintf(stderr,"%s: (this should not happen - sending SIGKILL)\n",__progname);
	     kill(dead,SIGKILL);
	     nerr ++;
	   }
	  else
	   { fprintf(stderr,"%s: ",__progname);
	     print_seq(seqs[i],stderr);
	     fprintf(stderr,": unclassifiable exit status %d\n",status);
	     nerr ++;
	   }
	  kids[i] = -1;
	  order[i] = nextorder ++;
	  livekids --;
	}
     }
    flags |= WNOHANG;
  }
}

static void catch_sig(int sig __attribute__((__unused__)))
{
}

static void setup_signals(void)
{
 struct sigaction sa;

 sa.sa_handler = catch_sig;
 sigemptyset(&sa.sa_mask);
 sa.sa_flags = 0;
 sigaction(SIGCHLD,&sa,0);
}

static void new_line(const char *p1, int l1, const char *p2, int l2)
{
 LINE *l;

 l = malloc(sizeof(LINE)+l1+l2);
 l->txt = (void *)(l+1);
 if (l1) bcopy(p1,l->txt,l1);
 if (l2) bcopy(p2,l->txt+l1,l2);
 l->link = 0;
 *linetail = l;
 linetail = &l->link;
}

static void geteof(void)
{
 ateof = 1;
 if (partlen > 0)
  { fprintf(stderr,"%s: missing newline supplied\n",__progname);
    new_line(part,partlen,0,0);
  }
}

static void append_part(const char *p, int l)
{
 if (partlen+l > partall) part = realloc(part,partall=partlen+l);
 bcopy(p,part+partlen,l);
 partlen += l;
}

static void readsome(void)
{
 char rbuf[8192];
 int rlen;
 char *nl;
 int x;

 rlen = read(0,&rbuf[0],sizeof(rbuf));
 if (rlen < 0)
  { if (errno != EINTR)
     { fprintf(stderr,"%s: input read error: %s\n",__progname,strerror(errno));
       geteof();
     }
    return;
  }
 if (rlen == 0)
  { geteof();
    return;
  }
 x = 0;
 while (x < rlen)
  { nl = memchr(&rbuf[x],'\n',rlen-x);
    if (nl == 0) break;
    new_line(part,partlen,&rbuf[x],nl-&rbuf[x]);
    partlen = 0;
    x = (nl+1) - &rbuf[0];
  }
 if (x < rlen) append_part(&rbuf[x],rlen-x);
}

static char *k_sub(const char *s, const char *rep)
{
 int sl;
 int rl;
 int o;
 int f;
 int l;
 char *t;
 char *tp;

 sl = strlen(s);
 rl = strlen(rep);
 l = 0;
 o = 0;
 while (1)
  { f = searchstr(s+o,sl-o,ktbl,klen);
    if (f < 0) break;
    l += f + rl;
    o += f + klen;
  }
 l += sl - o;
 t = malloc(l+1);
 tp = t;
 o = 0;
 while (1)
  { f = searchstr(s+o,sl-o,ktbl,klen);
    if (f < 0) break;
    if (f > 0) bcopy(s+o,tp,f);
    tp += f;
    if (rl > 0) bcopy(rep,tp,rl);
    tp += rl;
    o += f + klen;
  }
 if (sl > o)
  { bcopy(s+o,tp,sl-o);
    tp += sl - o;
  }
 if (tp != t + l) abort();
 *tp = '\0';
 return(t);
}

static void run_kid(int x)
{
 pid_t kid;
 char **av;
 int i;
 LINE *l;
 int xp[2];
 int r;

 if (kids[x] >= 0) abort();
 l = lines;
 lines = l->link;
 if (! lines) linetail = &lines;
 if (kflag)
  { av = malloc((seqlen[x]+1)*sizeof(char *));
    for (i=0;i<seqlen[x];i++) av[i] = k_sub(seqs[x][i],l->txt);
  }
 else
  { av = malloc((seqlen[x]+2)*sizeof(char *));
    for (i=0;i<seqlen[x];i++) av[i] = seqs[x][i];
    av[i++] = l->txt;
  }
 av[i] = 0;
 if (pipe(xp) < 0)
  { fprintf(stderr,"%s: pipe: %s\n",__progname,strerror(errno));
    exit(1);
  }
 kid = fork();
 if (kid < 0)
  { fprintf(stderr,"%s: fork: %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (kid == 0)
  { dup2(dnfd,0);
    close(dnfd);
    close(xp[0]);
    fcntl(xp[1],F_SETFD,1);
    execvp(av[0],av);
    i = errno;
    write(xp[1],&i,sizeof(int));
    exit(0);
  }
 close(xp[1]);
 r = 0;
 do <"exec">
  { while (r < sizeof(int))
     { int n;
       n = read(xp[0],((char *)&i)+r,sizeof(int)-r);
       if (n < 0)
	{ if (errno == EINTR) continue;
	  fprintf(stderr,"%s: exec pipe read error: %s\n",__progname,strerror(errno));
	  exit(1);
	}
       if (n == 0)
	{ if (r == 0)
	   { kids[x] = kid;
	     livekids ++;
	     if (vflag)
	      { print_seq(av,stderr);
		fprintf(stderr,"\n");
	      }
	     break <"exec">;
	   }
	  fprintf(stderr,"%s: exec pipe protocol failure\n",__progname);
	  exit(1);
	}
       r += n;
     }
    fprintf(stderr,"%s: exec ",__progname);
    print_seq(av,stderr);
    fprintf(stderr,": %s\n",strerror(i));
    nerr ++;
  } while (0);
 if (kflag) for (i=0;i<seqlen[x];i++) free(av[i]);
 free(av);
 close(xp[0]);
}

int main(int, char **);
int main(int ac, char **av)
{
 int i;
 int j;

 ac --;
 av ++;
 nerr = 0;
 while (ac > 0)
  { if (!strcmp(av[0],sep)) break;
    if (!strcmp(av[0],"-v"))
     { vflag = 1;
       ac --;
       av ++;
       continue;
     }
    if (!strncmp(av[0],"-s",2))
     { sep = av[0] + 2;
       ac --;
       av ++;
       continue;
     }
    if (!strncmp(av[0],"-k",2))
     { kflag = av[0] + 2;
       klen = strlen(kflag);
       if (klen == 0)
	{ fprintf(stderr,"%s: zero-length -k argument\n",__progname);
	  exit(1);
	}
       if (klen > 255)
	{ fprintf(stderr,"%s: -k argument too long\n",__progname);
	  exit(1);
	}
       ktbl = malloc(klen*sizeof(*ktbl));
       searchstr_maketbl(ktbl,kflag,klen,&equiv[0]);
       ac --;
       av ++;
       continue;
     }
    if (av[0][0] == '-')
     { fprintf(stderr,"%s: unrecognized flag: %s\n",__progname,av[0]);
       nerr ++;
     }
    break;
  }
 if (nerr) usage();
 if (ac < 1) usage();
 findseqs(ac,av);
 if (nseqs == 0) usage();
 if (seqs[0][0] == 0)
  { seqs ++;
    nseqs --;
    if (nseqs == 0) usage();
  }
 kids = malloc(nseqs*sizeof(*kids));
 order = malloc(nseqs*sizeof(*order));
 seqlen = malloc(nseqs*sizeof(*seqlen));
 for (i=0;i<nseqs;i++)
  { kids[i] = -1;
    order[i] = 0;
    for (j=0;seqs[i][j];j++) ;
    if (j == 0)
     { fprintf(stderr,"%s: empty sequence\n",__progname);
       exit(1);
     }
    seqlen[i] = j;
  }
 dnfd = open(devnull,O_RDWR,0);
 if (dnfd < 0)
  { fprintf(stderr,"%s: can't open %s: %s\n",__progname,devnull,strerror(errno));
    exit(1);
  }
 while (dnfd < 3) dnfd = dup(dnfd);
 nextorder = 1;
 livekids = 0;
 part = 0;
 partall = 0;
 partlen = 0;
 lines = 0;
 linetail = &lines;
 setup_signals();
 while (1)
  { if (lines)
     { j = -1;
       for (i=0;i<nseqs;i++) if (kids[i] < 0) if ((j < 0) || (order[i] < order[j])) j = i;
       if (j < 0)
	{ kidwait(0);
	  continue;
	}
       run_kid(j);
       continue;
     }
    if (! ateof)
     { readsome();
       kidwait(WNOHANG);
       continue;
     }
    if (livekids < 1) break;
    kidwait(0);
  }
 exit((nerr>127)?127:nerr);
}