#include "SkImage.h"
#include "SkPixmap.h"
#include "SkStream.h"
#include "image_cmp.hpp"

#include <algorithm>
#include <assert.h>
#include <cstdint>
#include <vector>

static int byte0(uint32_t p) { return (p >> 0) & 0xFF; }
static int byte1(uint32_t p) { return (p >> 8) & 0xFF; }
static int byte2(uint32_t p) { return (p >> 16) & 0xFF; }
static int byte3(uint32_t p) { return (p >> 24) & 0xFF; }

static uint32_t pack(unsigned a, unsigned b, unsigned c, unsigned d)
{
  return (a << 24) | (b << 16) | (c << 8) | (d << 0);
}

static ImageCmp pixmap_cmp(const SkPixmap& pa, const SkPixmap& pb)
{
  assert(pa.width() == pb.width());
  assert(pa.height() == pb.height());

  size_t pixcount = pa.width() * pa.height();
  int maxdiff = 0;
  size_t rgb_diffcount = 0;
  size_t pixel_diffcount = 0;

  for (int y = 0; y < pa.height(); ++y)
  {
    const uint32_t* row0 = pa.addr32(0, y);
    const uint32_t* row1 = pb.addr32(0, y);

    for (int x = 0; x < pa.width(); ++x)
    {
      const uint32_t p0 = row0[x];
      const uint32_t p1 = row1[x];
      if (p0 == p1)
      {
        continue;
      }
      // Find the max diff in this pixel's r,g,b,a.
      int rgb_diff = std::abs(byte0(p0) - byte0(p1));
      rgb_diff = std::max(rgb_diff, std::abs(byte1(p0) - byte1(p1)));
      rgb_diff = std::max(rgb_diff, std::abs(byte2(p0) - byte2(p1)));
      rgb_diff = std::max(rgb_diff, std::abs(byte3(p0) - byte3(p1)));
      maxdiff = std::max(maxdiff, rgb_diff);
      rgb_diffcount += rgb_diff;
      if (rgb_diff != 0)
      {
        ++pixel_diffcount;
      }
    }
  }

  return ImageCmp::diff(
      maxdiff,
      (float)rgb_diffcount / (pixcount * 255),
      pixel_diffcount);
}

ImageCmp image_cmp(sk_sp<SkImage> a, sk_sp<SkImage> b)
{
  if (a->width() != b->width() || a->height() != b->height())
  {
    return ImageCmp::dimensionMismatch();
  }

  // ensure we can peek at the pixels
  a = a->makeRasterImage();
  b = b->makeRasterImage();

  SkPixmap pa, pb;
  if (!a->peekPixels(&pa) || !b->peekPixels(&pb))
  {
    return ImageCmp::failedToLoad();
  }

  return pixmap_cmp(pa, pb);
}

uint32_t compute_diff0(uint32_t p0, uint32_t p1)
{
  auto a = std::abs(byte0(p0) - byte0(p1));
  auto b = std::abs(byte1(p0) - byte1(p1));
  auto c = std::abs(byte2(p0) - byte2(p1));
  return pack(0xFF, a, b, c);
}

uint32_t compute_diff1(uint32_t p0, uint32_t p1)
{
  auto d = p0 == p1 ? 0 : 0xFF;
  return pack(0xFF, d, d, d);
}

sk_sp<SkImage> make_diff_image(
    sk_sp<SkImage> a,
    sk_sp<SkImage> b,
    uint32_t (*compute_diff)(uint32_t a, uint32_t b))
{
  // ensure we can peek at the pixels
  a = a->makeRasterImage();
  b = b->makeRasterImage();

  SkPixmap pa, pb;
  if (!a->peekPixels(&pa) || !b->peekPixels(&pb))
  {
    throw "failed to peek";
  }

  const auto info = SkImageInfo::MakeN32Premul(pa.width(), pa.height());
  std::vector<uint32_t> pixels(pa.width() * pa.height());
  auto df = SkPixmap(info, pixels.data(), pa.width() * sizeof(uint32_t));

  for (int y = 0; y < df.height(); ++y)
  {
    for (int x = 0; x < df.width(); ++x)
    {
      *df.writable_addr32(x, y) = compute_diff(
          *pa.addr32(x, y),
          *pb.addr32(x, y));
    }
  }

  return SkImage::MakeRasterCopy(df);
}