#include <stdlib.h>
#include <malloc.h>
#include <png.h>

#include <SDL/SDL.h>
#include <SDL/SDL_image.h>

#if SDL_BYTEORDER == SDL_BIG_ENDIAN
# define RMASK  0xff000000
# define BMASK  0x00ff0000
# define GMASK  0x0000ff00
# define AMASK  0x000000ff
#else
# define RMASK  0x000000ff
# define BMASK  0x0000ff00
# define GMASK  0x00ff0000
# define AMASK  0xff0000
#endif
#define RGBAMASK RMASK, GMASK, BMASK, AMASK

/* Logging macros. */
#define LOG(str, args...)(fprintf(stdout, str"\n", ## args))
#define WARN(str, args...)(fprintf(stderr, "Warning: "str"\n", ## args))
#define ERR(str, args...) { fprintf(stderr, "ERROR %s:%d: "str"\n", \
      __FILE__, __LINE__, ## args); return EXIT_FAILURE; }

#define WS 6
#define HS 6

static int save_png(SDL_Surface* surface, const char* file);
static int write_png(const char* file_name, png_bytep* rows, int w, int h, int colourtype, int bitdepth);

int main(int argc, char* argv[]) {
  int i, ws, hs;
  unsigned int sflags, salpha;
  char file[8];
  SDL_Surface* final, *tmp, **load;
  SDL_Rect r;

  /* Init variables. */
  r.w = r.h = 0;

  if(argc == 2) {
    ws = hs = atoi(argv[1]);
  }
  else if(argc == 3) {
    ws = atoi(argv[1]);
    hs = atoi(argv[2]);
  } else {
    ws = WS;
    hs = HS;
  }

  load  = NULL;
  tmp   = NULL;
  final = NULL;

  /* Init SDL. */
  if(SDL_Init(SDL_INIT_VIDEO)) ERR("Initializing SDL: %s", SDL_GetError());
  /* Create the window. */
  tmp = SDL_SetVideoMode(320, 240, 0, SDL_NOFRAME);
  if(tmp == NULL) ERR("Initializing video surface: %s", SDL_GetError());

  /* Open RAM for the images. */
  load = malloc(sizeof(SDL_Surface*)*(ws*hs));
  if(load == NULL) ERR("Out of RAM");

  /* Load all the images to RAM. */
  for(i = 0; i < (ws*hs); i++) {
    /* Filenames will be in the sequence of: 000.png, 001.png, ..., 045.png, etc. */
    sprintf(file, "%d%d%d.png", i/100, (i%100)/10, i%10);

    /* Load the image properly. */
    tmp  = IMG_Load(file);
    if(tmp == NULL) ERR("Problem loading file '%s': %s", file, IMG_GetError());
    sflags = tmp->flags & (SDL_SRCALPHA | SDL_SRCCOLORKEY);
    salpha = tmp->format->alpha;
    if(sflags & SDL_SRCALPHA)
      SDL_SetAlpha(tmp, 0, SDL_ALPHA_OPAQUE);
    if(sflags & SDL_SRCCOLORKEY)
      SDL_SetColorKey(tmp, 0, tmp->format->colorkey);

    load[i] = tmp;

    /* Check if size has changed. */
    if(r.w == 0 && r.h == 0) {
      r.w = load[i]->w;
      r.h = load[i]->h;
    }
    else if((r.w != load[i]->w) || (r.h != load[i]->h))
      ERR("File '%s' is not of the same dimensions as the files before!", file);

    /* Create the suface if it hasn't been created already. */
    if(!final) {
      final = SDL_CreateRGBSurface(SDL_SWSURFACE | SDL_SRCALPHA, ws*r.w, hs*r.h,
            load[i]->format->BitsPerPixel, RGBAMASK);
      if(!final) ERR("Problem creating RGB Surface: %s", SDL_GetError());
      tmp = final;
    }
    /* New position. */
    r.y = r.h * (i/ws);
    r.x = r.w * (i%ws);
    if(SDL_BlitSurface(load[i], NULL, final, &r))
      ERR("Problem blitting surface '%s' to final surface: %s", file, SDL_GetError());
    SDL_FreeSurface(load[i]);
  }
  /* Draw the result and cleanup. */
  save_png(final, "sprite.png");
  SDL_FreeSurface(final);
  free(load);

  SDL_Quit();

  return EXIT_SUCCESS;
}

