/*
 *  a dvd copy tool
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License version 2 as
 *  published by the Free Software Foundation;
 *
 *
 *  Some parts of this source code were taken from the following projects:
 *  transcode    by Dr. Thomas Oestreich
 *  lsdvd        by Chris Phillips
 *  mplex        by mjpeg-tools developers group
 *
 */


#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <getopt.h>
#include <dvdread/ifo_read.h>
#include <dvdread/ifo_types.h>
#include <dvdread/dvd_reader.h>

#define VERSION        "0.4"
#define MAXSTREAMS     50
#define PROBEBLOCKS    30720

#define MPEG_VIDEO  0
#define MPEG_AUDIO  1
#define AC3_AUDIO   2
#define DTS_AUDIO   3
#define LPCM_AUDIO  4
#define SUBPICTURE  5


int debug                          = 0;
double dvd_size_bytes              = 4700000000.;
char *stream_type_names[6]         = { "MPEG Video", "MPEG Audio", "AC3 Audio", "DTS Audio", "LPCM-Audio", "Subpicture" };
double frames_per_s[4]             = {-1.0, 25.00, -1.0, 29.97};

// bitrate tables stolen from mplex
unsigned int ac3_bitrate_index[32] = { 32,40,48,56,64,80,96,112,128,160,192,224,256,320,384,448,512,576,640,0,0,0,0,0,0,0,0,0,0,0,0,0 };
unsigned int dts_bitrate_index[32] = { 32,56,64,96,112,128,192,224,256,320,384,448,512,576,640,768,960,1024,1152,1280,1344,1408,1411,1472,1536,1920,2048,3072,3840,0,0,0 };
unsigned int mpa_bitrate_index [2][4][16] =
{
        { /* MPEG audio V2 */
                {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0},
                {0,32,64,96,128,160,192,224,256,288,320,352,384,416,448,0},
                {0,32,48,56,64 ,80 ,96 ,112,128,160,192,224,256,320,384,0},
                {0,8 ,16,24,32 ,64 ,80 ,56 ,64 ,128,160,112,128,256,320,0}
        },
        { /* MPEG audio V1 */
                {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0},
                {0,32,64,96,128,160,192,224,256,288,320,352,384,416,448,0},
                {0,32,48,56,64 ,80 ,96 ,112,128,160,192,224,256,320,384,0},
                {0,32,40,48,56 ,64 ,80 ,96 ,112,128,160,192,224,256,320,0}
        }

};


struct streamdata {            // data for 1 stream
   unsigned char id;              // stream id
   long bitrate;                  // bitrate
   int type;                      // type of stream
   double size;                   // size of stream in bytes
};



/*
 * calc playback length in ms - stolen from lsdvd
 */

int dvdtime2msec(dvd_time_t *dt)
{
  double fps = frames_per_s[(dt->frame_u & 0xc0) >> 6];
  long   ms;
  ms  = (((dt->hour &   0xf0) >> 3) * 5 + (dt->hour   & 0x0f)) * 3600000;
  ms += (((dt->minute & 0xf0) >> 3) * 5 + (dt->minute & 0x0f)) * 60000;
  ms += (((dt->second & 0xf0) >> 3) * 5 + (dt->second & 0x0f)) * 1000;

  if(fps > 0)
  ms += ((dt->frame_u & 0x30) >> 3) * 5 + (dt->frame_u & 0x0f) * 1000.0 / fps;

  return ms;
}



/*
 *  Extract the real start and stop blocks to read from the ifo structure
 *  If start and stop chapters are 0, they'll be set to the right values (start=1, stop=lastchapter)
 */

