/*
 * Copyright (C) 2006 Red Hat, Inc.
 *
 * 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 2 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, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
 * 02111-1307, USA.

 * Authors:
 *      Mark McLoughlin <markmc@redhat.com>
 */

/* To compile:
 *   gcc -g -Wall -O0 $(pkg-config --cflags --libs ext2fs) e2cp.c -o e2cp
 */

/*
 * FIXME: Large file support
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <ext2fs/ext2fs.h>

typedef struct
{
  ext2_filsys  fs;
  int          src;
  int          dest;
  char        *blockbuf;
} copyctx;

typedef struct
{
  int start;
  int num;
} blkrange;


static int verbose_output = 0;
static int copied_blocks = 0;
static int zeroed_blocks = 0;

static void
verbose (const char *fmt, ...)
{
  va_list args;

  if (!verbose_output)
    return;

  va_start (args, fmt);

  vfprintf (stderr, fmt, args);

  va_end (args);
}

static inline int
check_all_zeros (const char *buf, int len)
{
  int i;

  for (i = 0; i < len && buf[i] == 0; i++);

  return i == len;
}

static int
io_blocks (copyctx *ctx, int skip, int n_blocks)
{
  int blocksize;

  blocksize = ctx->fs->blocksize;

  if (skip > 0)
    {
      if (lseek (ctx->src,  skip * blocksize, SEEK_CUR) < 0)
	return -1;
      if (lseek (ctx->dest, skip * blocksize, SEEK_CUR) < 0)
	return -1;
    }

  skip = 0;

  while (n_blocks > 0)
    {
      char    *p;
      int      left;
      ssize_t  n;

      p = ctx->blockbuf;
      left = blocksize;
      do
	{
	  if ((n = read (ctx->src, p, left)) < 0)
	    return n;
	  p += n;
	  left -= n;
	} while (left > 0);

      if (check_all_zeros (ctx->blockbuf, blocksize))
	{
	  skip++;
	}
      else
	{
	  if (skip > 0)
	    {
	      verbose ("Skipping %d zero blocks\n", skip);
	      if (lseek (ctx->dest, skip * blocksize, SEEK_CUR) < 0)
		return -1;
	      skip = 0;
	    }

	  p = ctx->blockbuf;
	  left = blocksize;
	  do
	    {
	      if ((n = write (ctx->dest, ctx->blockbuf, blocksize)) < 0)
		return n;
	      p += n;
	      left -= n;
	    } while (left > 0);
	}

      n_blocks--;
    }

  if (skip > 0)
    {
      verbose ("Skipping %d zero blocks\n", skip);
      if (lseek (ctx->dest, skip * blocksize, SEEK_CUR) < 0)
	return -1;
    }

  return 0;
}

static int
zero_blocks (copyctx *ctx, blkrange *range)
{
  int err;

  if (range->num == 1)
    verbose ("Zeroing block %d\n", range->start);
  else
    verbose ("Zeroing blocks %d-%d\n",
	     range->start, range->start + range->num - 1);

  zeroed_blocks += range->num;

  if ((err = io_blocks (ctx, range->num, 0)) < 0)
    {
      fprintf (stderr, "Failed to zero block range %d-%d: %s\n",
	       range->start, range->start + range->num - 1,
	       error_message (errno));
      return err;
    }

  return 0;
}

static int
copy_blocks (copyctx *ctx, blkrange *range)
{
  int err;

  if (range->num == 1)
    verbose ("Copying block %d\n", range->start);
  else
    verbose ("Copying blocks %d-%d\n",
	     range->start, range->start + range->num - 1);

  copied_blocks += range->num;

  if ((err = io_blocks (ctx, 0, range->num)) < 0)
    {
      fprintf (stderr, "Failed to copy block range %d-%d: %s\n",
	       range->start, range->start + range->num - 1,
	       error_message (errno));
      return err;
    }

  return 0;
}

static int
copy_ext2fs (copyctx *ctx)
{
  blkrange copy_range;
  blkrange zero_range;
  int      first_block;
  int      num_blocks;
  int      err;
  int      i;

  first_block = ctx->fs->super->s_first_data_block;
  num_blocks  = ctx->fs->super->s_blocks_count;

  err = ext2fs_read_block_bitmap (ctx->fs);
  if (err)
    {
      fprintf (stderr, "Error while reading block bitmap: %s\n",
               error_message (err));
      return 1;
    }

  if (first_block > 0)
    {
      zero_range.start = 0;
      zero_range.num   = first_block;
      if (zero_blocks (ctx, &zero_range) != 0)
	return 1;
    }

  copy_range.start = copy_range.num = -1;
  zero_range.start = zero_range.num = -1;

  for (i = first_block; i < num_blocks; i++)
    {
      if (ext2fs_fast_test_block_bitmap (ctx->fs->block_map, i))
	{
	  if (zero_range.start != -1)
	    {
	      if (zero_blocks (ctx, &zero_range) != 0)
		return 1;
	      zero_range.start = zero_range.num = -1;
	    }

	  if (copy_range.start == -1)
	    {
	      copy_range.start = i;
	      copy_range.num = 0;
	    }

	  copy_range.num++;
	}
      else
	{
	  if (copy_range.start != -1)
	    {
	      if (copy_blocks (ctx, &copy_range) != 0)
		return 1;
	      copy_range.start = copy_range.num = -1;
	    }

	  if (zero_range.start == -1)
	    {
	      zero_range.start = i;
	      zero_range.num = 0;
	    }

	  zero_range.num++;
	}
    }

  err = 0;

  if (copy_range.start != -1)
    err = copy_blocks (ctx, &copy_range);
  else if (zero_range.start != -1)
    err = zero_blocks (ctx, &zero_range);

  verbose ("Total %d, Copied %d, zeroed %d\n",
	   num_blocks, copied_blocks, zeroed_blocks);

  return err;
}

int
main (int argc, char **argv)
{
  const char  *src;
  const char  *dest;
  errcode_t    err;
  copyctx      ctx;
  int          mount_flags;

  initialize_ext2_error_table ();

  if (getenv ("E2CP_VERBOSE_OUTPUT"))
    verbose_output = 1;

  if (argc != 2 && argc != 3)
    {
      fprintf (stderr, "usage: %s <src> [<dest>]\n", argv[0]);
      return 1;
    }

  src = argv[1];

  if (argc == 2)
    dest = NULL; /* write to stdout */
  else
    dest = argv[2];

  err = ext2fs_check_if_mounted (src, &mount_flags);
  if (err)
    {
      fprintf (stderr, "Error determining whether '%s' is mounted: %s\n",
	       src, error_message (err));
      return 1;
    }

  if (mount_flags & (EXT2_MF_MOUNTED|EXT2_MF_BUSY))
    {
      fprintf (stderr, "Source '%s' is currently %s\n",
	       src, mount_flags & EXT2_MF_MOUNTED ? "mounted" : "busy");
      return 1;
    }

  err = ext2fs_open (src, 0, 0, 0, unix_io_manager, &ctx.fs);
  if (err)
    {
      fprintf (stderr, "Error while trying to open '%s': %s\n",
               src, error_message (err));
      return 1;
    }

  ctx.src = open (src, O_RDONLY);
  if (ctx.src < 0)
    {
      fprintf (stderr, "Error while trying to open '%s': %s\n",
	       src, error_message (errno));
      ext2fs_close (ctx.fs);
      return 1;
    }

  verbose ("Copying from %s\n", src);

  if (dest == NULL)
    {
      ctx.dest = 1;
      verbose ("Copying to stdout\n");
    }
  else
    {
      ctx.dest = open (dest, O_CREAT|O_TRUNC|O_WRONLY, 0600);
      if (ctx.dest < 0)
	{
	  fprintf (stderr, "Error while trying to open '%s': %s\n",
		   dest, error_message (errno));
	  if (ctx.src != 0)
	    close (ctx.src);
	  ext2fs_close (ctx.fs);
	  return 1;
	}
      verbose ("Copying to %s\n", dest);
    }

  verbose ("Block size is %d\n", ctx.fs->blocksize);

  ctx.blockbuf = (char *) malloc (ctx.fs->blocksize);
  if (ctx.blockbuf == NULL)
    {
      fprintf (stderr, "Failed to allocate block buffer\n");
      if (ctx.src != 0)
	close (ctx.src);
      if (ctx.dest != 1)
	close (ctx.dest);
      ext2fs_close (ctx.fs);
      return 1;
    }

  err = copy_ext2fs (&ctx);

  free (ctx.blockbuf);
  if (ctx.src != 0)
    close (ctx.src);
  if (ctx.dest != 1)
    close (ctx.dest);
  ext2fs_close (ctx.fs);

  return err;
}