static int save_png(SDL_Surface* surface, const char* file) {
  static unsigned char** ss_rows;
  static int ss_size;
  static int ss_w, ss_h;
  SDL_Surface* ss_surface;
  SDL_Rect ss_rect;
  int r, i;
  int alpha = 0;
  int pixel_bits = 32;

  unsigned surf_flags;
  unsigned surf_alpha;

  ss_rows = 0;
  ss_size = 0;
  ss_surface = 0;

  ss_w = surface->w;
  ss_h = surface->h;

  if(surface->format->Amask) {
    alpha = 1;
    pixel_bits = 32;
  } else {
    pixel_bits = 24;
  }

  ss_surface = SDL_CreateRGBSurface(SDL_SWSURFACE | SDL_SRCALPHA, ss_w, ss_h, pixel_bits, RGBAMASK);
  
  if(ss_surface == 0) {
    return -1;
  }

  surf_flags = surface->flags & (SDL_SRCALPHA | SDL_SRCCOLORKEY);
  surf_alpha = surface->format->alpha;
  if(surf_flags & SDL_SRCALPHA)
    SDL_SetAlpha(surface, 0, SDL_ALPHA_OPAQUE);
  if(surf_flags & SDL_SRCCOLORKEY)
    SDL_SetColorKey(surface, 0, surface->format->colorkey);

  ss_rect.x = 0;
  ss_rect.y = 0;
  ss_rect.w = ss_w;
  ss_rect.h = ss_h;
  SDL_BlitSurface(surface, &ss_rect, ss_surface, 0);

  if(ss_size == 0) {
    ss_size = ss_h;
    ss_rows = (unsigned char**)malloc(sizeof(unsigned char*) * ss_size);
    if(ss_rows == 0) {
      return -1;
    }
  }
  if(surf_flags & SDL_SRCALPHA)
    SDL_SetAlpha(surface, SDL_SRCALPHA, (Uint8)surf_alpha);
  if(surf_flags & SDL_SRCCOLORKEY)
    SDL_SetColorKey(surface, SDL_SRCCOLORKEY, surface->format->colorkey);

  for(i = 0; i < ss_h; i++) {
    ss_rows[i] = ((unsigned char*)ss_surface->pixels) + i * ss_surface->pitch;
  }

  if(alpha) {
    r = write_png(file, ss_rows, surface->w, surface->h, PNG_COLOR_TYPE_RGB_ALPHA, 8);
  } else {
    r = write_png(file, ss_rows, surface->w, surface->h, PNG_COLOR_TYPE_RGB, 8);
  }

  free(ss_rows);
  SDL_FreeSurface(ss_surface);
  ss_surface = NULL;

  return r;
}

static int write_png(const char* file_name, png_bytep* rows, int w, int h, int colourtype, int bitdepth) {
  png_structp png_ptr;
  png_infop info_ptr;
  FILE* fp = NULL;
  char* doing = "Open for writing";

  if(!(fp = fopen(file_name, "wb"))) goto fail;

  doing = "Create png info struct";
  if(!(info_ptr = png_create_info_struct(png_ptr))) goto fail;
  if(setjmp(png_jmpbuf(png_ptr))) goto fail;

  doing = "Init IO";
  png_init_io(png_ptr, fp);
  png_set_IHDR(png_ptr, info_ptr, w, h, bitdepth, colourtype, PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
  
  doing = "Write info";
  png_write_end(png_ptr, info_ptr);

  doing = "Write image";
  png_write_image(png_ptr, rows);

  doing = "Write end";
  png_write_end(png_ptr, NULL);

  doing = "Closing file";
  if(0 != fclose(fp)) goto fail;

  return 0;

fail:
  WARN("write_png: Could not %s", doing);
  return -1;
}