// gcc -g -Ilibs/libavutil -Ilibs/libavformat/ -Ilibs/libavcodec -o mpegfix mpegfix-0.0.1.c lib/libmythavformat-0.18.so

#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <netinet/in.h>
#include <assert.h>

#include "avformat.h"
#include "avcodec.h"

#define LLSIZE 80
#define WRITE
int debug = 0;

typedef struct FRAME_PTR {
  char *data;
  int size;
  int frame_num;
  struct FRAME_PTR *next;
} FRAME_PTR;

typedef struct {
  int fh;
  long ptrSize;
  int64_t ptsIncrement;
  int adjustFrameCount;
  int lastFrameNum;
  char lastFrameType;
  char *lastFramePos;
  int isSequence;
  int isGop;
} AVPktData;

typedef struct AVPackList {
  AVPacket pkt;
  int size;
  struct AVPackList *next;
  int mem_size;
} AVPackList;
//#define MPEG2trans_DEBUG
#define MATCH_HEADER(ptr) (((ptr)[0] == 0x00) && ((ptr)[1] == 0x00) && ((ptr)[2] == 0x01))

AVFormatContext *inputFC = NULL;
AVCodecContext  *pCodecCtx;
AVCodec         *pCodec;

void DPRINTF(int level, const char *format, ...)
{
  if (level <= debug) {
    va_list ap;
    va_start(ap, format);
    vprintf(format, ap);
    va_end(ap);
  }
}

unsigned long GETBITS (unsigned char * ptr, int num)
{
  unsigned long value = 0, gb_long;
  static unsigned char *gb_ptr = 0;
  static unsigned int gb_pos;
  int offset, offset_r, offset_b;
  if (ptr != 0) {
     gb_ptr = ptr;
     gb_pos = 0;
  }
  offset = gb_pos >> 3;
  offset_r = gb_pos & 0x07;
  offset_b = 32 - offset_r;
  gb_long = ntohl(*((unsigned long *) (gb_ptr + offset)));
  value = gb_long >> (offset_b - num);
  value = value & (0xffffffff >> (32 - num));
  gb_pos+=num;
  return value;
}

void SETBITS(unsigned char *ptr, long value, int num)
{
  static int sb_pos;
  static unsigned char *sb_ptr = 0;
  unsigned long sb_long, mask;
  int offset, offset_r, offset_b;
   if(ptr != 0) {
     sb_ptr = ptr;
     sb_pos = 0;
   }
   offset = sb_pos >> 3;
   offset_r = sb_pos & 0x07;
   offset_b = 32 - offset_r;
   mask = ~(((1 << num) -1)<<(offset_b - num));
   sb_long = ntohl(*((unsigned long *) (sb_ptr + offset)));
   value = value << (offset_b - num);
   sb_long = (sb_long & mask) + value;
   *((unsigned long *)(sb_ptr + offset)) = htonl(sb_long);
}

int InitAV(char *inputfile)
{
    int ret;

    // Open recording
    DPRINTF(0, "Opening %s\n", inputfile);
    ret = av_open_input_file(&inputFC, inputfile, NULL, 0, NULL);
    if (ret != 0)
    {
        fprintf(stderr,"Couldn't open input file, error #%d\n", ret);
        return 0;
    }

    // Getting stream information
    ret = av_find_stream_info(inputFC);
    if (ret < 0)
    {
        fprintf(stderr, "Couldn't get stream info, error #%1", ret);
        av_close_input_file(inputFC);
        inputFC = NULL;
        return 0;
    }

    // Dump stream information
    dump_format(inputFC, 0, inputfile, 0);

    pCodecCtx=inputFC->streams[0]->codec;
    pCodec=avcodec_find_decoder(pCodecCtx->codec_id);
    if(pCodec==NULL)
      return 0;
    pCodecCtx->flags |=CODEC_FLAG_LOW_DELAY; //force all frames to be i/p
    if(avcodec_open(pCodecCtx, pCodec)<0)
      return 0;
    
    return 1;
}

