/*
 * bfl - BitFLip image format
 * Copyright (C) 2025  Slendi <slendi@socopon.com>

 * 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 <https://www.gnu.org/licenses/>.
 *
 * Include the following in a source file before using the library:
 *   #define BFL_IMPLEMENTATION
 * to create the implementation.
 *
 * Build flags:
 * - BFL_NO_LZSS:     disable LZSS on encode.
 * - BFL_NO_COINFLIP: disable coinflip on encode.
 * - BFL_NO_COMPRESS: disable all compression on encode.
 */

#pragma once

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <span>
#include <vector>

namespace bfl {

enum : uint8_t {
  FLAG_HAS_ALPHA = 0x01,
  FLAG_IMG_RAW = 0x02,
  FLAG_TRA_RAW = 0x04,
  FLAG_IMG_NOLZ = 0x08,
  FLAG_TRA_NOLZ = 0x10,
};

struct Header {
  uint16_t w{}, h{};
  uint8_t flags{};
  uint32_t img_len{}, tra_len{};
};

struct View {
  Header hdr{};
  std::span<uint8_t const> img_c{};
  std::span<uint8_t const> tra_c{};
};

auto parse_bfl(std::span<uint8_t const> data) -> std::optional<View>;

struct Bitmap {
  uint16_t width{}, height{};
  std::vector<bool> image_data;
  std::optional<std::vector<bool>> transparency_data;

  [[nodiscard]] static auto from_rgba(std::span<uint32_t const> data, int w,
                                      int h) -> Bitmap;
  [[nodiscard]] static auto decode(std::span<uint8_t const> data) -> Bitmap;

  [[nodiscard]] auto to_rgba() const -> std::vector<uint32_t>;
  [[nodiscard]] auto encode() const -> std::vector<uint8_t>;
};

} // namespace bfl

#ifdef BFL_IMPLEMENTATION

#ifdef BFL_NO_COMPRESS
#ifndef BFL_NO_LZSS
#define BFL_NO_LZSS
#endif // BFL_NO_LZSS
#ifndef BFL_NO_COINFLIP
#define BFL_NO_COINFLIP
#endif // BFL_NO_COINFLIP
#endif // BFL_NO_COMPRESS

#include <algorithm>
#include <type_traits>
#include <utility>