int get_start_stop_blocks (dvd_reader_t *dvd, int title, uint8_t *ts_nr, int *chapstart, int *chapstop, uint32_t *startblock, uint32_t *lastblock, int* runtime ) {
   int ok = 0;
   ifo_handle_t   *ifo_title;
   ifo_handle_t   *ifo_main;
   vts_ptt_srpt_t *vts_ptt_srpt;
   tt_srpt_t      *tt_srpt;
   pgc_t          *pgc, *full_pgc;
   int titleid =  title -1, ttn, pgc_id, pgn, start_cell, end_cell, i;

   if( title == 0 ) {                                                     // don't rip title 0
      fprintf(stderr, "Invalid title number %d\n", title);
      return 0; 
   }

   ifo_main = ifoOpen(dvd, 0);                                            // open main ifo file
   if( !ifo_main ) {
      fprintf(stderr, "Error opening main ifo file\n");
      return(0); 
   }

   if( titleid > ifo_main->tt_srpt->nr_of_srpts ) {                       // title nr. < max num. of titles ?
      fprintf(stderr, "Invalid title number %d\n", title);
      return(0); 
   }

   tt_srpt = ifo_main->tt_srpt;
   *ts_nr = tt_srpt->title[ titleid ].title_set_nr;                       // save title set nr

   ifo_title = ifoOpen(dvd, tt_srpt->title[ titleid ].title_set_nr);      // open title ifo file
   if( !ifo_title ) {
      fprintf(stderr, "Error opening ifo file for title %d\n", title);
      ifoClose(ifo_main);
      return(0);
   }

   if( ifo_title->vtsi_mat ) {                                            // get start/stop block for selected chapters
      if( *chapstart == 0 ) *chapstart = 1;                                    // set first chapter if unset
      if( *chapstop == 0 ) *chapstop = tt_srpt->title[ titleid ].nr_of_ptts;   // set last chapter if unset

      ttn          = tt_srpt->title[ titleid ].vts_ttn;                   // mostly stolen from transcode
      vts_ptt_srpt = ifo_title->vts_ptt_srpt;

      pgc_id       = vts_ptt_srpt->title[ttn - 1].ptt[*chapstart-1].pgcn;
      pgn          = vts_ptt_srpt->title[ttn - 1].ptt[*chapstart-1].pgn;
      pgc          = ifo_title->vts_pgcit->pgci_srp[pgc_id - 1].pgc;
      full_pgc     = pgc;

      start_cell   = pgc->program_map[pgn - 1] - 1;
      *startblock  = pgc->cell_playback[start_cell].first_sector;         // first block of start cell

      if( *chapstop == tt_srpt->title[ titleid ].nr_of_ptts ) {           // if stop chapter = last chapter
         end_cell     = pgc->nr_of_cells -1;
         *lastblock   = pgc->cell_playback[end_cell].last_sector;
      }
      else {
         pgc_id       = vts_ptt_srpt->title[ttn - 1].ptt[*chapstop].pgcn;
         pgn          = vts_ptt_srpt->title[ttn - 1].ptt[*chapstop].pgn;
         pgc          = ifo_title->vts_pgcit->pgci_srp[pgc_id - 1].pgc;
         end_cell     = pgc->program_map[pgn - 1] - 2;
         *lastblock   = pgc->cell_playback[end_cell].last_sector;         // last block of end cell
      }

      if( (*chapstart == 1) && (*chapstop == tt_srpt->title[ titleid ].nr_of_ptts) ) {   // all chapters selected ?
         *runtime = dvdtime2msec(&full_pgc->playback_time) / 1000; }                        // runtime = movie runtime
      else {                                                                             // only some chapters selected ?
         *runtime = 0;
         for( i=start_cell ; i<=end_cell ; i++ ) {                                          // sum runtime of selected cells
            *runtime += dvdtime2msec(&pgc->cell_playback[i].playback_time) / 1000; }
      }

      ok = 1;
   }

   ifoClose(ifo_title);
   ifoClose(ifo_main);

   return(ok);
}