unsigned int setFrameCount(char *ptr, int num) {
  SETBITS(ptr+4, num, 10);
}

int getFrameType(AVPacket *pkt, char **offset)
{
  uint8_t *ptr = pkt->data;
  int len = pkt->size;
  while (len >= 4) {
    if(MATCH_HEADER(ptr) && ptr[3] == 0x00) {
      GETBITS(ptr+4, 10);
      if(offset)
        *offset = ptr;
      return(GETBITS(0, 3));
    }
    ptr++;
    len=pkt->size - (ptr - pkt->data);
    len++;
  }
  return 0;
}

FRAME_PTR *buildPframe(AVPackList *lVpkt_first, AVPackList *lVpkt_last,
                 FRAME_PTR *frame_head, int insert_count,
                 int lastFrameNum)
{
  char *cpy=NULL, *offset;
  int cpy_size=0, used;
  int cur_cnt=0;
  int finished=0;
  char *outbuf;
  int outbuf_size=1000000;
  int out_size, out_size1, out_type;
  AVCodecContext *c=NULL;
  AVCodec *out_codec;
  AVFrame *pFrame;
  AVPacket pkt;

  outbuf=(char *)malloc(outbuf_size);
  pFrame = avcodec_alloc_frame(); 
  while(lVpkt_first != lVpkt_last->next) {
    int frame_type = getFrameType(&lVpkt_first->pkt, &offset);
    int tmp_size = lVpkt_first->pkt.size;
    // - (offset - (char *)lVpkt_first->pkt.data);
    if (frame_type == 1 || frame_type == 2) {
      if(cpy_size < tmp_size) {
        cpy = (char *)realloc(cpy, tmp_size);
        cpy_size=tmp_size;
      }
      memcpy(cpy, lVpkt_first->pkt.data, tmp_size);
      setFrameCount(cpy + (offset - (char *)lVpkt_first->pkt.data), cur_cnt++);
      used=avcodec_decode_video(pCodecCtx, pFrame, &finished, cpy, cpy_size);
      if(used != cpy_size || !finished) {
        fprintf(stderr, "Failed to decode frame!\n");
        assert(0);
        return;
      }
    }
    lVpkt_first=lVpkt_first->next;
  }
  if(!finished) {
    fprintf(stderr, "Didn't find frame to decode!\n");
    assert(0);
    return;
  }
  
  out_codec = avcodec_find_encoder(pCodecCtx->codec_id);
  if(! out_codec) {
    fprintf(stderr, "Couldn't find MPEC2 encoder\n");
    assert(0);
    return;
  }
  c=avcodec_alloc_context();
  c->bit_rate = pCodecCtx->bit_rate;
  c->width = pCodecCtx->width;
  c->height = pCodecCtx->height;
  c->time_base= pCodecCtx->time_base;
  c->pix_fmt = pCodecCtx->pix_fmt;
  c->max_b_frames=0;
  c->gop_size=100;
  pFrame->pts = AV_NOPTS_VALUE;
  if (avcodec_open(c, out_codec) < 0) {
        fprintf(stderr, "could not open codec\n");
        exit(1);
  }
  out_size=avcodec_encode_video(c, outbuf, outbuf_size, pFrame);
  out_size1=avcodec_encode_video(c, outbuf, outbuf_size, pFrame);
  pkt.data=outbuf;
  pkt.size=out_size1;
  out_type=getFrameType(&pkt, &offset);
  {
     char frame_type;
     switch (out_type) {
       case 1: frame_type = 'I'; break;
       case 2: frame_type = 'P'; break;
       case 3: frame_type = 'B'; break;
       default: frame_type = 'X'; break;
     }
     DPRINTF(2,"Preparing %c-frame size: old-size: %d new-size: %d\n", out_type, cpy_size, out_size1);
  }
  while(insert_count) {
    frame_head->data=(char *)malloc(out_size1);
    setFrameCount(offset, ++lastFrameNum);
    memcpy(frame_head->data, outbuf, out_size1);
    frame_head->size=out_size1;
    frame_head->frame_num = lastFrameNum;
    frame_head=frame_head->next;
    insert_count--;
  }
  free(outbuf);
  avcodec_close(c);
  free(c);
  free(pFrame);
  if(cpy_size)
    free(cpy);
  return frame_head;
}

