// Inspired a lot by https://github.com/jeffman/MOTHER-3-Funland

#include <fstream>
#include <string_view>
#include <stdexcept>
#include <string>
#include <vector>
#include <iostream>
#include <Magick++.h>
#include <sstream>
#include "common.h"

class PaletteSet {
public:
  PaletteSet(BufferView m3) : m3_(m3) {
    for (int i=0; i<16; i++) {
      defaultPals_.emplace_back(m3, GetPointer(0) + (i << 5));
    }
  }

  const Palette& GetPalette(int spriteindex, int palnum = -1) {
    int a = GetPointer(spriteindex);

    if ((spriteindex > 0xFF) && (spriteindex < 0x26C)) {
      specialPals_[spriteindex] = Palette(m3_, a);
      return specialPals_[spriteindex];
    } else {
      if (palnum == -1) {
        palnum = static_cast<unsigned char>(m3_.ReadByte(0x1433D7C + 1 + ((spriteindex + 1) * 12))) & 0xF;
      }

      return defaultPals_.at(palnum);
    }
  }

private:

  int GetPointer(int spriteindex) {
    const int ADDR = 0x1A41548;

    int index = 0;

    if ((spriteindex > 0xFF) && (spriteindex < 0x26C)) {
      // These sprites use the (spriteindex - 0xFF)th entry
      index = spriteindex - 0xff;
    }

    int a = m3_.ReadFourBytes(ADDR + 4 + (index << 2));
    if (a == -1) return -1;
    return ADDR + a;
  }

  BufferView m3_;
  std::vector<Palette> defaultPals_;
  std::map<int, Palette> specialPals_;
};

struct Subsprite {
  int x;
  int y;
  int tile;
  bool flipH;
  bool flipV;
  int objSize;
  int objShape;

  int Width() const { return WIDTHS[objShape][objSize]; }

  int Height() const { return HEIGHTS[objShape][objSize]; }

private:

  static constexpr int WIDTHS[3][4] = { { 8, 16, 32, 64 }, { 16, 32, 32, 64 }, { 8, 8, 16, 32 } };
  static constexpr int HEIGHTS[3][4] = { { 8, 16, 32, 64 }, { 8, 8, 16, 32 }, { 16, 32, 32, 64 } };
};

struct FrameOutput {
  Magick::Image image;
  int width;
  int height;
  int centerX;
  int centerY;
};

struct Sprite {
  std::vector<Subsprite> subsprites;

  FrameOutput render(BufferView m3, const Palette& palette, const int gfxPtr) const {
    FrameOutput output;

    if (subsprites.empty()) return output;

    int minX = subsprites[0].x;
    int minY = subsprites[0].y;
    int maxX = minX + subsprites[0].Width();
    int maxY = minY + subsprites[0].Height();

    for (int j=1; j<subsprites.size(); j++) {
      const Subsprite& o = subsprites[j];

      if (o.x < minX) minX = o.x;
      if (o.y < minY) minY = o.y;
      if ((o.x + o.Width()) > maxX) maxX = o.x + o.Width();
      if ((o.y + o.Height()) > maxY) maxY = o.y + o.Height();
    }

    int width = maxX - minX;
    int height = maxY - minY;
    int centerX = -minX;
    int centerY = -minY;

    output.width = width;
    output.height = height;
    output.centerX = centerX;
    output.centerY = centerY;
    output.image = Magick::Image(Magick::Geometry(width, height), "transparent");
    output.image.modifyImage();

    Magick::Pixels view(output.image);

    for (const Subsprite& o : subsprites) {
      int tilePointer = o.tile << 5;
      for (int y=0; y<o.Height(); y += 8) {
        int actualY = o.flipV ? (o.Height() - y - 8) : y;

        for (int x=0; x<o.Width(); x += 8) {
          int tileData[8][8];

          int tileSrcOffset = gfxPtr + tilePointer;
          for (int ty=0; ty<8; ty++) {
            for (int tx=0; tx<4; tx++) {
              unsigned char ch = m3.ReadByte(tileSrcOffset++);
              tileData[tx*2][ty] = static_cast<unsigned char>(ch & 0xF);
              tileData[tx*2+1][ty] = static_cast<unsigned char>((ch >> 4) & 0xF);
            }
          }

          tilePointer += 0x20;

          int actualX = o.flipH ? (o.Width() - x - 8) : x;

          int destX = o.x + centerX + actualX;
          int destY = o.y + centerY + actualY;

          Magick::PixelPacket* pixels = view.get(destX,destY,8,8);

          for (int ty=0; ty<8; ty++) {
            int actualTy = o.flipV ? (7-ty) : ty;

            for (int tx=0; tx<8; tx++) {
              int actualTx = o.flipH ? (7-tx) : tx;

              if (tileData[actualTx][actualTy] != 0) {
                //auto& c = palette.Colors().at(tileData[actualTx][actualTy]);
                //std::cout << c.redQuantum() << "," << c.greenQuantum() << "," << c.blueQuantum() << std::endl;
                *pixels = palette.Colors().at(tileData[actualTx][actualTy]);
                //std::cout << tileData[actualTx][actualTy] << std::endl;
              }
              pixels++;
            }
          }

          view.sync();
        }
      }
    }

    return output;
  }
};