int probe_dvd ( dvd_file_t *vob, struct streamdata **list, uint32_t start_block, uint32_t stop_block ) {
   int x, y, z, found, count = 0;
   unsigned char buffer[DVD_VIDEO_LB_LEN], *searchptr;
   uint8_t packet_type;
   struct streamdata *new_data;
   uint32_t maxblocks = stop_block, i;

   if( (stop_block-start_block) > PROBEBLOCKS ) {
      maxblocks = start_block + PROBEBLOCKS; }

   for( i=start_block ; i<maxblocks ; i++ ) {                                  // read 10mb from dvd
      if( !DVDReadBlocks(vob, i, 1, buffer) ) {
         fprintf(stderr, "ERROR probing for streams\n");
         exit(1);
      }

      found = 0;
      packet_type = buffer[17];                                                // get stream id
      if( packet_type == 0xBD ) {                                              // if private stream
         packet_type = buffer[23+buffer[22]];                                     // get sub-id
         //if(debug) fprintf(stderr, "found stream 0x%x\n", packet_type);
      }
                                                                               // check for video, audio, subpicture
      if( ((packet_type >= 0xE0) && (packet_type <= 0xEF)) ||
          ((packet_type >= 0x20) && (packet_type <= 0x3F)) ||
          ((packet_type >= 0xC0) && (packet_type <= 0xCF)) ||
          ((packet_type >= 0xA0) && (packet_type <= 0xAF)) ||
          ((packet_type >= 0x80) && (packet_type <= 0x8F)) ) {
         for( x=0 ; x<count ; x++ ) {                                          // scan if we already know this stream
            if( list[x]->id == packet_type ) {                                 // makes sure, we only get headers at stream start
               found = 1;
               break;
            }
         }
         if( !found ) {                                                        // unknown stream, add to list
            new_data = malloc(sizeof(struct streamdata));
            if( new_data == NULL ) {
               fprintf(stderr, "Out of memory\n");
               exit(1);
            }

            new_data->id = packet_type;

            if( (packet_type >= 0xE0) && (packet_type <= 0xEF) ) {
               new_data->type = MPEG_VIDEO;
               new_data->bitrate = 0;                                                      // dummy entry for video
            }
            else if( (packet_type >= 0xC0) && (packet_type <= 0xCF) ) {
               new_data->type = MPEG_AUDIO;
               new_data->bitrate = 0;
               searchptr = buffer +17;
               searchptr = (unsigned char*)memchr(searchptr, 0xFF, searchptr-buffer);      // search for 1st byte of MPA Syncword 
               while( searchptr != NULL ) {
                  if( searchptr <= (buffer+2045) ) {                                       // check for search hit at end of buffer
                     if( (searchptr[1] & 0xF0) == 0xF0 ) {                                 // if next 4 bits set, we found the sync word
                        x = (searchptr[1] & 0x08) >> 3;                                    // version id
                        y = (searchptr[1] & 0x06) >> 1;                                    // layer code
                        z = (searchptr[2] & 0xF0) >> 4;                                    // bitrate code
                        new_data->bitrate = mpa_bitrate_index[x][y][z] * 1024;
                        if( debug) fprintf(stderr, "MPEG_AUDIO bitrate: %d\n", mpa_bitrate_index[x][y][z]);
                        break;
                     }
                  }
                  searchptr++;
                  searchptr = (unsigned char*)memchr(searchptr, 0xFF, searchptr-buffer);
               }
               if( new_data->bitrate == 0 ) {                                              // no syncword found, try next packet
                  fprintf(stderr, "unable to get MPEG-Audio bitrate\n");
                  continue;
               }
            }
            else if( (packet_type >= 0x80) && (packet_type <= 0x87) ) {
               y = buffer[22] + buffer[buffer[22]+26] + (unsigned char)26;                 // calc start offsets
               if( (buffer[y] != 0x0B) || (buffer[y+1] != 0x77) ) {                        // check for AC3 Syncword
                  continue; }
               new_data->type = AC3_AUDIO;
               x = buffer[y + 4] & 0x3F;                                                   // 6 lowest bits
               new_data->bitrate = ac3_bitrate_index[ x>>1 ] * 1024;                       // shift index >> 1 (same as /2 )
               if( debug) fprintf(stderr, "AC3_AUDIO bitrate: %d\n", ac3_bitrate_index[ x>>1 ]);
            }
            else if( (packet_type >= 0x88) && (packet_type <= 0x8F) ) {
               new_data->type = DTS_AUDIO;
               x = (buffer[buffer[22] + 35] & 0x3) << 3;                                   // 3 lowest bits of byte
               x = x | ((buffer[buffer[22] + 36] & 0xE0) >> 5);                            // 2 highest bits of next byte
               new_data->bitrate = dts_bitrate_index[x] * 1000;                            // kbit * 1000 fits better to filesize, so what ...
               if( debug) fprintf(stderr, "DTS_AUDIO bitrate: %d\n", dts_bitrate_index[x]);
            }
            else if( (packet_type >= 0xA0) && (packet_type <= 0xAF) ) {
               new_data->type = LPCM_AUDIO;
               new_data->bitrate = 2000000;                                                // FIXME - hard coded, no detection yet
            }
            else if( (packet_type >= 0x20) && (packet_type <= 0x3F) ) {
               new_data->type = SUBPICTURE;
               new_data->bitrate = 0;                                                      // dummy entry for subpicture
            }

            list[count] = new_data;
            count++;
         }
      }
   }

   return(count);
}




