/*
   gop_fixup:  Fixes timestamp of MPEG2 GOP headers.

   (c) 2001-2004 Mark Rages <markrages@mlug.missouri.edu>
   
   Released under the terms of the GNU GPL.  You know what they are. 
*/

#include "sysdep.h"
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#include <assert.h>
#include <error.h>

int fixed_gop=0;
int fixed_df=0;

unsigned int framerate=0;

void gop2time(unsigned char* gop, unsigned char* time) {

  time[0]=(gop[0]&(64+32+16+8+4))>>2;
  time[1]=((gop[0]&(2+1))<<4)
    +((gop[1]&(128+64+32+16))>>4);
  time[2]=((gop[1]&(4+2+1))<<3)
    +((gop[2]&(128+64+32))>>5);
  time[3]=((gop[2]&(16+8+4+2+1))<<1)
    +((gop[3]&(128))>>7);
}

void printtime(unsigned char* time) {
  printf("%2.2d:%2.2d:%2.2d.%2.2d                 \r",time[0],time[1],time[2],time[3]);
}

void time2gop(unsigned char* gop, unsigned char* time) {

  gop[0] = (gop[0]&128)+             // 7
    (time[0]<<2)+                    // 6-2
    ((time[1]&(32+16))>>4);          // 1-0
  gop[1] = ((time[1]&(8+4+2+1))<<4)+ // 7-4
    (gop[1]&8)+                      // 3
    ((time[2]&(32+16+8))>>3);        // 2-0
  gop[2] = ((time[2]&(4+2+1))<<5)+
    ((time[3]&(32+16+8+4+2))>>1);
  gop[3] = ((time[3]&1)<<7)+
    (gop[3]&(64+32+16+8+4+2+1));
}

void frame2time(int frame, unsigned char* time) {
  time[3] = frame%framerate;
  time[2] = (frame/framerate)%60;
  time[1] = ((frame/framerate)/60)%60;
  time[0] = ((frame/framerate)/60)/60;
}

void goptime(unsigned char* mem, int i){

  unsigned char oldtime[4]; // h,m,s,f
  unsigned char newtime[4]; // h,m,s,f

  gop2time(mem,oldtime);
  frame2time(i,newtime);

  if (mem[0]&128) {
    printf(" Setting non-DF\n");
    mem[0]&=~128;
    assert(!(mem[0]&128));
    fixed_df++;
  }

  if (memcmp(oldtime,newtime,4) != 0) {
    fixed_gop++;
    printf(" Fixing ");
    time2gop(mem,newtime);
  }
  printtime(newtime);
  fflush(stdout);
}

/* Let's list some startcodes... */
#define PICTURE_START_CODE   0x00
#define USER_DATA_START_CODE 0xb2
#define SEQUENCE_HEADER_CODE 0xb3
#define SEQUENCE_ERROR_CODE  0xb4
#define EXTENSION_START_CODE 0xb5

#define SEQUENCE_END_CODE    0xb7
#define GROUP_START_CODE     0xb8

#define MIN(a,b) ((a)<(b)?(a):(b))

#define MAP_PAGES 1000

int main(int argc, char ** argv) {
  int fd;
  int frame_count;
  int sequence_count=0;
  int fi;
  int fr_index;
  unsigned char *mem;
  unsigned int frameratetab[]={0,24,24,25,30,30,30,50,60,60};
  unsigned int new_framerate;

  OFF_T len, fpos, map_len, map_start, map_size;

  map_size=getpagesize()*MAP_PAGES;

  for (fi=1; fi<argc; fi++) {
    fd = open_rw(argv[fi]);
    if (fd==-1) {
      perror("opening file"); exit(3);
    }
    
    len = lseek(fd,0,SEEK_END);
    if (errno) perror("Finding Length");
    lseek(fd,0,SEEK_SET);
    
    map_start=0;
    frame_count = 0;
    
    while (len>0) {

      map_len=MIN(len,map_size);
      /*
      printf("\nMap size: %d\n",map_size);
      printf("len: %f\n",(float)len);
      printf("map_len: %d\n",map_len);
      printf("map_start: %d\n",map_start);
      */
      
      /* We an extra page, so when we get to the end of the map, 
         we're still looking at good data */

      mem = mmap(0,map_size+getpagesize(),PROT_READ|PROT_WRITE,MAP_SHARED,fd,map_start);
      
      if ((int)mem==-1) { perror("\nmmap()ing"); exit(1); }      

      for (fpos=0; fpos<map_len; fpos++) {

	if ((mem[fpos]==0) &&
	    (mem[fpos+1]==0) &&
	    (mem[fpos+2]==1))
	  {
	    switch(mem[fpos+3]) 
	      {
	      case SEQUENCE_END_CODE:
		break;
	      case SEQUENCE_HEADER_CODE:
		fr_index=mem[fpos+4+3]&0xf;
		assert(fr_index<10);
		new_framerate=frameratetab[fr_index];
		if (framerate) {
		  if (framerate!=new_framerate) {
		    printf("Segments in stream have mismatched framerates.\n");
		    printf("%d != %d\n",framerate,new_framerate);
		    printf("Whatever, continuing...\n");
		    //exit(1);
		  }
		} else {
		  framerate=new_framerate;
		  printf("Found sequence with nominal framerate %d\n",framerate);
		}
		break;
	      case GROUP_START_CODE:
		if (!framerate) { 
		  printf("No framerate information found before GOP.\n"); 
		  exit(1);
		}
		printf("%s ",argv[fi]);
		goptime(mem+fpos+4,frame_count);
		break;
	      case PICTURE_START_CODE:
		frame_count++;
		break;
	      }
	  }
      }
      munmap(mem,map_size+getpagesize());
      
      map_start+=(map_len);
      len-=(map_len);
    }

    printf("\nfinished %s \n",argv[fi]);
    printf("fixed %d GOP timestamps\n",fixed_gop);
    printf("set %d timestamps to non-drop-frame\n",fixed_df);
    printf("encountered %d pictures\n",frame_count);
    fixed_gop=fixed_df=0;
    close(fd);
  }
  return 0;
}