struct SpriteSheet {
  int spriteindex = 0;
  int gfxPtr;
  std::vector<Sprite> sprites;

  void render(BufferView m3, PaletteSet& palettes, std::vector<FrameOutput>& frames) const {
    const Palette& palette = palettes.GetPalette(spriteindex);
    for (int i=0; i<sprites.size(); i++) {
      FrameOutput f = sprites[i].render(m3, palette, gfxPtr);
      frames.push_back(std::move(f));
    }
  }
};

class Bank {
public:

  Bank(BufferView m3, const int baseAddr, const int gfxAddr) : baseAddr_(baseAddr), gfxAddr_(gfxAddr) {
    int numEntries = m3.ReadFourBytes(baseAddr);
    std::cout << numEntries << std::endl;

    for (int i = 0; i < numEntries; i++) {
      int sheetAddr = GetPointerToSheet(m3, i);
      if (sheetAddr == -1) continue;

      SpriteSheet ss;
      ss.spriteindex = i;
      ss.gfxPtr = GetGfxPointer(m3, i);

      m3.Seek(sheetAddr);
      m3.SeekAdd(8);

      unsigned short spriteCount = m3.ReadNextTwoBytes();
      //std::cout << i << ":" << spriteCount << " at " << std::hex << sheetAddr << std::endl;
      for (int j=0; j<spriteCount; j++) {
        Sprite sprite;

        unsigned short subspriteCount = m3.ReadNextTwoBytes();
        for (int k=0; k<subspriteCount; k++) {
          Subsprite subsprite;
          subsprite.y = static_cast<char>(m3.ReadNextByte());
          subsprite.x = static_cast<char>(m3.ReadNextByte());

          unsigned short ch = m3.ReadNextTwoBytes();
          subsprite.tile = static_cast<unsigned short>(ch & 0x3FF);
          subsprite.flipH = (ch & 0x400) != 0;
          subsprite.flipV = (ch & 0x800) != 0;
          subsprite.objSize = static_cast<unsigned char>((ch >> 12) & 3);
          subsprite.objShape = static_cast<unsigned char>((ch >> 14) & 3);

          /*std::cout << subsprite.y << std::endl;
          std::cout << subsprite.x << std::endl;
          std::cout << subsprite.flipH << "," << subsprite.flipV << std::endl;*/

          sprite.subsprites.push_back(std::move(subsprite));
        }

        m3.SeekAdd(2);

        ss.sprites.push_back(std::move(sprite));
      }

      spritesheets_[i] = std::move(ss);
      //if (i == 15) return;
    }
  }