unsigned int process_video(AVPacket *pkt, AVPktData *pkt_data)
{
  uint8_t *ptr = pkt->data;
  int len = pkt->size;
  pkt_data->isSequence = 0;
  pkt_data->isGop = 0;
  while (len >= 4) {
    if(MATCH_HEADER(ptr) &&
       (ptr[3] == 0xb3 | ptr[3] == 0xb8 | ptr[3] == 0x00)) {
      switch (ptr[3]) {
      case 0xb3 : // sequence header
      {
        pkt_data->isSequence = 1;
        GETBITS(ptr+7, 4);
        switch(GETBITS(0, 4)) {
          case 1:
            pkt_data->ptsIncrement = 3754;
            break;
          case 2:
            pkt_data->ptsIncrement = 3750; //24
            break;
          case 3:
            pkt_data->ptsIncrement = 3600; // 25
            break;
          case 4:
            pkt_data->ptsIncrement = 3003; //29.997
            break;
          case 5:
            pkt_data->ptsIncrement = 3000; //30
        }
        ptr+=11;
        GETBITS(ptr, 6);
        if(GETBITS(0, 1)) {
          ptr+=64;
        }
        GETBITS(ptr, 7);
        if(GETBITS(0, 1)) {
          ptr+=64;
        }
        DPRINTF(4,"Rate: %d\n", pkt_data->ptsIncrement);
      } break;
      case 0xb8 : // gop header
      {
//        printf("GOP\n");
        ptr+=7;
        pkt_data->adjustFrameCount=0;
        pkt_data->isGop=1;
      } break;
      case 0x00: // picture header
      {
        int num, type;
        char pictype;
        num = GETBITS(ptr+4, 10);
        type = GETBITS(0, 3);
        switch (type) {
          case 1: pictype = 'I'; break;
          case 2: pictype = 'P'; break;
          case 3: pictype = 'B'; break;
          case 4: pictype = 'D'; break;
          default: pictype = 'X'; break;
        }
        if(num == 1023) {
          pkt_data->adjustFrameCount=0;
        }
          
        if(pkt_data->adjustFrameCount) {
          num+=pkt_data->adjustFrameCount;
          setFrameCount(ptr, num);
          DPRINTF(4,"Adjust: %d == %d\n", num,  GETBITS(ptr+4, 10));
          num = GETBITS(ptr+4, 10);
        }
        pkt_data->lastFrameType = pictype;
        pkt_data->lastFrameNum = num;
        pkt_data->lastFramePos = ptr;
        return;
      } break;
      default: {
      } break;
    }
  }
  ptr++;
  len=pkt->size - (ptr - pkt->data);
  }
  return 0;
}
#define VID_STREAM 0
#define AUD_STREAM 1

