// $Id: tsstrip.c,v 1.2 2010/07/07 00:20:47 tosy Exp $

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>
#include <fcntl.h>

#define STMS 10   // number of streams per PMT
#define PMTS 10   // number of PMTs per program

#define NPKT 200000  // maximum buffering packets (~= 20sec @terr.)

typedef struct {
  uint16_t pmt;
  uint16_t sid;
  int      cnt;
  uint8_t  stm[STMS];
  uint16_t pid[STMS];
  uint32_t crc[STMS];
  int      div;
  uint8_t  tmp[188+184];
} pmt_t;

typedef struct {
  uint8_t  pkt[188];
  void     *nxt;
} pktbuf_t;

pmt_t pmt[PMTS];

int curpat;
int pktcnt;
int patcnt;
int chgpat;
int outcnt;

uint16_t sid;
int dbg;

uint32_t
GetBits (uint8_t *buf, int bits)
{
  uint32_t val;

  val = (*buf++) & ((2 << ((bits - 1) & 0x7)) - 1);
  for( ; bits > 8; bits -= 8) {
    val = (val << 8) | (uint32_t)(*buf++);
  }

  return val;
}

uint32_t
GetCrc32 (uint8_t *buf, int len)
{
  uint32_t crc;
  uint32_t x;
  int i, j;

  crc = ~0;
  for (i = 0; i < len; i++) {
    x = (*buf++) << 24;

    for (j = 0; j < 8; j++) {
      crc = (crc << 1) ^ (((crc ^ x) & 0x80000000)? 0x04c11db7: 0);
      x <<= 1;
    }
  }
  return crc;
}

void
AnalyzePat (uint8_t *buf)
{
  int      f, i, p;
  int      siz;
  uint16_t w;

  f = 0;

  //---- continuous ----
  if (patcnt != (buf[3] & 0x0f)) {
    f |= 4;
    patcnt = buf[3] & 0x0f;
  }
  patcnt = (patcnt + 1) & 0x0f;

  //---- CRC ----
  siz = GetBits(&buf[6], 12) - 1;

  if (GetCrc32(&buf[5], siz) != GetBits(&buf[siz+5], 32)) {
    f |= 2;
  }

  //---- analyze ----
  else {
    i = 0;
    for( p = 0x11; p < (siz + 5); p += 4 ) {
      w = GetBits(&buf[p+2], 13);
      if ((w & 0x1ffe) != 0x1fc8) { //ignore one-seg
        if (w != pmt[i].pmt) {
          pmt[i].sid = GetBits(&buf[p], 16);
          pmt[i].pmt = w;
          f |= 1;
        }
        i++;
      }
    }
    for( ; i < PMTS; i++ ) {
      if (pmt[i].pmt != 0)
        f |= 1;
      pmt[i].pmt = 0;
    }
  }

  //---- result ----
  if (f && dbg) {
    fprintf(stderr, "%8d: PAT: ", pktcnt);

    if (f & 2) {
      fprintf(stderr, "[PAT CRC error]");
      if (f & 4)
        fprintf(stderr, "[PAT uncontinuous]");
    }
    else {
      for( i = 0; i < PMTS; i++ ) {
        if (pmt[i].pmt != 0)
          fprintf(stderr, "0x%04x[0x%04x]%s", pmt[i].pmt, pmt[i].sid,
            ((i==(PMTS-1))||(pmt[i+1].pmt==0))? " ": ", ");
      }
      if (f & 4)
        fprintf(stderr, "[PAT uncontinuous]");
      if (f & 1)
        fprintf(stderr, "[PAT changed]");
    }
    fprintf(stderr, "\n");
  }
}

