#include <cassert>
#include <filesystem>
#include <fstream>
#include <optional>
#include <print>
#include <span>
#include <vector>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

// Helpers

inline void store_le16(std::vector<uint8_t> &buf, uint16_t val) {
  buf.push_back(static_cast<uint8_t>(val & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
}

inline void store_le32(std::vector<uint8_t> &buf, uint32_t val) {
  buf.push_back(static_cast<uint8_t>(val & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 16) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 24) & 0xFF));
}

inline void store_le64(std::vector<uint8_t> &buf, uint64_t val) {
  buf.push_back(static_cast<uint8_t>(val & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 16) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 24) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 32) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 40) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 48) & 0xFF));
  buf.push_back(static_cast<uint8_t>((val >> 56) & 0xFF));
}

static inline uint16_t read_le16(std::span<const uint8_t> s, size_t off) {
  return static_cast<uint16_t>(s[off] | (s[off + 1] << 8));
}

static inline uint32_t read_le32(std::span<const uint8_t> s, size_t off) {
  return static_cast<uint32_t>(s[off] | (s[off + 1] << 8) | (s[off + 2] << 16) |
                               (s[off + 3] << 24));
}

static std::vector<uint8_t> lzss_compress(std::span<const uint8_t> in) {
  const int W = 4096;
  const int LA = 18;
  const int MIN = 3;

  std::vector<uint8_t> out;
  size_t i = 0;

  while (i < in.size()) {
    uint8_t flag = 0;
    size_t flag_pos = out.size();
    out.push_back(0);

    for (int bit = 0; bit < 8 && i < in.size(); ++bit) {
      size_t best_len = 0;
      size_t best_off = 0;

      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 = 0;
        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]);
        ++i;
      }
    }

    out[flag_pos] = flag;
  }

  return out;
}

// "Library"

constexpr uint8_t FLAG_HAS_ALPHA = 0x01;
constexpr uint8_t FLAG_IMG_RAW = 0x02;  // pre-LZSS image was RAW bitpack
constexpr uint8_t FLAG_TRA_RAW = 0x04;  // pre-LZSS mask was RAW bitpack
constexpr uint8_t FLAG_IMG_NOLZ = 0x08; // image stream stored without LZSS
constexpr uint8_t FLAG_TRA_NOLZ = 0x10; // mask stream stored without LZSS

static std::vector<uint8_t>
rle_encode_with_transparency(const std::vector<bool> &data, bool initial,
                             const std::vector<bool> *transparency,
                             bool allow_trick) {
  std::vector<uint8_t> out;
  out.push_back(static_cast<uint8_t>(initial));

  bool state = initial;
  size_t count = 0;
  size_t i = 0;

  auto flush_emit = [&](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_emit(255);
          state = !state;
          count = 0;
        }
        ++count;
        ++i;
      }
      continue;
    }

    bool bit = data[i];
    if (bit == state) {
      ++count;
      ++i;

      if (count == 255 && i < data.size()) {
        flush_emit(255);
        out.push_back(0);
        count = 0;
      }
    } else {
      flush_emit(static_cast<uint8_t>(count));
      state = !state;
      count = 1;
      ++i;
    }
  }

  if (count > 0)
    flush_emit(static_cast<uint8_t>(count));
  return out;
}

static std::vector<uint8_t> lzss_decompress(std::span<const uint8_t> in) {
  std::vector<uint8_t> out;
  size_t i = 0;

  while (i < in.size()) {
    uint8_t flag = in[i++];
    for (int bit = 0; bit < 8 && i < in.size(); ++bit) {
      bool is_lit = (flag >> bit) & 1;
      if (is_lit) {
        out.push_back(in[i++]);
      } else {
        if (i + 1 >= in.size())
          return out;
        uint8_t b0 = in[i++];
        uint8_t b1 = in[i++];
        size_t len = (b0 >> 4) + 3;
        size_t off = ((b0 & 0x0F) << 8) | b1; // 1..4095

        if (off == 0 || off > out.size())
          return out; // invalid
        size_t src = out.size() - off;
        for (size_t k = 0; k < len; ++k) {
          out.push_back(out[src + k]);
        }
      }
    }
  }
  return out;
}