int main(int argc, char **argv) 
{
  char *infile;
  AVPackList *lApkt_tail, *lApkt_head,
               *lVpkt_tail, *lVpkt_head, *lVpkt_tail_orig;
  AVPackList lApkt[LLSIZE], lVpkt[LLSIZE];
  AVPacket pkt, *Vpkt_last;
  AVPktData pkt_data;
  int ret, i;
  int state=0;
  int64_t deltaPTS, initPTS, expectedPTS = 0, tmpVidPts = 0, maxVidPts = 0;
  int aud_fh;
  FRAME_PTR *frame_tail, *frame_head, frame[10];

  if (argc != 2) {
      fprintf(stderr, "Usage: %s <filename>\n", argv[0]);
      exit(1);
  }
  infile = argv[1];

  frame_tail=frame_head=frame;
  for(i=0;i<9;i++) {
    frame[i].next=&frame[i+1];
  }
  frame[9].next = frame_tail;

  memset(&pkt_data, 0, sizeof(pkt_data));
  aud_fh = open64("temp.mp2", O_WRONLY|O_CREAT, S_IRWXU);
  pkt_data.fh = open64("temp.m2v", O_WRONLY|O_CREAT, S_IRWXU);

  memset(lApkt, 0, sizeof(AVPackList)*LLSIZE);
  memset(lVpkt, 0, sizeof(AVPackList)*LLSIZE);
  for(i=0; i<LLSIZE-1; i++) {
    lApkt[i].next = &lApkt[i+1];
    lVpkt[i].next = &lVpkt[i+1];
  }
  lApkt_tail = lApkt_head = lApkt[LLSIZE-1].next = &lApkt[0];
  lVpkt_tail = lVpkt_head = lVpkt[LLSIZE-1].next = &lVpkt[0];
  av_register_all();
  if(! InitAV(infile)) {
    exit(-1);
  }

  av_init_packet(&pkt);

  while(1)
  {
     /* read packet */
     AVPackList *curPkt;
     ret = av_read_frame(inputFC, &pkt);
     if (ret < 0) break;
     {
       uint8_t *ptr;
       if(pkt.stream_index==VID_STREAM) {
         curPkt = lVpkt_head;
         lVpkt_head = lVpkt_head->next;
       } else {
         curPkt = lApkt_head;
         lApkt_head=lApkt_head->next;
       }
       ptr=curPkt->pkt.data;
       curPkt->pkt = pkt;
       if(curPkt->mem_size < pkt.size) {
         curPkt->pkt.data = (uint8_t *)realloc(ptr, pkt.size);
         curPkt->mem_size = pkt.size;
       } else {
         curPkt->pkt.data = ptr;
       }
       memcpy(curPkt->pkt.data, pkt.data, pkt.size);
       av_free_packet(&pkt);
     }
     if(state == 0) {
        while(lVpkt_tail != lVpkt_head) {
          process_video(&lVpkt_tail->pkt, &pkt_data);
          if(pkt_data.isSequence)
            break;
          DPRINTF(3,"Dropping V packet\n");
          lVpkt_tail = lVpkt_tail->next;
        }
        if(lVpkt_tail != lVpkt_head) {
          while(lApkt_tail != lApkt_head) {
            if(lApkt_tail->pkt.pts < lVpkt_tail->pkt.pts &&
               lApkt_tail->next != lApkt_head) {
              if(lApkt_tail->next->pkt.pts > lVpkt_tail->pkt.pts) {
                state=1;
                break;
              } else {
                DPRINTF(3,"Dropping A packet\n");
                lApkt_tail = lApkt_tail->next;
                continue;
              }
            } else if(lApkt_tail->pkt.pts == lVpkt_tail->pkt.pts) {
              state=1;
              break;
            } else if(lApkt_tail->pkt.pts > lVpkt_tail->pkt.pts) {
              state=1;
              break;
            }
            if(lApkt_tail->next == lApkt_head) {
              break;
            }
          }
        }
        if(state==1) {
          deltaPTS=lVpkt_tail->pkt.pts - lApkt_tail->pkt.pts;
          if(deltaPTS < 0) {
            initPTS=lVpkt_tail->pkt.pts - 16200;
          } else {
            initPTS=lApkt_tail->pkt.pts - 16200;
          }
          DPRINTF(1,"%lld %lld %lld %d\n", lVpkt_tail->pkt.pts, initPTS, pkt_data.ptsIncrement,  pkt_data.lastFrameNum);
          expectedPTS = lVpkt_tail->pkt.pts - initPTS - 
                        (pkt_data.ptsIncrement * pkt_data.lastFrameNum);
          DPRINTF(1,"V: %lld A: %lld delta: %lld\n", lVpkt_tail->pkt.pts - initPTS,
                                                  lApkt_tail->pkt.pts - initPTS,
                                                  deltaPTS);
          DPRINTF(0,"Delta: %fms\n", 1000.0*deltaPTS/90000.0);
        }
      }
      if(state==1) {
        if(curPkt->pkt.stream_index==VID_STREAM) {
          char *offset;
          DPRINTF(4,"Frame type: %d pos: %lx\n", getFrameType(&curPkt->pkt, &offset), curPkt->pkt.data);
        if(getFrameType(&curPkt->pkt, NULL)==1) {
          lVpkt_tail_orig = lVpkt_tail;
          while(lVpkt_tail->next != lVpkt_head) {
            lVpkt_tail->pkt.pts -= initPTS;
            process_video(&lVpkt_tail->pkt, &pkt_data);
            if(pkt_data.isGop) {
               expectedPTS += maxVidPts;
               maxVidPts = 0;
            }
            while(frame_tail != frame_head &&
                  frame_tail->frame_num <= pkt_data.lastFrameNum &&
                  (pkt_data.lastFrameType == 'I' ||
                   pkt_data.lastFrameType == 'P')) {
              //time to insert a new p-frame
              DPRINTF(2,"Sending new P-Frame(1): %d\n", frame_tail->frame_num);
#ifdef WRITE
              write(pkt_data.fh, frame_tail->data,
                    frame_tail->size);
#endif
              free(frame_tail->data);
              frame_tail=frame_tail->next;
              setFrameCount(pkt_data.lastFramePos, ++pkt_data.lastFrameNum);
              pkt_data.adjustFrameCount++;
            }
            tmpVidPts =  pkt_data.ptsIncrement * pkt_data.lastFrameNum;
            if(tmpVidPts + pkt_data.ptsIncrement > maxVidPts)
              maxVidPts = tmpVidPts + pkt_data.ptsIncrement;
            DPRINTF(1,"VID: %c-Frame #:%d pts: %lld exp: %lld pos: %lx\n",
                   pkt_data.lastFrameType, pkt_data.lastFrameNum, 
                   lVpkt_tail->pkt.pts, expectedPTS + tmpVidPts, pkt_data.lastFramePos);

            if(tmpVidPts + pkt_data.ptsIncrement == maxVidPts &&
               lVpkt_tail->pkt.pts > expectedPTS + tmpVidPts + 
                                     (pkt_data.ptsIncrement >> 1)) {
              if(lVpkt_tail->pkt.pts < expectedPTS + tmpVidPts + 10*pkt_data.ptsIncrement) {
              int insert_count;
              insert_count = (lVpkt_tail->pkt.pts + 
                              (pkt_data.ptsIncrement >> 1) -
                              (expectedPTS + tmpVidPts)) /
                              pkt_data.ptsIncrement;
              frame_head=buildPframe(lVpkt_tail_orig, lVpkt_tail, frame_head,
                                     insert_count, pkt_data.lastFrameNum);
            }
            }
#ifdef WRITE
            write(pkt_data.fh, lVpkt_tail->pkt.data, lVpkt_tail->pkt.size);
#endif
            lVpkt_tail=lVpkt_tail->next;
          }
          while(frame_tail != frame_head) {
            //insert a new p-frame before the next I-frame
            DPRINTF(2,"Sending new P-Frame(2): %d\n", frame_tail->frame_num);
#ifdef WRITE
            write(pkt_data.fh, frame_tail->data,
                  frame_tail->size);
#endif
            free(frame_tail->data);
            frame_tail=frame_tail->next;
            maxVidPts += pkt_data.ptsIncrement;
          }
          // can't free packets until we are done incase we need to build
          // a new p-frame along the line
          while(lVpkt_tail_orig != lVpkt_tail) {
            lVpkt_tail_orig=lVpkt_tail_orig->next;
          }
        }
        }
        while(lApkt_tail != lApkt_head) {
          lApkt_tail->pkt.pts -= initPTS;
          // write_audio(lApkt_tail->pkt, initPTS);
          DPRINTF(1,"AUD: pts: %lld\n", lApkt_tail->pkt.pts);
#ifdef WRITE
          write(aud_fh, lApkt_tail->pkt.data, lApkt_tail->pkt.size);
#endif
          lApkt_tail=lApkt_tail->next;
        }
     }
  }
  av_close_input_file(inputFC);
  close(aud_fh);
  close(pkt_data.fh);
  return 0;
}