/*
 * pes.c:
 *
 * See the main source file 'vdr.c' for copyright information and
 * how to reach the author.
 *
 */

#include <stdlib.h>

#include "pes.h"
#include <vdr/ringbuffer.h>

//#define DEBUG(x...) printf(x)
#define DEBUG(x...)

//#define PD(x...) printf(x)
#define PD(x...)

// --- cPES --------------------------------------------------------------------

cPES::cPES(eRule ru)
{
  rb=new cRingBufferFrame(KILOBYTE(50));
  defaultRule=ru;
  Reset();
}

cPES::~cPES()
{
  delete rb;
}

void cPES::Reset()
{
  for(int i=0 ; i<NUM_RULESETS ; i++) SetDefaultRule(defaultRule,i);
  UseRuleset(0);
  ClearSeen();
  Clear();
}

void cPES::Clear()
{
  Lock();
  mode=pmNewSync; frame=0; mpegType=2;
  rb->Clear();
  Unlock();
}

bool cPES::ValidRuleset(const int num)
{
  if(num>=0 && num<NUM_RULESETS) return true;
  DEBUG("PES: illegal ruleset %d\n",num);
  return false;
}

void cPES::UseRuleset(const int num)
{
  if(ValidRuleset(num)) {
    currRules=rules[num];
    currNum=num;
    }
}

int cPES::CurrentRuleset()
{
  return currNum;
}

void cPES::SetDefaultRule(eRule ru, const int num)
{
  if(ValidRuleset(num))
    for(int i=0 ; i<NUM_RULES ; i++) rules[num][i]=ru;
}

void cPES::SetRule(int type, eRule ru, const int num)
{
  if(ValidRuleset(num)) rules[num][type]=ru;
}

void cPES::SetRuleR(int ltype, int htype, eRule ru, const int num)
{
  if(ValidRuleset(num)) {
    if(ltype<htype) for( ; ltype<=htype ; ltype++) rules[num][ltype]=ru;
    else DEBUG("PES: bad range %x-%x\n",ltype,htype);
    }
}

unsigned int cPES::Seen(int type) const
{
  return seen[type];
}

void cPES::ClearSeen()
{
  memset(seen,0,sizeof(seen));
  totalBytes=totalSkipped=totalZeros=0;
  skipped=zeros=0;
}

void cPES::Skip(uchar *data, int count)
{
  if(data) {
    Skipped(data,count);
    while(count>0) {
      skipped++;
      if(!*data++) zeros++;
      count--;
      }
    }
  else if(skipped) {
   totalSkipped+=skipped;
   if(skipped==zeros) totalZeros+=zeros;
   else DEBUG("PES: skipped %d bytes\n",skipped);
   skipped=zeros=0;
   }
}

void cPES::Statistics()
{
  if(totalBytes) {
    DEBUG("PES: Stats %lld bytes total, %lld skipped, %lld zero-gaps\n",
          totalBytes,totalSkipped,totalZeros);
    for(int type=0 ; type<=0xFF ; type++)
     if(seen[type])
       DEBUG("PES: Stats %02X: %d packets\n",type,seen[type]);
    }
}

void cPES::ModifyPaketSize(int mod)
{
  if(SOP) {
    int size=header[4]*256+header[5]+mod;
    header[4]=(size>>8)&0xFF;
    header[5]=(size   )&0xFF;
    }
  else DEBUG("PES: modify paket size called in middle of packet\n");
}

void cPES::Redirect(eRule ru)
{
  if(SOP) {
    currRule=ru;
    redirect=true;
    }
  else DEBUG("PES: redirect called in middle of packet\n");
}

int cPES::HeaderSize(uchar *head, int len)
{
  if(len<PES_MIN_SIZE) return -PES_MIN_SIZE;
  switch(head[3]) {
    default:
    // Program end
    case 0xB9:
    // Programm stream map
    case 0xBC:
    // video stream start codes
    case 0x00 ... 0xB8:
    // reserved
    case 0xF0 ... 0xFF:
      return PES_MIN_SIZE;

    // Pack header
    case 0xBA:
      if(len<5) return -5;
      switch(head[4]&0xC0) {
        default:
          DEBUG("PES: unknown mpegType in pack header (0x%02x)\n",head[4]);
          // fall through
        case 0x00:
          mpegType=1;
          return 12;
        case 0x40:
          mpegType=2;
          if(len<14) return -14;
          return 14+(head[13]&0x07); // add stuffing bytes
        }

    // System header
    case 0xBB:
      if(len<6) return -6;
      return 6+head[4]*256+head[5]; //XXX+ is there a difference between mpeg1 & mpeg2??

    // Padding stream
    case 0xBE:
    // Private stream2 (navigation data)
    case 0xBF:
      return 6; //XXX+ is there a difference between mpeg1 & mpeg2??

    // Private stream1
    case 0xBD:
    // all the rest (the real packets)
    case 0xC0 ... 0xCF:
    case 0xD0 ... 0xDF:
    case 0xE0 ... 0xEF:
      if(len<7) return -7;
      int index=6;
      while((head[index]&0xC0)==0xC0) { // skip stuffing bytes
        index++; if(index>=len) return -(index+1);
        }
      if((head[index]&0xC0)==0x80) { // mpeg2
        mpegType=2;
        index+=2; if(index>=len) return -(index+1);
        return index+1+head[index]; // mpeg2 header data bytes
        }
      mpegType=1;
      if((head[index]&0xC0)==0x40) { // mpeg1 buff size
        index+=2; if(index>=len) return -(index+1);
        }
      switch(head[index]&0x30) {
        case 0x30: index+=9; break; // mpeg1 pts&dts
        case 0x20: index+=4; break; // mpeg1 pts
        case 0x10: DEBUG("PES: bad pts/dts flags in MPEG1 header (0x%02x)\n",head[index]); break;
        }
      return index+1;
    }
}