  const std::map<int, SpriteSheet>& SpriteSheets() const { return spritesheets_; }

private:

  int GetPointerToSheet(BufferView m3, int index) {
    int readAt = baseAddr_ + 4 + (index << 2);
    int a = m3.ReadFourBytes(readAt);
    //std::cout << readAt << " :: " << a << std::hex << std::endl;
    if (a == 0) return -1;
    return a + baseAddr_;
  }

  int GetGfxPointer(BufferView m3, int index) {
    return m3.ReadFourBytes(gfxAddr_ + 4 + (index << 2)) + gfxAddr_;
  }

  int baseAddr_;
  int gfxAddr_;
  std::map<int, SpriteSheet> spritesheets_;
};

int main(int argc, char** argv) {
  if (argc < 3) {
    std::cout << "Usage: ./sprite_dumper [path to rom] {sheet IDs}" << std::endl;
    std::cout << "sheet IDs should be space separated references to the sheets to concatenate" << std::endl;
    std::cout << "the format of the ID is BankNum.SheetNum" << std::endl;
    return -1;
  }

  Magick::InitializeMagick(nullptr);

  Rom m3(argv[1]);
  PaletteSet palettes(m3.buffer());
  Bank b1(m3.buffer(), 0x1A442A4, 0x14383E4);
  Bank b2(m3.buffer(), 0x1AE0638, 0x194BC30);
  Bank b3(m3.buffer(), 0x1AEE4C4, 0x1A012B8);
  Bank b4(m3.buffer(), 0x1AF1ED0, 0x1A36AA0);

  std::vector<FrameOutput> frames;

  for (int i=2; i<argc; i++) {
    std::stringstream argfmt;
    argfmt << argv[i];

    int bankNum;
    char ch;
    int sheetNum;
    argfmt >> bankNum;
    argfmt >> ch; //.
    argfmt >> sheetNum;

    const Bank& bankToRead = [&](){
      switch (bankNum) {
        case 0: return b1;
        case 1: return b2;
        case 2: return b3;
        case 3: return b4;
        default: throw std::invalid_argument("Invalid bank num: " + std::to_string(bankNum));
      }
    }();

    bankToRead.SpriteSheets().at(sheetNum).render(m3.buffer(), palettes, frames);
  }

  int maxWidth = 0;
  int maxHeight = 0;
  for (const FrameOutput& f : frames) {
    if (f.width > maxWidth) maxWidth = f.width;
    if (f.height > maxHeight) maxHeight = f.height;
  }

  const int FRAMES_PER_ROW = 10;
  int sheetWidth;
  int sheetHeight;

  std::ofstream datafile("out.txt");
  datafile << maxWidth << "," << maxHeight << " cell size" << std::endl;
  datafile << FRAMES_PER_ROW << " frames per row" << std::endl;
  datafile << frames.size() << " frames" << std::endl;
  datafile << std::endl;

  if (frames.size() < FRAMES_PER_ROW) {
    sheetWidth = frames.size() * maxWidth;
    sheetHeight = maxHeight;
  } else {
    sheetWidth = FRAMES_PER_ROW * maxWidth;
    sheetHeight = (frames.size() / FRAMES_PER_ROW + 1) * maxHeight;
  }

  Magick::Image sheet(Magick::Geometry(sheetWidth, sheetHeight), "transparent");
  for (int i=0; i<frames.size(); i++) {
    const FrameOutput& f = frames.at(i);
    sheet.composite(f.image, (i%FRAMES_PER_ROW)*maxWidth, (i/FRAMES_PER_ROW)*maxHeight, Magick::OverCompositeOp);
    datafile << f.width << "," << f.height << "," << f.centerX << "," << f.centerY << std::endl;
  }

  //Magick::Image im = b2.SpriteSheets().at(3).render(m3, palettes);
  sheet.magick("png");
  sheet.write("out.png");

  return 0;
}