void
AnalyzePmt (uint8_t *buf, pmt_t *pmt)
{
  int f;
  int siz;
  int p;
  int i;

  uint8_t *tmp;
  uint32_t crc;

  f = 0;
  p = 5;

  //---- continuous ----
  if (pmt->cnt != (buf[3] & 0x0f)) {
    f |= 4;
    pmt->cnt = buf[3] & 0x0f;
  }
  pmt->cnt = (pmt->cnt + 1) & 0x0f;

  //---- concatinate ----
  if (pmt->div == 1) {
    memcpy(&(pmt->tmp[188]), &buf[4], 184);
    tmp = pmt->tmp;
    if ((tmp[3] & 0x30) == 0x30)
      p += tmp[4]+1; //adaptation field
    siz = GetBits(&tmp[p+1], 12) - 1;
    pmt->div = 0;
  }
  else {
    tmp = buf;
    if ((tmp[3] & 0x30) == 0x30)
      p += tmp[4]+1; //adaptation field
    siz = GetBits(&tmp[p+1], 12) - 1;

    if ((siz+p+3) > 188) {
      pmt->div = 1;
      memcpy(pmt->tmp, buf, 188);

      if ((siz+p+3) > (188+184)) {
        if (dbg)
          fprintf(stderr, "%8d: PMT 0x%04x: [PMT too large (ignored)]\n", pktcnt, pmt->pmt);
        pmt->div = 0;
      }
      return; //to be continued
    }
  }

  //---- CRC ----
  if (GetCrc32(&tmp[p], siz) != GetBits(&tmp[siz+p], 32)) {
    f |= 2;
  }

  //---- analyze ----
  else {
    i = 0;
    siz += p;
    p += 7;
    for( ; p < siz; p += tmp[p+4] + 5) {
      if (tmp[p] != 0x0d) {
        crc = GetCrc32(&tmp[p], tmp[p+4]+5);
        if (pmt->crc[i] != crc) {
          pmt->stm[i] = tmp[p];
          pmt->pid[i] = GetBits(&tmp[p+1], 13);
          pmt->crc[i] = crc;
          f |= 8;
        }
        i++;
      }
    }
    for( ; i < STMS; i++ ) {
      if (pmt->pid[i] != 0)
        f |= 8;
      pmt->stm[i] = 0;
      pmt->pid[i] = 0;
    }
  }

  //---- result ----
  if (f && dbg) {
    fprintf(stderr, "%8d: PMT 0x%04x: ", pktcnt, pmt->pmt);
    if (f & 2) {
      fprintf(stderr, "[PMT CRC error]");
      if (f & 4)
        fprintf(stderr, "[PMT uncontinuous]");
    }
    else if (f & 8) {
      for( i = 0; i < STMS; i++ ) {
        if (pmt->pid[i] != 0) {
          if (i != 0)
            fprintf(stderr, ", ");
          fprintf(stderr, "0x%04x[0x%02x]", pmt->pid[i], pmt->stm[i]);
        }
      }
      fprintf(stderr, " [PMT changed (%d)]", curpat);
    }
    fprintf(stderr, "\n");
  }

  if (!(f & 2) && (f & 8) && (chgpat >= 0))
    chgpat = curpat;
}

int
AnalyzeTS (int fi, uint8_t *buf)
{
  uint16_t pid;
  int i;

  if (read(fi, buf, 188) <= 0)
    return -1;

  pid = GetBits(&buf[1], 13);

  if (pid == 0x0000) { //PAT
    AnalyzePat(buf);
    curpat = pktcnt;
  }
  else {
    for( i = 0; i < PMTS; i++ ) {
      if ((pmt[i].pmt != 0) && (pmt[i].pmt == pid))
        break;
    }
    if ((i < PMTS) && (pid == pmt[i].pmt)) { //PMT
      AnalyzePmt(buf, &pmt[i]);
    }
  }
  return 0;
}

int
IsFiltered (uint16_t pid)
{
  int i, j;

  if (pid == 0x0000)
    return 1;

  for( i = 0; i < PMTS; i++ ) {
    if (pmt[i].pmt == pid)
      return (pmt[i].sid == sid);

    if ((pmt[i].pmt != 0) && (pmt[i].sid == sid)) {
      for( j = 0; j < STMS; j++ )
        if (pmt[i].pid[j] == pid)
          return 1;
    }
  }
  return 0;
}

void
ProcessTS (int fo, uint8_t *buf)
{
  uint16_t pid;
  uint32_t crc;
  int i, p;

  if (--chgpat >= 0) { //skipping...
    return;
  }

  pid = GetBits(&buf[1], 13);

  if (pid == 0x0000) { //PAT
    p = 0x11;
    for( i = 0; i < PMTS; i++ ) {
      if ((pmt[i].pmt != 0) && (pmt[i].sid == sid)) {
        buf[p++] =  pmt[i].sid >> 8;
        buf[p++] =  pmt[i].sid & 0xff;
        buf[p++] = (pmt[i].pmt >> 8) | 0xe0;
        buf[p++] =  pmt[i].pmt & 0xff;
      }
    }
    buf[7] = p - 4;
    crc = GetCrc32(&buf[5], p - 5);
    buf[p++] =  crc >> 24;
    buf[p++] = (crc >> 16) & 0xff;
    buf[p++] = (crc >>  8) & 0xff;
    buf[p++] =  crc        & 0xff;
    memset(&buf[p], ~0, 188 - p);
  }
  else if (!IsFiltered(pid))
    return;

  outcnt++;

  if (write(fo, buf, 188) == -1) {
    fprintf(stderr, "*** OUTPUT ERROR ***\n");
  }
}

void
ShowUsage (void)
{
  fprintf(stderr, "tsstrip - strip MPEG-2 TS stream\n");
  fprintf(stderr, "usage:  tsstrip [-bd] [-s sid] [input.m2t [output.m2t]]\n");
  fprintf(stderr, "	-b	: enable buffering\n");
  fprintf(stderr, "	-d	: print debug message\n");
  fprintf(stderr, "	-s sid	: select service by sid\n");
}