namespace bfl {

template <class T> inline void put_le(std::vector<uint8_t> &buf, T v) noexcept {
  static_assert(std::is_unsigned_v<T>);
  for (size_t i{}; i < sizeof(T); i++)
    buf.push_back(static_cast<uint8_t>((v >> (i * 8)) & 0xFF));
}

template <class T>
inline T get_le(std::span<uint8_t const> s, size_t off) noexcept {
  static_assert(std::is_unsigned_v<T>);
  T v = 0;
  for (size_t i{}; i < sizeof(T); i++)
    v |= (static_cast<T>(s[off + i]) << (i * 8));
  return v;
}

#ifndef BFL_NO_LZSS
std::vector<uint8_t> lzss_compress(std::span<uint8_t const> in) {
  int const W{4096}, LA{18}, MIN{3};
  std::vector<uint8_t> out;
  out.reserve(in.size() / 8 + 16);
  size_t i{};
  while (i < in.size()) {
    uint8_t flag{};
    size_t flag_pos{out.size()};
    out.push_back(0);
    for (int bit{}; bit < 8 && i < in.size(); bit++) {
      size_t best_len{}, best_off{};
      size_t wnd_start{(i > (size_t)W) ? i - W : 0};
      size_t max_len{std::min((size_t)LA, in.size() - i)};
      for (size_t p{i}; p-- > wnd_start;) {
        size_t l{};
        while (l < max_len && in[p + l] == in[i + l])
          l++;
        if (l >= (size_t)MIN && l > best_len) {
          best_len = l;
          best_off = i - p;
          if (best_len == (size_t)LA)
            break;
        }
      }
      if (best_len >= (size_t)MIN && best_off >= 1 && best_off <= 4095) {
        uint8_t b0 =
            static_cast<uint8_t>(((best_len - 3) << 4) | (best_off >> 8));
        uint8_t b1{static_cast<uint8_t>(best_off & 0xFF)};
        out.push_back(b0);
        out.push_back(b1);
        i += best_len;
      } else {
        flag |= (1u << bit);
        out.push_back(in[i++]);
      }
    }
    out[flag_pos] = flag;
  }
  return out;
}
#endif // BFL_NO_LZSS

std::vector<uint8_t> lzss_decompress(std::span<uint8_t const> in) {
  std::vector<uint8_t> out;
  out.reserve(in.size() * 2);
  size_t i{};
  while (i < in.size()) {
    uint8_t flag{in[i++]};
    for (int bit{}; bit < 8 && i < in.size(); bit++) {
      if ((flag >> bit) & 1) {
        out.push_back(in[i++]);
      } else {
        if (i + 1 >= in.size())
          return out;
        uint8_t b0{in[i++]}, b1{in[i++]};
        size_t len{static_cast<size_t>((b0 >> 4) + 3)};
        size_t off{static_cast<size_t>(((b0 & 0x0F) << 8) | b1)};
        if (off == 0 || off > out.size())
          return out;
        size_t src{out.size() - off};
        for (size_t k{}; k < len; k++)
          out.push_back(out[src + k]);
      }
    }
  }
  return out;
}

std::vector<uint8_t> bits_to_bytes(std::vector<bool> const &data) noexcept {
  std::vector<uint8_t> bytes((data.size() + 7) / 8, 0);
  for (size_t i{}; i < data.size(); i++)
    if (data[i])
      bytes[i / 8] |= (1u << (i % 8));
  return bytes;
}

std::vector<bool> decode_raw(std::span<uint8_t const> in, size_t total_bits) {
  std::vector<bool> out(total_bits, false);
  for (size_t i{}; i < total_bits; i++) {
    uint8_t byte = (i / 8 < in.size()) ? in[i / 8] : 0;
    out[i] = ((byte >> (i % 8)) & 1) != 0;
  }
  return out;
}

#ifndef BFL_NO_COINFLIP
std::vector<uint8_t> coinflip_encode(std::vector<bool> const &data,
                                     bool initial,
                                     std::vector<bool> const *transparency,
                                     bool allow_trick) {
  std::vector<uint8_t> out;
  out.reserve(data.size() / 4 + 8);
  out.push_back(static_cast<uint8_t>(initial));

  bool state{initial};
  size_t count{}, i{};

  auto flush{[&](uint8_t n) { out.push_back(n); }};

  while (i < data.size()) {
    bool is_transparent{transparency && (*transparency)[i]};

    if (allow_trick && is_transparent) {
      while (i < data.size() && (*transparency)[i]) {
        if (count == 255) {
          flush(255);
          state = !state;
          count = 0;
        }
        count++;
        i++;
      }
      continue;
    }

    bool bit{data[i]};
    if (bit == state) {
      count++;
      i++;
      if (count == 255 && i < data.size()) {
        flush(255);
        out.push_back(0);
        count = 0;
      }
    } else {
      flush(static_cast<uint8_t>(count));
      state = !state;
      count = 1;
      i++;
    }
  }
  if (count)
    flush(static_cast<uint8_t>(count));
  return out;
}
#endif // BFL_NO_COINFLIP

std::vector<bool> rle_decode(std::span<uint8_t const> in, size_t total_bits) {
  std::vector<bool> out(total_bits, false);
  if (in.size() < 2)
    return out;

  bool state{in[0] != 0};
  size_t produced{0}, j{1};

  while (produced < total_bits && j < in.size()) {
    uint8_t n{in[j++]};
    if (n == 0)
      continue;
    size_t emit{std::min<size_t>(n, total_bits - produced)};
    for (size_t k{}; k < emit; k++)
      out[produced + k] = state;
    produced += emit;
    if (j < in.size() && in[j] == 0)
      j++;
    else
      state = !state;
  }
  return out;
}

auto parse_bfl(std::span<uint8_t const> data) -> std::optional<View> {
  if (data.size() < 3 + 2 + 2 + 1 + 4 + 4)
    return std::nullopt;
  if (!(data[0] == 'B' && data[1] == 'F' && data[2] == 'L'))
    return std::nullopt;

  size_t off{3};
  View v;
  v.hdr.w = get_le<uint16_t>(data, off);
  off += 2;
  v.hdr.h = get_le<uint16_t>(data, off);
  off += 2;
  v.hdr.flags = data[off++];
  v.hdr.img_len = get_le<uint32_t>(data, off);
  off += 4;
  v.hdr.tra_len = get_le<uint32_t>(data, off);
  off += 4;

  if (off + v.hdr.img_len > data.size())
    return std::nullopt;
  if (off + v.hdr.img_len + v.hdr.tra_len > data.size())
    return std::nullopt;

  v.img_c = data.subspan(off, v.hdr.img_len);
  off += v.hdr.img_len;
  v.tra_c = data.subspan(off, v.hdr.tra_len);
  return v;
}

auto Bitmap::to_rgba() const -> std::vector<uint32_t> {
  auto mk{[](uint8_t r, uint8_t g, uint8_t b, uint8_t a) -> uint32_t {
    return (uint32_t)r | ((uint32_t)g << 8) | ((uint32_t)b << 16) |
           ((uint32_t)a << 24);
  }};
  size_t n{(size_t)width * (size_t)height};
  std::vector<uint32_t> out(n, 0);
  for (size_t i{}; i < n; i++) {
    bool tr{transparency_data && i < transparency_data->size() &&
            (*transparency_data)[i]};
    bool wh{i < image_data.size() && image_data[i]};
    out[i] =
        tr ? mk(0, 0, 0, 0) : (wh ? mk(255, 255, 255, 255) : mk(0, 0, 0, 255));
  }
  return out;
}

auto Bitmap::encode() const -> std::vector<uint8_t> {
  auto choose_rle_or_raw{
      [&](std::vector<bool> const &bits, bool &raw_flag, bool allow_trick) {
#ifdef BFL_NO_COINFLIP
        (void)allow_trick;
        raw_flag = true;
        return bits_to_bytes(bits);
#else
        std::vector<bool> const *mask{transparency_data ? &(*transparency_data)
                                                        : nullptr};
        auto a{coinflip_encode(bits, false, mask, allow_trick)};
        auto b{coinflip_encode(bits, true, mask, allow_trick)};
        std::vector<uint8_t> best{(a.size() <= b.size()) ? std::move(a)
                                                         : std::move(b)};
        if (best.size() * 8 > bits.size()) {
          best = bits_to_bytes(bits);
          raw_flag = true;
        }
        return best;
#endif
      }};

  std::vector<uint8_t> out;
  out.push_back('B');
  out.push_back('F');
  out.push_back('L');
  put_le<uint16_t>(out, width);
  put_le<uint16_t>(out, height);

  bool img_raw{}, tra_raw{};
  auto img_bytes =
      choose_rle_or_raw(image_data, img_raw, transparency_data.has_value());
  std::vector<uint8_t> tra_bytes;
  if (transparency_data)
    tra_bytes = choose_rle_or_raw(*transparency_data, tra_raw, false);

  auto maybe_lz{[](std::vector<uint8_t> v, bool &no_lz) {
#ifndef BFL_NO_LZSS
    auto c{lzss_compress(v)};
    if (c.size() < v.size())
      return c;
#endif // BFL_NO_LZSS
    no_lz = true;
    return v;
  }};
  bool img_no_lz{}, tra_no_lz{};
  auto img_stream{maybe_lz(std::move(img_bytes), img_no_lz)};
  std::vector<uint8_t> tra_stream;
  if (!tra_bytes.empty())
    tra_stream = maybe_lz(std::move(tra_bytes), tra_no_lz);

  uint8_t flags{};
  if (transparency_data)
    flags |= FLAG_HAS_ALPHA;
  if (img_raw)
    flags |= FLAG_IMG_RAW;
  if (tra_raw)
    flags |= FLAG_TRA_RAW;
  if (img_no_lz)
    flags |= FLAG_IMG_NOLZ;
  if (tra_no_lz)
    flags |= FLAG_TRA_NOLZ;

  out.push_back(flags);
  put_le<uint32_t>(out, static_cast<uint32_t>(img_stream.size()));
  put_le<uint32_t>(out, static_cast<uint32_t>(tra_stream.size()));
  out.insert(out.end(), img_stream.begin(), img_stream.end());
  out.insert(out.end(), tra_stream.begin(), tra_stream.end());
  return out;
}

auto Bitmap::from_rgba(std::span<uint32_t const> data, int w, int h) -> Bitmap {
  assert(w > 0 && h > 0);
  assert(static_cast<int>(data.size()) == w * h);

  Bitmap bm;
  bm.width = (uint16_t)w;
  bm.height = (uint16_t)h;
  size_t n{(size_t)w * (size_t)h};
  bm.image_data.resize(n);
  bm.transparency_data.emplace();
  bm.transparency_data->resize(n);

  for (size_t i{}; i < n; i++) {
    uint32_t px{data[i]};
    uint8_t r = (px >> 0) & 0xFF, g = (px >> 8) & 0xFF, b = (px >> 16) & 0xFF,
            a = (px >> 24) & 0xFF;
    bool white{(r == 255 && g == 255 && b == 255 && a != 0)};
    bool black{(r == 0 && g == 0 && b == 0 && a != 0)};
    bool tr{(!white && !black) || (a == 0)};
    (*bm.transparency_data)[i] = tr;
    bm.image_data[i] = white;
  }
  return bm;
}

auto Bitmap::decode(std::span<uint8_t const> data) -> Bitmap {
  Bitmap bm{};
  auto view{parse_bfl(data)};
  if (!view)
    return bm;

  auto &h{view->hdr};
  bool has_alpha{(h.flags & FLAG_HAS_ALPHA) != 0};
  bool img_was_raw{(h.flags & FLAG_IMG_RAW) != 0};
  bool tra_was_raw{(h.flags & FLAG_TRA_RAW) != 0};
  bool img_no_lz{(h.flags & FLAG_IMG_NOLZ) != 0};
  bool tra_no_lz{(h.flags & FLAG_TRA_NOLZ) != 0};

  std::vector<uint8_t> img_stream =
      img_no_lz ? std::vector<uint8_t>(view->img_c.begin(), view->img_c.end())
                : lzss_decompress(view->img_c);
  std::vector<uint8_t> tra_stream =
      has_alpha
          ? (tra_no_lz
                 ? std::vector<uint8_t>(view->tra_c.begin(), view->tra_c.end())
                 : lzss_decompress(view->tra_c))
          : std::vector<uint8_t>{};

  size_t total_bits{(size_t)h.w * (size_t)h.h};
  auto img_bits{img_was_raw ? decode_raw(img_stream, total_bits)
                            : rle_decode(img_stream, total_bits)};
  std::optional<std::vector<bool>> tra_bits;
  if (has_alpha)
    tra_bits = tra_was_raw ? decode_raw(tra_stream, total_bits)
                           : rle_decode(tra_stream, total_bits);

  bm.width = h.w;
  bm.height = h.h;
  bm.image_data = std::move(img_bits);
  bm.transparency_data = std::move(tra_bits);
  return bm;
}

} // namespace bfl

#endif // BFL_IMPLEMENTATION