void usage (char *progname, int exitcode) {
   fprintf(stderr, "usage:  %s -i device -t title [-c chapter|chapter range] [-s stream(s)] [-v] [-h]\n", progname);
   fprintf(stderr, "           device : dvd device or directory containing dvd structure\n");
   fprintf(stderr, "           title  : title number\n");
   fprintf(stderr, "           chapter: chapter number or range        (default: all chapter)\n");
   fprintf(stderr, "           streams: video/audio/subtitle stream(s) (default: all streams)\n");
   fprintf(stderr, "\n");
   exit(exitcode);
}


void version ( void ) {
   fprintf(stderr, "   streamanalize v%s - (C) 2003 by Reinhardt Wolf\n", VERSION);
   exit(0);
}



int main(int argc, char *argv[]) {

   dvd_reader_t *dvd;
   dvd_file_t *vob;
   char *opt_inputname = NULL;
   int opt_title = 0;
   int opt_chapter_start = 0;
   int opt_chapter_stop = 0;
   int opt_streamlist[MAXSTREAMS];
   int opt_stream_count = 0;

   struct streamdata *streamdata_list[MAXSTREAMS];
   int stream_count = 0;
   uint8_t ts_nr;
   int i, c, x;
   int *b;
   int runtime;
   double video_size, overhead_size, nonvideo_size = 0, selected_size = 0, temp;
   uint32_t start_block = 0, stop_block = 0;



   while( (c = getopt(argc, argv, "i:t:c:s:d:hv")) != -1 ) {            // process options
      switch(c) {
         case 'h': usage(argv[0], 0);
                   break;

         case 'v': version();
                   break;

         case 'i': opt_inputname = optarg;
                   break;

         case 't': opt_title = atoi(optarg);
                   break;

         case 's': b = opt_streamlist;
                   if( (x = sscanf(optarg, "%x,%x,%x,%x,%x,%x,%x,%x,%x,%x", &b[0],&b[1],&b[2],&b[3],&b[4],&b[5],&b[6],&b[7],&b[8],&b[9])) > 0 ) {
                      opt_stream_count = opt_stream_count + x; 
                   }
                   break;

         case 'c': x = sscanf(optarg,"%d-%d", &opt_chapter_start, &opt_chapter_stop);
                   if( x != 2 ) {
                      x = sscanf(optarg,"%d", &opt_chapter_start);
                      if( x != 1 ) {
                         fprintf(stderr, "invalid parameter for option -c\n");
                         exit(1);
                      }
                      opt_chapter_stop = opt_chapter_start;
                   }
                   break;
       
         default: usage(argv[0], 1);
      }
   }
   if( (opt_inputname == NULL) || (opt_title == 0) ) usage(argv[0], 1);

   dvd = DVDOpen(opt_inputname);
   if( !dvd ) {
      fprintf(stderr, "error opening dvd at %s\n", opt_inputname); 
      exit(1);
   }



   if( get_start_stop_blocks(dvd, opt_title, &ts_nr, &opt_chapter_start, &opt_chapter_stop, &start_block, &stop_block, &runtime) ) {
      vob = DVDOpenFile(dvd, ts_nr, DVD_READ_TITLE_VOBS);
      if( !vob ) {
         fprintf(stderr, "ERROR: opening vobs for title %d failed\n", ts_nr);
         DVDClose(dvd);
         exit(1);
      }

      stream_count = probe_dvd(vob, streamdata_list, start_block, stop_block); // probe some blocks to find all streams
      DVDCloseFile(vob);

      if( !opt_stream_count ) {                                                // if no streamlist given as parameter
         opt_stream_count = stream_count;                                         // write found streams to streamlist
         for( i=0 ; i<stream_count ; i++ ) {
            opt_streamlist[i] = streamdata_list[i]->id; }
      }

      fprintf(stderr, "\n\n");
      fprintf(stderr, "Title %d - ", opt_title);
      fprintf(stderr, "%d Chapters (%d Blocks / %11.0f Bytes) - ", opt_chapter_stop-opt_chapter_start+1, stop_block-start_block, (double)(stop_block-start_block) * 2048);
      fprintf(stderr, "Runtime %d sec.\n\n", runtime);
      fprintf(stderr, "Track List:\n");
      for( i=0 ; i<stream_count ; i++ ) {
         switch( streamdata_list[i]->type ) {
            case MPEG_VIDEO : streamdata_list[i]->size = 0;
                              break;
            case SUBPICTURE : streamdata_list[i]->size = (double)(stop_block-start_block) * 2048 / 1450;
                              break;
            case LPCM_AUDIO :
            case MPEG_AUDIO :
            case AC3_AUDIO  :
            case DTS_AUDIO  : streamdata_list[i]->size = streamdata_list[i]->bitrate / 8 * runtime;
                              break;
         }
         nonvideo_size += streamdata_list[i]->size;
      }

      overhead_size = (stop_block-start_block) * 2048 / 50;                    // rough guess about space wasted by padding, pack headers, ...
      nonvideo_size += overhead_size;                                          // sum of all non video stuff
      video_size = ((double)(stop_block-start_block) * 2048) - nonvideo_size;  // video size = vob size - non video stuff

      for( i=0 ; i<stream_count ; i++ ) {
         if( streamdata_list[i]->type == MPEG_VIDEO ) streamdata_list[i]->size = video_size;  // add video size

         for( x=0 ; x<opt_stream_count ; x++ ) {
            if( opt_streamlist[x] == streamdata_list[i]->id ) {
               fprintf(stderr, "X"); 
               selected_size += streamdata_list[i]->size;
            }
         }

         fprintf(stderr, "   0x%x  %10s", streamdata_list[i]->id, stream_type_names[streamdata_list[i]->type]);
         fprintf(stderr, "  %10.0f Bytes\n", streamdata_list[i]->size);
      }

      fprintf(stderr, "\n");
      fprintf(stderr, "Size of selected streams: %11.0f Bytes\n", selected_size+overhead_size);
      fprintf(stderr, "Max. target size        : %11.0f Bytes\n", dvd_size_bytes);
      if( (selected_size+overhead_size) <= dvd_size_bytes ) {
         fprintf(stderr, "Factor                  :  not needed\n"); }
      else {
         temp = (selected_size+overhead_size) - dvd_size_bytes;
         temp = video_size - temp;
         fprintf(stderr, "Factor                  :  %5.3f\n", video_size / temp);
      }
   }



   DVDClose(dvd); 


   return 0;
}