#include "util/map_image_util.hpp"

#include "util/string_util.hpp"

#include <algorithm>
#include <cctype>
#include <sstream>

#define STB_IMAGE_WRITE_IMPLEMENTATION
#define STB_IMAGE_WRITE_STATIC
#include "third_party/stb_image_write.h"

namespace lm {

namespace {

std::string lowerExt(const std::string& filename)
{
  const auto dot = filename.rfind('.');
  if (dot == std::string::npos)
    return "";
  std::string ext = filename.substr(dot);
  return StringUtil::toLower(ext);
}

bool isPngExtension(const std::string& ext)
{
  return ext == ".png";
}

bool isPgmExtension(const std::string& ext)
{
  return ext == ".pgm";
}

bool pngMagic(const std::string& bytes)
{
  static const unsigned char kSig[] = {0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A};
  if (bytes.size() < 8)
    return false;
  return std::equal(std::begin(kSig), std::end(kSig), bytes.begin());
}

bool pgmMagic(const std::string& bytes)
{
  return bytes.size() >= 2 && bytes[0] == 'P' && (bytes[1] == '5' || bytes[1] == '2');
}

void applyNegate(std::vector<uint8_t>& gray)
{
  for (auto& px : gray)
    px = static_cast<uint8_t>(255 - px);
}

}  // namespace

bool MapImageUtil::decodePgm(const std::string& bytes, std::vector<uint8_t>& gray, ImageSize& size, std::string& err)
{
  if (bytes.size() < 4 || bytes[0] != 'P')
  {
    err = "not a PGM image";
    return false;
  }
  if (bytes[1] != '5')
  {
    err = "only binary PGM (P5) is supported";
    return false;
  }

  std::size_t i = 2;
  auto skipSpace = [&]() {
    while (i < bytes.size() && (bytes[i] == ' ' || bytes[i] == '\t' || bytes[i] == '\r' || bytes[i] == '\n'))
      ++i;
  };
  auto readInt = [&](int& out) -> bool {
    for (;;)
    {
      skipSpace();
      if (i < bytes.size() && bytes[i] == '#')
      {
        while (i < bytes.size() && bytes[i] != '\n')
          ++i;
        continue;
      }
      if (i >= bytes.size() || !std::isdigit(static_cast<unsigned char>(bytes[i])))
        return false;
      out = 0;
      while (i < bytes.size() && std::isdigit(static_cast<unsigned char>(bytes[i])))
      {
        out = out * 10 + (bytes[i] - '0');
        ++i;
      }
      return true;
    }
  };

  skipSpace();
  if (!readInt(size.width) || !readInt(size.height))
  {
    err = "invalid PGM header";
    return false;
  }
  int maxval = 0;
  if (!readInt(maxval))
  {
    err = "invalid PGM header";
    return false;
  }
  if (size.width <= 0 || size.height <= 0 || maxval <= 0 || maxval > 65535)
  {
    err = "invalid PGM dimensions";
    return false;
  }

  if (i < bytes.size() && bytes[i] == '\r')
    ++i;
  if (i < bytes.size() && bytes[i] == '\n')
    ++i;

  const std::size_t data_size = static_cast<std::size_t>(size.width) * static_cast<std::size_t>(size.height);
  if (bytes.size() - i < data_size)
  {
    err = "PGM data truncated";
    return false;
  }

  gray.assign(bytes.begin() + static_cast<std::ptrdiff_t>(i),
              bytes.begin() + static_cast<std::ptrdiff_t>(i + data_size));
  if (maxval != 255)
  {
    const double scale = 255.0 / static_cast<double>(maxval);
    for (auto& px : gray)
      px = static_cast<uint8_t>(std::min(255.0, px * scale));
  }
  return true;
}

bool MapImageUtil::pngDimensions(const std::string& bytes, ImageSize& size, std::string& err)
{
  if (!pngMagic(bytes))
  {
    err = "not a PNG image";
    return false;
  }
  if (bytes.size() < 24)
  {
    err = "PNG header truncated";
    return false;
  }
  auto readU32 = [&](std::size_t pos) -> int {
    return (static_cast<unsigned char>(bytes[pos]) << 24) | (static_cast<unsigned char>(bytes[pos + 1]) << 16) |
           (static_cast<unsigned char>(bytes[pos + 2]) << 8) | static_cast<unsigned char>(bytes[pos + 3]);
  };
  size.width = readU32(16);
  size.height = readU32(20);
  if (size.width <= 0 || size.height <= 0)
  {
    err = "invalid PNG dimensions";
    return false;
  }
  return true;
}

bool MapImageUtil::toPngBytes(const std::string& bytes,
                              const std::string& filename_hint,
                              int negate,
                              std::string& png_out,
                              ImageSize& size,
                              std::string& err)
{
  const std::string ext = lowerExt(filename_hint);
  const bool as_png = isPngExtension(ext) || pngMagic(bytes);
  const bool as_pgm = isPgmExtension(ext) || (!as_png && pgmMagic(bytes));

  if (as_png)
  {
    if (!pngDimensions(bytes, size, err))
      return false;
    if (negate == 0)
    {
      png_out = bytes;
      return true;
    }
    err = "PNG negate not supported";
    return false;
  }

  if (!as_pgm)
  {
    err = "image must be PNG or PGM";
    return false;
  }

  std::vector<uint8_t> gray;
  if (!decodePgm(bytes, gray, size, err))
    return false;
  if (negate != 0)
    applyNegate(gray);

  int out_len = 0;
  unsigned char* png_mem = stbi_write_png_to_mem(gray.data(), size.width, size.width, size.height, 1, &out_len);
  if (!png_mem || out_len <= 0)
  {
    err = "failed to encode PNG";
    return false;
  }
  png_out.assign(reinterpret_cast<char*>(png_mem), static_cast<std::size_t>(out_len));
  STBIW_FREE(png_mem);
  return true;
}

}  // namespace lm