int
main (int argc, char **argv)
{
  int      fi, fo;
  int      i;
  int      usg;
  int      bfr;
  pktbuf_t *top, *buf, *p;
  char     *ifn, *ofn;
  char     c;

  curpat = 0;
  pktcnt = 0;
  patcnt = 0;
  chgpat = 0;
  outcnt = 0;

  ifn = (char *)NULL;
  ofn = (char *)NULL;
  usg = 0;
  bfr = 0;
  dbg = 0;
  sid = 0;

  while((*++argv != NULL) && (usg == 0)) {
    if ((**argv == '-') && ((*argv)[1] != 0)) {
      ++(*argv);
      while((c = *(*argv)++) != 0) {
        switch(c) {
        case 'b':
          bfr = 1;
          break;
        case 'd':
          dbg = 1;
          break;
        case 's':
          if (**argv != 0)
            sid = strtol(*argv, argv, 0);
          else if ( *++argv != NULL )
            sid = strtol(*argv, argv, 0);
          else
            usg = 1;
          break;
        default:
          usg = 1;
        }
      }
    }
    else {
      if (ifn == NULL)
        ifn = *argv;
      else if (ofn == NULL)
        ofn = *argv;
      else
        usg = 1;
    }
  }

  if (usg == 1) {
    ShowUsage();
    return 1;
  }

  if (dbg) {
    fprintf(stderr, "***  input file: '%s'\n", ifn);
    fprintf(stderr, "*** output file: '%s'\n", ofn);
    if (sid == 0)
      fprintf(stderr, "***  select sid: AUTO\n");
    else
      fprintf(stderr, "***  select sid: %d\n", sid);
    fprintf(stderr, "***     verbose: %s\n", (dbg)? "ON": "OFF");
  }

  if ((ifn == NULL)||(strcmp(ifn, "-") == 0))
    fi = fcntl(STDIN_FILENO, F_DUPFD, O_RDONLY);
  else
    fi = open(ifn, O_RDONLY);
  if ((ofn == NULL)||(strcmp(ofn, "-") == 0))
    fo = fcntl(STDOUT_FILENO, F_DUPFD, O_WRONLY);
  else
    fo = open(ofn, O_WRONLY|O_CREAT|O_TRUNC, 0666);

  if (fi == -1) {
    fprintf(stderr, "*** Cannot open '%s' for input. ***", ifn);
    return 1;
  }
  if (fo == -1) {
    fprintf(stderr, "*** Cannot open '%s' for output. ***", ofn);
    return 1;
  }

  memset(pmt, ~0, sizeof(pmt));

  top = (pktbuf_t *)malloc(sizeof(pktbuf_t));
  buf = top;
  buf->nxt = (pktbuf_t *)NULL;

  //---- Analyze (read only) ----
  while(!AnalyzeTS(fi, buf->pkt)) {

    buf->nxt = (pktbuf_t *)malloc(sizeof(pktbuf_t));
    buf = buf->nxt;
    buf->nxt = (pktbuf_t *)NULL;

    if (++pktcnt >= NPKT)
      break;
  }

  //---- Prepare (select SID) ----
  for( i = 0; i < PMTS; i++ ) {
    if (pmt[i].pmt != 0) {
      if (sid == 0) {
        sid = pmt[i].sid;
        break;
      }
    }
  }

  if (dbg) {
    fprintf(stderr, "%8d: ***START*** SID:0x%04x PID:0x0000", chgpat, sid);
    for( i = 1; i < 0x2000; i++ ) {
      if (IsFiltered(i))
        fprintf(stderr, ",0x%04x", i);
    }
    fprintf(stderr, "\n");
  }

  //---- Process (read/write) ----
  while(!AnalyzeTS(fi, buf->pkt)) {
    pktcnt++;

    if (bfr) {
      if (buf->nxt == (pktbuf_t *)NULL) { //last?
        for( p = top; p != (pktbuf_t *)NULL; p = p->nxt) {
          ProcessTS(fo, p->pkt);
        }
        buf = top; //rewind
      }
    }
    else {
      ProcessTS(fo, top->pkt);

      if (top->nxt != (pktbuf_t *)NULL) {
        p = top;
        top = p->nxt;
        free(p);
        ProcessTS(fo, top->pkt);
      }

      if (top->nxt != (pktbuf_t *)NULL) {
        p = top;
        top = p->nxt;
        buf->nxt = p;
        buf = buf->nxt;
        buf->nxt = (pktbuf_t *)NULL;
      }
    }
  }

  if (dbg) {
    fprintf(stderr, "%8d: ****END****\n", pktcnt);

    i = (int)((uint64_t)outcnt * 10000 / (uint64_t)pktcnt);
    fprintf(stderr, "*** %d packets out (%d.%02d%%) ***\n", outcnt, i/100, i%100);
  }

  close(fi);
  close(fo);

  return 0;
}