/*  Zutils - Utilities dealing with compressed files
    Copyright (C) 2009, 2010, 2011, 2012, 2013 Antonio Diaz Diaz.

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#define _FILE_OFFSET_BITS 64

#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#include <stdint.h>
#include <unistd.h>
#include <sys/wait.h>

#include "zutils.h"
#include "rc.h"


const char * invocation_name = 0;
const char * util_name = program_name;

int verbosity = 0;


int parse_format_type( const std::string & arg )
  {
  for( int i = 0; i < num_formats; ++i )
    if( arg == format_names[i] )
      return i;
  show_error( "Bad argument for '--format' option." );
  std::exit( 1 );
  }


// Returns the number of bytes really read.
// If (returned value < size) and (errno == 0), means EOF was reached.
//
int readblock( const int fd, uint8_t * const buf, const int size )
  {
  int rest = size;
  errno = 0;
  while( rest > 0 )
    {
    const int n = read( fd, buf + size - rest, rest );
    if( n > 0 ) rest -= n;
    else if( n == 0 ) break;				// EOF
    else if( errno != EINTR && errno != EAGAIN ) break;
    errno = 0;
    }
  return size - rest;
  }


// Returns the number of bytes really written.
// If (returned value < size), it is always an error.
//
int writeblock( const int fd, const uint8_t * const buf, const int size )
  {
  int rest = size;
  errno = 0;
  while( rest > 0 )
    {
    const int n = write( fd, buf + size - rest, rest );
    if( n > 0 ) rest -= n;
    else if( n < 0 && errno != EINTR && errno != EAGAIN ) break;
    errno = 0;
    }
  return size - rest;
  }


bool feed_data( const int infd, const int outfd,
                const uint8_t * magic_data, const int magic_size )
  {
  if( magic_size && writeblock( outfd, magic_data, magic_size ) != magic_size )
    { show_error( "Write error", errno ); return false; }
  enum { buffer_size = 4096 };
  uint8_t buffer[buffer_size];
  while( true )
    {
    const int size = readblock( infd, buffer, buffer_size );
    if( size != buffer_size && errno )
      { show_error( "Read error", errno ); return false; }
    if( size > 0 && writeblock( outfd, buffer, size ) != size )
      { show_error( "Write error", errno ); return false; }
    if( size < buffer_size ) break;
    }
  return true;
  }


bool good_status( const Children & children, const bool finished )
  {
  bool error = false;
  for( int i = 0; i < 2; ++i )
    {
    const pid_t pid = children.pid[i];
    if( pid )
      {
      const char * const msg =
        ( i & 1 ) ? children.compressor_name : "data feeder";
      if( !finished )
        {
        const int tmp = child_status( pid, msg );
        if( tmp < 0 ) kill( pid, SIGTERM );	// child not terminated
        else if( tmp != 0 ) error = true;	// child status != 0
        }
      else
        if( wait_for_child( pid, msg ) != 0 ) error = true;
      }
    }
  return !error;
  }


bool set_data_feeder( int * const infdp, Children & children, int format_index )
  {
  const uint8_t * magic_data = 0;
  int magic_size = 0;
  if( format_index < 0 )
    format_index = test_format( *infdp, &magic_data, &magic_size );
  children.compressor_name = get_compressor_name( format_index );

  if( children.compressor_name )	// compressed
    {
    int fda[2];				// pipe from feeder
    int fda2[2];			// pipe from compressor
    if( pipe( fda ) < 0 || pipe( fda2 ) < 0 )
      { show_error( "Can't create pipe", errno ); return false; }
    const int old_infd = *infdp;
    *infdp = fda2[0];
    const pid_t pid = fork();
    if( pid == 0 )			// child 1 (compressor feeder)
      {
      if( close( fda[0] ) != 0 ||
          close( fda2[0] ) != 0 || close( fda2[1] ) != 0 ||
          !feed_data( old_infd, fda[1], magic_data, magic_size ) )
        _exit( 2 );
      if( close( fda[1] ) != 0 )
        { show_close_error( "data feeder" ); _exit( 2 ); }
      _exit( 0 );
      }
    if( pid < 0 )			// parent
      { show_fork_error( "data feeder" ); return false; }

    const pid_t pid2 = fork();
    if( pid2 == 0 )			// child 2 (compressor)
      {
      if( dup2( fda[0], STDIN_FILENO ) >= 0 &&
          dup2( fda2[1], STDOUT_FILENO ) >= 0 &&
          close( fda[0] ) == 0 && close( fda[1] ) == 0 &&
          close( fda2[0] ) == 0 && close( fda2[1] ) == 0 )
        {
        const std::vector< std::string > & compressor_args =
          get_compressor_args( format_index );
        const int size = compressor_args.size();
        const char ** const argv = new const char *[size+3];
        argv[0] = children.compressor_name;
        for( int i = 0; i < size; ++i )
          argv[i+1] = compressor_args[i].c_str();
        argv[size+1] = ( verbosity >= 0 ) ? "-d" : "-dq";
        argv[size+2] = 0;
        execvp( argv[0], (char **)argv );
        }
      show_exec_error( children.compressor_name );
      _exit( 2 );
      }
    if( pid2 < 0 )			// parent
      { show_fork_error( children.compressor_name ); return false; }

    close( fda[0] ); close( fda[1] ); close( fda2[1] );
    children.pid[0] = pid;
    children.pid[1] = pid2;
    }
  else					// uncompressed
    {
    int fda[2];				// pipe from feeder
    if( pipe( fda ) < 0 )
      { show_error( "Can't create pipe", errno ); return false; }
    const int old_infd = *infdp;
    *infdp = fda[0];
    const pid_t pid = fork();
    if( pid == 0 )			// child (feeder)
      {
      if( close( fda[0] ) != 0 ||
          !feed_data( old_infd, fda[1], magic_data, magic_size ) )
        _exit( 2 );
      if( close( fda[1] ) != 0 )
        { show_close_error( "data feeder" ); _exit( 2 ); }
      _exit( 0 );
      }
    if( pid < 0 )			// parent
      { show_fork_error( "data feeder" ); return false; }
    close( fda[1] );
    children.pid[0] = pid;
    children.pid[1] = 0;
    }
  return true;
  }


void show_help_addr()
  {
  std::printf( "\nReport bugs to zutils-bug@nongnu.org\n"
               "Zutils home page: http://www.nongnu.org/zutils/zutils.html\n" );
  }


void show_version( const char * const Util_name )
  {
  if( !Util_name || !Util_name[0] )
    std::printf( "%s %s\n", Program_name, PROGVERSION );
  else
    std::printf( "%s (%s) %s\n", Util_name, program_name, PROGVERSION );
  std::printf( "Copyright (C) %s Antonio Diaz Diaz.\n", program_year );
  std::printf( "License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>\n"
               "This is free software: you are free to change and redistribute it.\n"
               "There is NO WARRANTY, to the extent permitted by law.\n" );
  }


void show_error( const char * const msg, const int errcode, const bool help )
  {
  if( verbosity >= 0 )
    {
    if( msg && msg[0] )
      {
      std::fprintf( stderr, "%s: %s", util_name, msg );
      if( errcode > 0 )
        std::fprintf( stderr, ": %s", std::strerror( errcode ) );
      std::fprintf( stderr, "\n" );
      }
    if( help )
      std::fprintf( stderr, "Try '%s --help' for more information.\n",
                    invocation_name );
    }
  }


void show_error2( const char * const msg, const char * const name )
  {
  if( verbosity >= 0 )
    std::fprintf( stderr, "%s: %s '%s': %s.\n",
                  util_name, msg, name, std::strerror( errno ) );
  }


void show_close_error( const char * const prog_name )
  {
  if( verbosity >= 0 )
    std::fprintf( stderr, "%s: Can't close output of %s: %s.\n",
                  util_name, prog_name, std::strerror( errno ) );
  }


void show_exec_error( const char * const prog_name )
  {
  if( verbosity >= 0 )
    std::fprintf( stderr, "%s: Can't exec '%s': %s.\n",
                  util_name, prog_name, std::strerror( errno ) );
  }


void show_fork_error( const char * const prog_name )
  {
  if( verbosity >= 0 )
    std::fprintf( stderr, "%s: Can't fork '%s': %s.\n",
                  util_name, prog_name, std::strerror( errno ) );
  }


void internal_error( const char * const msg )
  {
  if( verbosity >= 0 )
    std::fprintf( stderr, "%s: internal error: %s.\n", util_name, msg );
  std::exit( 3 );
  }


int test_format( const int infd, const uint8_t ** const magic_datap,
                 int * const magic_sizep )
  {
  enum { buf_size = 5 };
  static uint8_t buf[buf_size];
  int i = 0;
  if( readblock( infd, buf, 1 ) == 1 )
    {
    ++i;
    if( buf[0] == bzip2_magic[0] )
      {
      if( readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == bzip2_magic[1] &&
          readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == bzip2_magic[2] )
        { *magic_datap = bzip2_magic; *magic_sizep = bzip2_magic_size;
          return fmt_bz2; }
      }
    else if( buf[0] == gzip_magic[0] )
      {
      if( readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == gzip_magic[1] )
        { *magic_datap = gzip_magic; *magic_sizep = gzip_magic_size;
          return fmt_gz; }
      }
    else if( buf[0] == lzip_magic[0] )
      {
      if( readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == lzip_magic[1] &&
          readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == lzip_magic[2] &&
          readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == lzip_magic[3] )
        { *magic_datap = lzip_magic; *magic_sizep = lzip_magic_size;
          return fmt_lz; }
      }
    else if( buf[0] == xz_magic[0] )
      {
      if( readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == xz_magic[1] &&
          readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == xz_magic[2] &&
          readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == xz_magic[3] &&
          readblock( infd, &buf[i], 1 ) == 1 && buf[i++] == xz_magic[4] )
        { *magic_datap = xz_magic; *magic_sizep = xz_magic_size;
          return fmt_xz; }
      }
    }
  *magic_datap = buf; *magic_sizep = i;
  return -1;
  }


int wait_for_child( const pid_t pid, const char * const name,
                    const int eretval )
  {
  int status;
  while( waitpid( pid, &status, 0 ) == -1 )
    {
    if( errno != EINTR )
      {
      if( verbosity >= 0 )
        std::fprintf( stderr, "%s: Error waiting termination of '%s': %s.\n",
                      util_name, name, std::strerror( errno ) );
      _exit( eretval );
      }
    }

  if( WIFEXITED( status ) )
    {
    const int tmp = WEXITSTATUS( status );
    if( eretval == 1 && tmp == 1 ) return 2;		// for ztest
    return tmp;
    }
  return eretval;
  }


int child_status( const pid_t pid, const char * const name,
                  const int eretval )
  {
  int status;
  while( true )
    {
    const int tmp = waitpid( pid, &status, WNOHANG );
    if( tmp == -1 && errno != EINTR )
      {
      if( verbosity >= 0 )
        std::fprintf( stderr, "%s: Error checking status of '%s': %s.\n",
                      util_name, name, std::strerror( errno ) );
      _exit( eretval );
      }
    if( tmp == 0 ) return -1;			// child not terminated
    if( tmp == pid ) break;			// child terminated
    }

  if( WIFEXITED( status ) ) return WEXITSTATUS( status );
  return eretval;
  }
