/* Copyright (C) 2008  Emmanuel Varoquaux

   This file is part of XOS.

   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/>. */

/* Construit une image de la disquette de demarrage.
 *
 * Utilisation : build [options] bootsect setup kernel [rootdir]
 *   bootsect  le secteur de boot
 *   setup     le code de chargement du noyau
 *   kernel    le noyau
 *   rootdir   le repertoire a copier sur la disquette
 * Options :
 *   -b filename  lire la liste des secteurs defectueux depuis le fichier
 *            filename. Si filename est -, lit la liste des secteurs
 *            defectueux depuis l'entree standard.
 *   -o filename  ecrire l'image dans le fichier filename, au lieu de la sortie
 *            standard.
 *   -q       mode silencieux. */

#define _GNU_SOURCE

#include "format.h"
#include "floppy.h"

#include <getopt.h>
#include <error.h>
#include <string.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <errno.h>

/* nombre de secteurs occupes */
#define N_SEC(size) (((size) + (BYTSPERSEC - 1)) / BYTSPERSEC)

struct write_reserved_args_struct {
  const char *setup_filename;
  const char *kernel_filename;
};

struct build_args_struct {
  char *bootsect_filename;
  char *setup_filename;
  char *kernel_filename;
  char *rootdir_filename;
  char *badblocks_filename;
  FILE *output;
};

struct stats_struct {
  off_t setup_size;
  off_t kernel_size;
  unsigned short reserved_sectors;
};

int quiet;

static void die(const char *s)
{
  error(EXIT_FAILURE, errno, s);
}

static void __attribute__ ((noreturn)) die_bad_usage(const char *s)
{
  if (s)
    error(0, 0, s);
  else
    fflush(stdout);
  fprintf(stderr, "Try `%s --help' for more information.\n", program_invocation_name);
  exit(EXIT_FAILURE);
}

static void print_help()
{
  printf("Usage: %s [OPTION]... BOOTSECT SETUP KERNEL [ROOTDIR]\n", program_invocation_name);
  puts("Create a bootable floppy disk image of XOS.\n"
       "\n"
       "Mandatory arguments to long options are mandatory for short options too.\n"
       "  -b, --badblocks-file=FILE  read bad blocks from FILE\n"
       "  -h, --help                 display this help and exit\n"
       "  -o, --output=FILE          write output in FILE instead of standard output\n"
       "  -q, --quiet                quiet output");
}

static void read_bootsect(const char *filename, char bootsect[])
{
  FILE *fp;
  int n;

  if (!(fp = fopen(filename, "r")))
    die(filename);
  n = fread(bootsect, 1, 512, fp);
  if (ferror(fp))
    die("fread");
  if (n < 512)
    memset(bootsect + n, 0, 512 - n);
  fclose(fp);
}

static inline void mark_badblock(char *badblocks, unsigned short n)
{
  if (n >= SECTORS)
    return;
  badblocks[n >> 3] |= 1 << (n & 7);
}

static void read_badblocks(const char *filename, char *badblocks)
{
  FILE *fp;
  unsigned short n;

  if (!strcmp(filename, "-"))
    fp = stdin;
  else
    if (!(fp = fopen(filename, "r")))
      die(filename);
  while (fscanf(fp, "%hu", &n) == 1)
    mark_badblock(badblocks, n);
  if (ferror(fp))
    die("fscanf");
  if (!feof(fp))
    error(EXIT_FAILURE, 0, "Failed to read bad blocks");
  if (fp != stdin)
    fclose(fp);
}

static unsigned short compute_reserved_sectors(off_t setup_size, off_t kernel_size)
{
  unsigned short reserved_sectors;

  reserved_sectors = N_SEC(512); /* secteur de demarrage */
  reserved_sectors += N_SEC(setup_size); /* setup */
  reserved_sectors += N_SEC(kernel_size); /* noyau */
  return reserved_sectors;
}

/* Ecrit un fichier sur la disquette, par secteurs entiers. */
static void write_file(const char *filename, struct output_struct *output, void (*write_sectors)(const void *, unsigned short, struct output_struct *))
{
  FILE *fp;
  char buf[BYTSPERSEC];
  size_t n;

  if (!(fp = fopen(filename, "r")))
    die(filename);
  do {
    n = fread(buf, 1, BYTSPERSEC, fp);
    if (ferror(fp))
      die("fread");
    if (n) {
      if (n < BYTSPERSEC)
        memset(buf + n, 0, BYTSPERSEC - n);
      write_sectors(buf, 1, output);
    }
  }
  while (!feof(fp));
  fclose(fp);
}

/* Ecrit la partie reservee de la disquette FAT, sauf le secteur de demarrage
 * qui est deja ecrit. */
static void write_reserved(void (*write_sectors)(const void *, unsigned short, struct output_struct *), void *arg, struct output_struct *output)
{
  struct write_reserved_args_struct *args;

  args = (struct write_reserved_args_struct *)arg;
  write_file(args->setup_filename, output, write_sectors);
  write_file(args->kernel_filename, output, write_sectors);
}