static std::vector<bool> rle_decode(std::span<const uint8_t> 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;
  size_t 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 = 0; k < emit; ++k)
      out[produced + k] = state;
    produced += emit;

    if (j < in.size() && in[j] == 0) {
      ++j;
    } else {
      state = !state;
    }
  }

  return out;
}

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

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

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

    std::vector<uint8_t> ret;

    // header
    ret.push_back('B');
    ret.push_back('F');
    ret.push_back('L');
    store_le16(ret, width);
    store_le16(ret, height);

    auto encode_best = [&](std::vector<bool> const &data, bool &raw_flag,
                           bool allow_trick) {
      const std::vector<bool> *mask =
          transparency_data ? &(*transparency_data) : nullptr;

      auto init_false =
          rle_encode_with_transparency(data, false, mask, allow_trick);
      auto init_true =
          rle_encode_with_transparency(data, true, mask, allow_trick);

      std::vector<uint8_t> result = (init_false.size() <= init_true.size())
                                        ? std::move(init_false)
                                        : std::move(init_true);

      if (result.size() * 8 > data.size()) {
        result = raw_encode(data);
        raw_flag = true;
      }
      return result;
    };

    std::vector<uint8_t> image_bytes, trans_bytes;
    bool image_raw = false, trans_raw = false;

    image_bytes =
        encode_best(image_data, image_raw, transparency_data.has_value());
    if (transparency_data) {
      trans_bytes = encode_best(*transparency_data, trans_raw, false);
    }

    bool img_no_lz = false, tra_no_lz = false;

    std::vector<uint8_t> img_stream_lz = lzss_compress(image_bytes);
    std::vector<uint8_t> img_stream =
        (img_stream_lz.size() < image_bytes.size())
            ? std::move(img_stream_lz)
            : (img_no_lz = true, image_bytes);

    std::vector<uint8_t> tra_stream;
    if (!trans_bytes.empty()) {
      auto tra_stream_lz = lzss_compress(trans_bytes);
      tra_stream = (tra_stream_lz.size() < trans_bytes.size())
                       ? std::move(tra_stream_lz)
                       : (tra_no_lz = true, trans_bytes);
    }

    // flags
    uint8_t flags = 0;
    if (transparency_data)
      flags |= FLAG_HAS_ALPHA;
    if (image_raw)
      flags |= FLAG_IMG_RAW;
    if (trans_raw)
      flags |= FLAG_TRA_RAW;
    if (img_no_lz)
      flags |= FLAG_IMG_NOLZ;
    if (tra_no_lz)
      flags |= FLAG_TRA_NOLZ;
    ret.push_back(flags);

    store_le32(ret, static_cast<uint32_t>(img_stream.size()));
    store_le32(ret, static_cast<uint32_t>(tra_stream.size()));

    // payload
    ret.insert(ret.end(), img_stream.begin(), img_stream.end());
    ret.insert(ret.end(), tra_stream.begin(), tra_stream.end());

    return ret;
  }

  static auto from_rgba(std::span<uint32_t const> const &data, int width,
                        int height) -> Bitmap;
  static auto decode(std::span<uint8_t const> const &data) -> Bitmap;
};