int cPES::PacketSize(uchar *head, int len)
{
  switch(head[3]) {
    default:
    // video stream start codes
    case 0x00 ... 0xB8:
    // Program end
    case 0xB9:
    // Pack header
    case 0xBA:
    // System header
    case 0xBB:
    // Programm stream map
    case 0xBC:
    // reserved
    case 0xF0 ... 0xFF:
      return len; // packet size = header size

    // Private stream1
    case 0xBD:
    // Padding stream
    case 0xBE:
    // Private stream2 (navigation data)
    case 0xBF:
    // all the rest (the real packets)
    case 0xC0 ... 0xCF:
    case 0xD0 ... 0xDF:
    case 0xE0 ... 0xEF:
      return 6+head[4]*256+head[5];
    }
}

int cPES::Return(int used, int len)
{
  PD("PES: return used=%d len=%d mode=%d\n",used,len,mode);
  if(SOP && unsavedHeader && used>=len) {
    // if we are about to finish the current data packet and we have
    // an unsaved header inside, we must save the header to the buffer
    memcpy(hbuff,header,headerSize);
    header=hbuff; unsavedHeader=false;
    PD("PES: header saved\n");
    }
  if(used>len) {
    DEBUG("PES: BUG! used %d > len %d\n",used,len);
    used=len;
    }
  if(used>0) totalBytes+=used;
  Unlock(); // release lock from Process()
  return used;
}