static void build(const struct build_args_struct *args, struct stats_struct *stats)
{
  struct stat statbuf;
  char bootsect[512];
  char badblocks[(SECTORS + 7) >> 3];
  struct write_reserved_args_struct write_reserved_args;

  /* lecture du secteur de demarrage */
  if (stat(args->bootsect_filename, &statbuf) == -1)
    die(args->bootsect_filename);
  if (S_ISDIR(statbuf.st_mode)) {
    errno = EISDIR;
    die(args->bootsect_filename);
  }
  read_bootsect(args->bootsect_filename, bootsect);

  /* lecture des blocs defectueux */
  memset(badblocks, 0, sizeof badblocks);
  if (args->badblocks_filename) {
    if (strcmp(args->badblocks_filename, "-")) {
      if (stat(args->badblocks_filename, &statbuf) == -1)
        die(args->badblocks_filename);
      if (S_ISDIR(statbuf.st_mode)) {
        errno = EISDIR;
        die(args->badblocks_filename);
      }
    }
    read_badblocks(args->badblocks_filename, badblocks);
  }

  /* taille de setup */
  if (stat(args->setup_filename, &statbuf) == -1)
    die(args->setup_filename);
  if (S_ISDIR(statbuf.st_mode)) {
    errno = EISDIR;
    die(args->setup_filename);
  }
  stats->setup_size = statbuf.st_size;

  /* taille du noyau */
  if (stat(args->kernel_filename, &statbuf) == -1)
    die(args->kernel_filename);
  if (S_ISDIR(statbuf.st_mode)) {
    errno = EISDIR;
    die(args->kernel_filename);
  }
  stats->kernel_size = statbuf.st_size;

  stats->reserved_sectors = compute_reserved_sectors(stats->setup_size, stats->kernel_size);
  write_reserved_args.setup_filename = args->setup_filename;
  write_reserved_args.kernel_filename = args->kernel_filename;

  format(args->output, stats->reserved_sectors, bootsect, write_reserved, &write_reserved_args, args->rootdir_filename, badblocks);
}

static void show_statistics(const struct stats_struct *stats, FILE *stream)
{
  fprintf(stream, "Boot sector:    %7lu bytes  %4hu sector(s)\n", 512UL, N_SEC(512));
  fprintf(stream, "Setup:          %7lu bytes  %4hu sector(s)\n", stats->setup_size, (unsigned short)N_SEC(stats->setup_size));
  fprintf(stream, "Kernel:         %7lu bytes  %4hu sector(s)\n", stats->kernel_size, (unsigned short)N_SEC(stats->kernel_size));
  fprintf(stream, "Total reserved:                %4hu sector(s)  %2u%%\n", stats->reserved_sectors, stats->reserved_sectors * 100 / SECTORS);
}

int main(int argc, char **argv)
{
  static const struct option longopts[] = {
    {"badblocks-file", required_argument, NULL, 'b'},
    {"help", no_argument, NULL, 'h'},
    {"output", required_argument, NULL, 'o'},
    {"quiet", no_argument, NULL, 'q'},
    {NULL, 0, NULL, 0}
  };

  char *badblocks_filename;
  char *output_filename;
  struct build_args_struct args;
  struct stats_struct stats;
  int c;

  /* options par defaut */
  quiet = 0;
  badblocks_filename = NULL;
  output_filename = NULL;

  /* lecture des options de la ligne de commande */
  while ((c = getopt_long(argc, argv, "b:ho:q", longopts, NULL)) != -1)
    switch (c) {
    case 'b':
      if (badblocks_filename)
        free(badblocks_filename);
      if (!(badblocks_filename = strdup(optarg)))
        die("strdup");
      break;
    case 'h':
      print_help();
      return EXIT_SUCCESS;
    case 'o':
      if (output_filename)
        free(output_filename);
      if (!(output_filename = strdup(optarg)))
        die("strdup");
      break;
    case 'q':
      quiet = 1;
      break;
    case '?':
      die_bad_usage(NULL);
    }

  /* lecture des autres arguments */
  if (optind == argc)
    die_bad_usage("missing bootsect operand");
  if (!(args.bootsect_filename = strdup(argv[optind++])))
    die("strdup");
  if (optind == argc)
    die_bad_usage("missing setup operand");
  if (!(args.setup_filename = strdup(argv[optind++])))
    die("strdup");
  if (optind == argc)
    die_bad_usage("missing kernel operand");
  if (!(args.kernel_filename = strdup(argv[optind++])))
    die("strdup");
  if (optind < argc) {
    if (!(args.rootdir_filename = strdup(argv[optind])))
      die("strdup");
  }
  else
    args.rootdir_filename = NULL;

  args.badblocks_filename = badblocks_filename;
  if (output_filename) {
    if (!(args.output = fopen(output_filename, "w+")))
      die(output_filename);
    free(output_filename);
  }
  else
    args.output = stdout;

  build(&args, &stats);

  if (!quiet)
    show_statistics(&stats, args.output == stdout ? stderr : stdout);

  free(args.bootsect_filename);
  free(args.setup_filename);
  free(args.kernel_filename);
  if (args.rootdir_filename)
    free(args.rootdir_filename);
  if (args.badblocks_filename)
    free(badblocks_filename);
  if (args.output != stdout)
    fclose(args.output);

  return EXIT_SUCCESS;
}