auto Bitmap::decode(std::span<uint8_t const> const &data) -> Bitmap {
  Bitmap bm{};

  if (data.size() < 3 + 2 + 2 + 1 + 4 + 4)
    return bm;
  if (!(data[0] == 'B' && data[1] == 'F' && data[2] == 'L'))
    return bm;

  size_t off = 3;
  uint16_t w = read_le16(data, off);
  off += 2;
  uint16_t h = read_le16(data, off);
  off += 2;

  uint8_t flags = data[off++];

  uint32_t img_c_len = read_le32(data, off);
  off += 4;
  uint32_t tra_c_len = read_le32(data, off);
  off += 4;

  if (off + img_c_len > data.size())
    return bm;
  if (off + img_c_len + tra_c_len > data.size())
    return bm;

  auto img_c = data.subspan(off, img_c_len);
  off += img_c_len;
  auto tra_c = data.subspan(off, tra_c_len);

  bool has_alpha = (flags & FLAG_HAS_ALPHA) != 0;
  bool img_was_raw = (flags & FLAG_IMG_RAW) != 0;
  bool tra_was_raw = (flags & FLAG_TRA_RAW) != 0;
  bool img_no_lz = (flags & FLAG_IMG_NOLZ) != 0;
  bool tra_no_lz = (flags & FLAG_TRA_NOLZ) != 0;

  std::vector<uint8_t> img_stream =
      img_no_lz ? std::vector<uint8_t>(img_c.begin(), img_c.end())
                : lzss_decompress(img_c);

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

  size_t total_bits = static_cast<size_t>(w) * static_cast<size_t>(h);

  std::vector<bool> img_bits = img_was_raw ? raw_decode(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 ? raw_decode(tra_stream, total_bits)
                            : rle_decode(tra_stream, total_bits));
  }

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

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

  Bitmap bm;
  bm.width = w;
  bm.height = h;
  bm.image_data.resize(static_cast<size_t>(w) * static_cast<size_t>(h));
  bm.transparency_data = std::vector<bool>{};
  bm.transparency_data->resize(static_cast<size_t>(w) * static_cast<size_t>(h));

  for (size_t i = 0; i < data.size(); ++i) {
    uint32_t px = data[i];

    // little-endian: 0xAABBGGRR
    uint8_t r = static_cast<uint8_t>((px >> 0) & 0xFF);
    uint8_t g = static_cast<uint8_t>((px >> 8) & 0xFF);
    uint8_t b = static_cast<uint8_t>((px >> 16) & 0xFF);
    uint8_t a = static_cast<uint8_t>((px >> 24) & 0xFF);

    bool is_white = (r == 255 && g == 255 && b == 255 && a != 0);
    bool is_black = (r == 0 && g == 0 && b == 0 && a != 0);

    bool is_transparent = (!is_white && !is_black) || (a == 0);

    (*bm.transparency_data)[i] = is_transparent;
    bm.image_data[i] = is_white;
  }

  return bm;
}

auto Bitmap::to_rgba() -> std::vector<uint32_t> {
  std::vector<uint32_t> out;
  out.resize(static_cast<size_t>(width) * static_cast<size_t>(height));

  auto make_rgba = [](uint8_t r, uint8_t g, uint8_t b, uint8_t a) -> uint32_t {
    return (static_cast<uint32_t>(r) << 0) | (static_cast<uint32_t>(g) << 8) |
           (static_cast<uint32_t>(b) << 16) | (static_cast<uint32_t>(a) << 24);
  };

  for (size_t i = 0; i < out.size(); ++i) {
    if ((transparency_data && i >= transparency_data->size()) ||
        i >= image_data.size()) {
      out[i] = make_rgba(0, 0, 0, 0); // safety
      continue;
    }

    if (transparency_data && (*transparency_data)[i]) {
      out[i] = make_rgba(0, 0, 0, 0);
    } else {
      if (image_data[i]) {
        out[i] = make_rgba(255, 255, 255, 255);
      } else {
        out[i] = make_rgba(0, 0, 0, 255);
      }
    }
  }

  return out;
}

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

static bool parse_bfl_header(std::span<const uint8_t> data, BFLHeader &hdr) {
  // minimal header: "BFL" + w + h + flags + lenI + lenT
  if (data.size() < 3 + 2 + 2 + 1 + 4 + 4)
    return false;
  if (!(data[0] == 'B' && data[1] == 'F' && data[2] == 'L'))
    return false;

  size_t off = 3;
  hdr.w = read_le16(data, off);
  off += 2;
  hdr.h = read_le16(data, off);
  off += 2;
  hdr.flags = data[off++];
  hdr.img_len = read_le32(data, off);
  off += 4;
  hdr.tra_len = read_le32(data, off);
  off += 4;

  if (off + hdr.img_len > data.size())
    return false;
  if (off + hdr.img_len + hdr.tra_len > data.size())
    return false;
  return true;
}