int cPES::Process(const uchar *data, int len, int pid)
{
  Lock(); // lock is released in Return()
  PD("PES: enter data=%p len=%d mode=%d have=%d need=%d old=%d\n",
      data,len,mode,have,need,old);
  int used=0;
  while(used<len) {
    uchar *c=(uchar *)data+used;
    int rest=len-used, n;
    switch(mode) {
      case pmNewSync:
        PD("PES: new sync from %d, rest %d\n",used,rest);
        have=old=need=0;
        unsavedHeader=false; outputHeader=true; SOP=true; redirect=false;
        mode=pmFastSync;
        // fall through

      case pmFastSync:
        // a short cut for the most common case
        // if matched here, header isn't copied around
        if(rest>=PES_MIN_SIZE) {
          PD("PES: fastsync try used=%d: %02x %02x %02x %02x %02x %02x %02x %02x %02x\n",
                used,c[0],c[1],c[2],c[3],c[4],c[5],c[6],c[7],c[8]);
          if(c[2]==0x01 && c[1]==0x00 && c[0]==0x00) {
            headerSize=HeaderSize(c,rest);
            if(headerSize>0 && rest>=headerSize) {
              // found a packet start :-)
              PD("PES: fastsync hit used=%d headerSize=%d rest=%d\n",used,headerSize,rest);
              header=c; unsavedHeader=true;
              used+=headerSize; mode=pmHeaderOk; continue;
              }
            }
          else if(c[2]!=0x00) { used+=3; Skip(c,3); continue; }
          else { used++; Skip(c); continue; }
          }
        // copy remaining bytes to buffer
        memcpy(hbuff,c,rest);
        have=old=rest; used+=rest;
        mode=pmSync;
        PD("PES: buffering started old=%d\n",old);
        break;

      case pmSync:
        PD("PES: slowsync have=%d old=%d\n",have,old);
        if(have<PES_MIN_SIZE && rest>=1) { hbuff[have++]=c[0]; used++; continue; }
        if(have>=PES_MIN_SIZE) {
          PD("PES: slowsync try used=%d: %02x %02x %02x %02x\n",
                used,hbuff[0],hbuff[1],hbuff[2],hbuff[3]);
          if(hbuff[0]==0x00 && hbuff[1]==0x00&& hbuff[2]==0x01) {
            need=abs(HeaderSize(hbuff,have));
            mode=pmGetHeader; continue;
            }
          // no sync found, move buffer one position ahead
          have--; Skip(hbuff);
          memmove(hbuff,hbuff+1,have);
          // if all bytes from previous data block used up, switch to FastSync
          if(!--old) {
            used=0; mode=pmFastSync;
            PD("PES: buffering ended\n");
            }
          continue;
          }
        break;

      case pmGetHeader:
        if(have<need) {
          n=min(need-have,rest);
          memcpy(hbuff+have,c,n);
          have+=n; used+=n;
          PD("PES: get header n=%d need=%d have=%d used=%d\n",n,need,have,used);
          continue;
          }
        if(have>=need) {
          need=abs(HeaderSize(hbuff,have));
          if(have<need) continue;
          // header data complete
          PD("PES: slowsync hit used=%d\n",used);
          if(have>need) DEBUG("PES: bug, buffered too much. have=%d need=%d\n",have,need);
          if(have>(int)sizeof(hbuff)) DEBUG("PES: bug, header buffer overflow. have=%d size=%d\n",have,sizeof(hbuff));
          headerSize=need;
          header=hbuff;
          mode=pmHeaderOk;
          }
        break;

      case pmHeaderOk:
        type=header[3]; seen[type]++; Skip(0);
        if(type<=0xB8) {
          // packet types 0x00-0xb8 are video stream start codes
          DEBUG("PES: invalid packet type 0x%02x, skipping\n",type);
          mode=pmNewSync;
          break;
          }
        payloadSize=PacketSize(header,headerSize)-headerSize;
        if(payloadSize<0) {
          DEBUG("PES: invalid payloadsize %d, skipping\n",payloadSize);
          mode=pmNewSync;
          break;
          }
        PD("PES: found sync at offset %d, type %02x, length %d, next expected %d\n",
              used-headerSize,type,headerSize+payloadSize,used+payloadSize);
        PD("PES: header type=%02x mpeg=%d header=%d payload=%d:",
              type,mpegType,headerSize,payloadSize);
        for(int i=0 ; i<headerSize ; i++) PD(" %02x",header[i]);
        PD("\n");
        currRule=currRules[type];
        have=need=0;
        mode=pmPayload;
        // fall through

      case pmPayload:
        n=min(payloadSize-have,rest);
        if(!n) {
          if(payloadSize==0) n=1;
          else break;
          }
        PD("PES: payload have=%d n=%d SOP=%d\n",have,n,SOP);
        switch(currRule) {
          default:
            DEBUG("PES: bug, unknown rule %d, assuming pass\n",currRule);
            // fall through
          case prPass: n=n; break;
          case prSkip: n=-n; break;
          case prAct1: n=Action1(type,pid,c,n); break;
          case prAct2: n=Action2(type,pid,c,n); break;
          case prAct3: n=Action3(type,pid,c,n); break;
          case prAct4: n=Action4(type,pid,c,n); break;
          case prAct5: n=Action5(type,pid,c,n); break;
          case prAct6: n=Action6(type,pid,c,n); break;
          case prAct7: n=Action7(type,pid,c,n); break;
          case prAct8: n=Action8(type,pid,c,n); break;
          }
        if(n==0) {
          if(redirect) { redirect=false; continue; }
          return Return(used,len);
          }
        need=n; SOP=false;
        mode=pmRingGet;
        // fall through

      case pmRingGet:
        frame=rb->Get();
        if(frame) {
          outCount=frame->Count();
          outData=(uchar *)frame->Data();
          PD("PES: ringbuffer got frame %p count=%d\n",frame,outCount);
          nextMode=pmRingDrop; mode=pmOutput; break;
          }
        mode=pmDataPut;
        // fall through

      case pmDataPut:
        if(need<0) {
          need=-need; outputHeader=false;
          mode=pmDataReady; continue;
          }
        if(outputHeader) {
          outData=header; outCount=headerSize; outputHeader=false;
          nextMode=pmDataPut;
          }
        else if(payloadSize) {
          outData=c; outCount=need;
          nextMode=pmDataReady;
          }
        else {
          mode=pmDataReady;
          continue;
          }
        mode=pmOutput;
        // fall through

      case pmOutput:
        for(;;) {
          PD("PES: output data=%p count=%d -> ",outData,outCount);
          n=Output(outData,outCount);
          PD("n=%d\n",n);
          if(n<0) return Return(-1,len);
          if(n==0) return Return(used,len);

          outCount-=n; outData+=n;
          if(outCount<=0) { mode=nextMode; break; }
          }
        break;

      case pmDataReady:
        if(payloadSize) { used+=need; have+=need; }
        PD("PES: data ready need=%d have=%d paySize=%d used=%d\n",
            need,have,payloadSize,used);
        if(have>=payloadSize) {
          PD("PES: packet finished\n");
          if(have>payloadSize) DEBUG("PES: payload exceeded, size=%d have=%d\n",payloadSize,have);
          mode=pmNewSync;
          }
        else mode=pmPayload;
        break;

      case pmRingDrop:
        PD("PES: ringbuffer drop %p\n",frame);
        rb->Drop(frame); frame=0;
        mode=pmRingGet;
        break;

      default:
        DEBUG("PES: bug, bad mode %d\n",mode);
        return Return(-1,len);
      }
    }
  PD("PES: leave\n");
  return Return(used,len);
}