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

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

template <class T> inline void put_le(std::vector<uint8_t> &buf, T v) {
  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) {
  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;
}

inline std::vector<uint8_t> read_file(std::filesystem::path const &p) {
  std::ifstream f(p, std::ios::binary);
  return std::vector<uint8_t>((std::istreambuf_iterator<char>(f)),
                              std::istreambuf_iterator<char>());
}
inline void write_file(std::filesystem::path const &p,
                       std::span<uint8_t const> d) {
  std::ofstream f(p, std::ios::binary);
  f.write(reinterpret_cast<char const *>(d.data()),
          static_cast<std::streamsize>(d.size()));
}

static 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;
}
static 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;
}

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

static std::vector<uint8_t> bits_to_bytes(std::vector<bool> const &data) {
  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;
}

static 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;
}

static std::vector<uint8_t>
rle_encode_with_transparency(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;
}

static 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;
}

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

struct BFLView {
  BFLHeader hdr{};
  std::span<uint8_t const> img_c{};
  std::span<uint8_t const> tra_c{};
};

static std::optional<BFLView> parse_bfl(std::span<uint8_t const> data) {
  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};
  BFLView 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;
}

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

  std::vector<uint32_t> to_rgba() const {
    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;
  }

  std::vector<uint8_t> encode() const {
    auto choose_rle_or_raw{[&](std::vector<bool> const &bits, bool &raw_flag,
                               bool allow_trick) {
      std::vector<bool> const *mask{transparency_data ? &(*transparency_data)
                                                      : nullptr};
      auto a{rle_encode_with_transparency(bits, false, mask, allow_trick)};
      auto b{rle_encode_with_transparency(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;
    }};

    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) {
      auto c{lzss_compress(v)};
      if (c.size() < v.size())
        return c;
      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;
  }

  static Bitmap from_rgba(std::span<uint32_t const> data, int w, int h) {
    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;
  }

  static Bitmap decode(std::span<uint8_t const> data) {
    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;
  }
};

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);
}

int main(int argc, char const *argv[]) {
  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") {
    auto bytes{read_file(argv[2])};
    auto v{parse_bfl(bytes)};
    if (!v) {
      std::println(stderr, "Invalid or truncated BFL: {}", argv[2]);
      return 2;
    }
    print_bfl_header(v->hdr);
    return 0;
  }

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

  std::filesystem::path in{argv[1]}, outp{argv[2]};
  Bitmap bmp;

  if (in.extension() == ".bfl") {
    bmp = Bitmap::decode(read_file(in));
  } else {
    int w, h, comp;
    unsigned char *raw{stbi_load(in.string().c_str(), &w, &h, &comp, 4)};
    if (!raw) {
      std::println(stderr, "stbi_load failed on {}", in.string());
      return 1;
    }
    std::span<uint32_t const> px{reinterpret_cast<uint32_t const *>(raw),
                                 (size_t)w * (size_t)h};
    bmp = Bitmap::from_rgba(px, w, h);
    stbi_image_free(raw);
  }

  if (outp.extension() == ".bfl") {
    auto bin{bmp.encode()};
    write_file(outp, std::span<uint8_t const>(bin.data(), bin.size()));
  } else {
    auto rgba{bmp.to_rgba()};
    stbi_write_png(outp.string().c_str(), bmp.width, bmp.height, 4, rgba.data(),
                   bmp.width * 4);
  }
  return 0;
}