static void print_bfl_header(BFLHeader const &h) {
  std::println("BFL header:");
  std::println("  width       : {}", h.w);
  std::println("  height      : {}", h.h);
  std::println("  flags (hex) : 0x{:02X}", h.flags);
  std::println("    has_alpha : {}", (h.flags & FLAG_HAS_ALPHA) ? "yes" : "no");
  std::println("    img_raw   : {}", (h.flags & FLAG_IMG_RAW) ? "yes" : "no");
  std::println("    tra_raw   : {}", (h.flags & FLAG_TRA_RAW) ? "yes" : "no");
  std::println("    img_no_lz : {}", (h.flags & FLAG_IMG_NOLZ) ? "yes" : "no");
  std::println("    tra_no_lz : {}", (h.flags & FLAG_TRA_NOLZ) ? "yes" : "no");
  std::println("  image bytes : {}", h.img_len);
  std::println("  alpha bytes : {}", h.tra_len);
}

// CLI

auto read_entire_file(std::filesystem::path const &path)
    -> std::vector<uint8_t> {
  std::ifstream stream(path, std::ios::in | std::ios::binary);
  std::vector<uint8_t> const contents((std::istreambuf_iterator<char>(stream)),
                                      std::istreambuf_iterator<char>());
  return contents;
}

auto write_entire_file(std::filesystem::path const &path,
                       std::span<uint8_t const> const &data) -> void {
  std::ofstream stream(path, std::ios::out | std::ios::binary);
  stream.write(reinterpret_cast<char const *>(data.data()), data.size());
}

auto main(int const argc, char const *argv[]) -> int {
  if (argc < 2) {
    std::println(stderr,
                 "Usage:\n"
                 "  {} <input> <output>\n"
                 "  {} header <file.bfl>",
                 argv[0], argv[0]);
    return 1;
  }

  if (argc == 3 && std::string_view(argv[1]) == "header") {
    std::filesystem::path filepath{argv[2]};
    auto bytes = read_entire_file(filepath);
    BFLHeader hdr{};
    if (!parse_bfl_header(bytes, hdr)) {
      std::println(stderr, "Invalid or truncated BFL: {}", filepath.string());
      return 2;
    }
    print_bfl_header(hdr);
    return 0;
  }

  if (argc != 3) {
    std::println(stderr,
                 "Usage:\n"
                 "  {} <input> <output>\n"
                 "  {} header <file.bfl>",
                 argv[0], argv[0]);
    return 1;
  }

  Bitmap input;

  std::filesystem::path output_filepath{argv[2]};
  std::filesystem::path filepath{argv[1]};
  if (filepath.extension() == ".bfl") {
    input = Bitmap::decode(read_entire_file(filepath));
  } else {
    int width, height, comp;
    unsigned char *raw_data =
        stbi_load(filepath.string().c_str(), &width, &height, &comp, 4);
    if (!raw_data) {
      std::println(stderr, "stbi_load failed on {}", filepath.string());
      return 1;
    }

    std::span<uint32_t const> data{reinterpret_cast<uint32_t const *>(raw_data),
                                   static_cast<size_t>(width) *
                                       static_cast<size_t>(height)};
    input = Bitmap::from_rgba(data, width, height);

    stbi_image_free(raw_data);
  }

  if (output_filepath.extension() == ".bfl") {
    auto const input_bin = input.encode();
    std::span<uint8_t const> d{input_bin.data(), input_bin.size()};
    write_entire_file(output_filepath, d);
  } else {
    auto rgba = input.to_rgba();
    stbi_write_png(output_filepath.string().c_str(), input.width, input.height,
                   4, rgba.data(), input.width * 4);
  }
  return 0;
}