#include <iostream>
#include <Magick++.h>
#include <vector>
#include <string>
#include <map>
#include <fstream>
#include "common.h"
#include "identifier.h"

constexpr int NUM_ROOMS = 1000;
constexpr int NUM_TILESETS = 12;

struct RoomInfo {
  int width;
  int height;

  static std::map<int, RoomInfo> ReadFromRom(BufferView m3) {
    const int BASE_ADDR = 0xD2E1D8 + 12;
    const int ENTRY_LEN = 28;

    std::map<int, RoomInfo> output;

    for (int i=0; i<NUM_ROOMS; i++) {
      RoomInfo& ri = output[i];

      ri.width = ((m3.ReadByte(BASE_ADDR + i*ENTRY_LEN + 20) & 7) + 1) << 4;
      ri.height = (((m3.ReadByte(BASE_ADDR + i*ENTRY_LEN + 24) & 0x3F) >> 3) + 1) << 4;
    }

    return output;
  }
};

struct RoomGfxPal {
  int paletteId;
  int tilesetId[NUM_TILESETS];

  static std::map<int, RoomGfxPal> ReadFromRom(BufferView m3) {
    const int BASE_ADDR = 0xD34F44 + 12;
    const int ENTRY_LEN = 26;

    std::map<int, RoomGfxPal> output;

    for (int i=0; i<NUM_ROOMS; i++) {
      RoomGfxPal& rgp = output[i];

      rgp.paletteId = static_cast<short>(m3.ReadTwoBytes(BASE_ADDR + i*ENTRY_LEN + 24));
      for (int tid=0; tid<NUM_TILESETS; tid++) {
        rgp.tilesetId[tid] = m3.ReadTwoBytes(BASE_ADDR + i*ENTRY_LEN + tid*2);
      }
    }

    return output;
  }

  std::vector<std::vector<char>> GetTilesets(BufferView m3) const {
    const int BASE_ADDR = 0xD3B4E0;

    std::vector<std::vector<char>> output(NUM_TILESETS);

    for (int i=0; i<NUM_TILESETS; i++) {
      if (tilesetId[i] < (NUM_ROOMS * 3)) {
        unsigned long gfxAddr = m3.ReadFourBytes(BASE_ADDR + 4 + (tilesetId[i] << 2)) + BASE_ADDR;
        output[i] = m3.Decompress(gfxAddr);
      }
    }

    return output;
  }

  std::vector<Palette> GetPalettes(BufferView m3) const {
    const int BASE_ADDR = 0xF3C344;

    std::vector<Palette> output;
    int offset = m3.ReadFourBytes(BASE_ADDR + 4 + (paletteId << 2));

    for (int i=0; i<16; i++) {
      unsigned long palAddr = BASE_ADDR + offset + (i << 5);
      output.emplace_back(m3, palAddr);
    }

    return output;
  }
};

std::vector<char> GetMapTiles(BufferView m3, int roomId) {
  const int BASE_ADDR = 0x104D9CC;
  int tilesAddr = m3.ReadFourBytes(BASE_ADDR + 4 + (roomId << 2)) + BASE_ADDR;
  return m3.Decompress(tilesAddr);
}

std::vector<std::vector<char>> GetMapLayers(BufferView m3, int roomId) {
  const int BASE_ADDR = 0xF9003C;
  std::vector<std::vector<char>> layers;
  for (int i=0; i<3; i++) {
    try {
      int lookAddr = BASE_ADDR + 4 + (((roomId * 3) + i) << 2);
      int layerAddr = m3.ReadFourBytes(lookAddr) + BASE_ADDR;
      layers.push_back(m3.Decompress(layerAddr));
    } catch (const std::domain_error&) {
      // ignore
    }
  }
  return layers;
}

struct TileUse {
  size_t id;
  bool tflipx = false;
  bool tflipy = false;
};

unsigned short stripFlipInfo(unsigned short ch) {
  return ch & ~(0x4000 | 0x8000);
}

class Map {
public:

  Map(
    BufferView m3,
    int roomNum,
    RoomInfo roomInfo,
    std::vector<Palette> palettes,
    std::vector<std::vector<char>> tilesets)
      : roomNum_(roomNum),
        width_(roomInfo.width),
        height_(roomInfo.height),
        palettes_(palettes),
        tilesets_(tilesets),
        mapTiles_(GetMapTiles(m3, roomNum)),
        mapLayers_(GetMapLayers(m3, roomNum))
  {
    for (int layer=0; layer<mapLayers_.size(); layer++) {
      const std::vector<char>& ml = mapLayers_[layer];
      if (ml.empty()) continue;

      std::vector<TileUse> newLayer;

      for (int mapy = 0; mapy < height_; mapy++) {
        for (int mapx = 0; mapx < width_; mapx++) {
          unsigned short ch = BufferView(ml).ReadTwoBytes((mapx + (mapy * width_)) * 2);
          TileUse tu;
          tu.id = metatiles_.add(stripFlipInfo(ch));
          tu.tflipx = (ch & 0x4000) != 0;
          tu.tflipy = (ch & 0x8000) != 0;
          newLayer.push_back(std::move(tu));
        }
      }

      itemised_.push_back(std::move(newLayer));
    }
  }

  Magick::Image renderTile(size_t metatile_id, bool tflipx, bool tflipy) const {
    Magick::Image result("16x16", "transparent");

    unsigned short ch = metatiles_.get(metatile_id);
    unsigned short tile16 = ch & 0x3FF;
    if ((tile16 >> 6) >= 12) return result;

    result.modifyImage();
    Magick::Pixels view(result);

    int tpal = (ch >> 10) & 0xF;
    const Palette& palette = palettes_[tpal];

    int tile8[2][2];
    bool sflipx[2][2];
    bool sflipy[2][2];

    unsigned int magic = BufferView(mapTiles_).ReadFourBytes(tile16 * 8);

    tile8[0][0] = mapTiles_[(tile16 * 8) + 4];
    tile8[0][1] = mapTiles_[(tile16 * 8) + 5];
    tile8[1][0] = mapTiles_[(tile16 * 8) + 6];
    tile8[1][1] = mapTiles_[(tile16 * 8) + 7];

    for (int i=0; i<2; i++) {
      for (int j=0; j<2; j++) {
        sflipx[i][j] = (tile8[i][j] & 0x40) != 0;
        sflipy[i][j] = (tile8[i][j] & 0x80) != 0;

        tile8[i][j] &= 0x3f;
        tile8[i][j] |= (ch & 0x3c0);
      }
    }

    unsigned int mask = (magic >> 16) & 0xf;
    if ((mask & 0x1) == 0) tile8[0][0] = -1;
    if ((mask & 0x2) == 0) tile8[0][1] = -1;
    if ((mask & 0x4) == 0) tile8[1][0] = -1;
    if ((mask & 0x8) == 0) tile8[1][1] = -1;

    for (int tiley=0; tiley<2; tiley++) {
      for (int tilex=0; tilex<2; tilex++) {
        if (tile8[tiley][tilex] < 0) continue;

        int tileset = tile8[tiley][tilex] >> 6;
        int subtile = tile8[tiley][tilex] & 0x3f;

        int tileData[8][8];
        BufferView tilesetData(tilesets_[tileset]);
        tilesetData.Seek(subtile << 5);
        for (int ty=0; ty<8; ty++) {
          for (int tx=0; tx<4; tx++) {
            unsigned char vvvv = tilesetData.ReadNextByte();
            tileData[tx*2][ty] = static_cast<unsigned char>(vvvv & 0xF);
            tileData[tx*2+1][ty] = static_cast<unsigned char>((vvvv >> 4) & 0xF);
          }
        }

        int stx = tflipx ? 1 - tilex : tilex;
        int sty = tflipy ? 1 - tiley : tiley;

        int destX = /*(mapx << 4) +*/ (stx << 3);
        int destY = /*(mapy << 4) +*/ (sty << 3);

        bool reallyFlipX = (sflipx[tiley][tilex] ^ tflipx);
        bool reallyFlipY = (sflipy[tiley][tilex] ^ tflipy);

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

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

          for (int tx=0; tx<8; tx++) {
            int actualTx = reallyFlipX ? (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 result;
  }

  const std::vector<std::vector<TileUse>>& getItemisedLayers() const {
    return itemised_;
  }

  int getWidth() const { return width_; }

  int getHeight() const { return height_; }

  int getRoomNum() const { return roomNum_; }

private:

  int roomNum_;
  int width_;
  int height_;
  std::vector<Palette> palettes_;
  std::vector<std::vector<char>> tilesets_;
  std::vector<char> mapTiles_;
  std::vector<std::vector<char>> mapLayers_;

  identifier<unsigned short> metatiles_;
  std::vector<std::vector<TileUse>> itemised_;
};

struct GlobalTile {
  Magick::Image image;
  std::string base64;

  GlobalTile(Magick::Image input) : image(input) {
    Magick::Blob blob;
    input.write(&blob);
    base64 = blob.base64();
  }
};

class GlobalTileKeyExtract {
public:

  const std::string& operator()(const GlobalTile& globaltile) const {
    return globaltile.base64;
  }
};

using globaltile_identifier = identifier<GlobalTile, GlobalTileKeyExtract>;
using globaltile_id = globaltile_identifier::id_type;

int main(int argc, char** argv) {
  if (argc < 3) {
    std::cout << "Usage: ./tileset_dumper [path to rom] [--palette=path to override palette in ACT format] {map ID}" << std::endl;
    return -1;
  }

  Magick::InitializeMagick(nullptr);

  int mapArgStart = 2;

  bool shouldOverridePalette = false;
  std::vector<Palette> overridePalettes;
  if (std::string(argv[2]).substr(0, 10) == "--palette=") {
    mapArgStart++;
    shouldOverridePalette = true;

    std::string filename = std::string(argv[2]).substr(10);
    std::ifstream palfile(filename, std::ios::binary);
    if (!palfile.is_open()) {
      throw std::invalid_argument("Could not find palette file: " + filename);
    }

    palfile.seekg(0, palfile.end);
    int length = palfile.tellg();
    palfile.seekg(0, palfile.beg);
    std::vector<char> paletteData(length, 0);
    palfile.read(paletteData.data(), length);

    BufferView palbuf(paletteData);
    for (int i=0; i<12; i++) {
      std::vector<Magick::Color> colors;

      for (int j=0; j<16; j++) {
        unsigned char r = palbuf.ReadNextByte();
        unsigned char g = palbuf.ReadNextByte();
        unsigned char b = palbuf.ReadNextByte();
        colors.push_back(Magick::ColorRGB(r/256.0, g/256.0, b/256.0));
      }

      overridePalettes.emplace_back(std::move(colors));
    }
  }

  Rom m3(argv[1]);
  auto roomInfos = RoomInfo::ReadFromRom(m3.buffer());
  auto roomGfxPals = RoomGfxPal::ReadFromRom(m3.buffer());

  std::list<Map> maps;
  for (int i=mapArgStart; i<argc; i++) {
    int roomNum = std::stoi(argv[i]);
    RoomGfxPal& roomGfxPal = roomGfxPals[roomNum];

    maps.emplace_back(
      m3.buffer(),
      roomNum,
      roomInfos[roomNum],
      shouldOverridePalette ? overridePalettes : roomGfxPal.GetPalettes(m3.buffer()),
      roomGfxPal.GetTilesets(m3.buffer()));
  }

  globaltile_identifier globaltiles;
  for (const Map& map : maps) {
    std::map<int, int> translatedTileIds;

    // Generate map datafile.
    std::ofstream mapfile("out" + std::to_string(map.getRoomNum()) + ".tmx");
    mapfile << R"(<map version="1.0" orientation="orthogonal" renderorder="right-down" width=")";
    mapfile << map.getWidth();
    mapfile << R"(" height=")";
    mapfile << map.getHeight();
    mapfile << R"(" tilewidth="16" tileheight="16">)" << std::endl;
    mapfile << R"(  <tileset firstgid="1" source="out.tsx" />)" << std::endl;

    for (int layer = map.getItemisedLayers().size()-1; layer >= 0; layer--) {
      mapfile << R"(  <layer id=")";
      mapfile << layer;
      mapfile << R"(" name="Layer )";
      mapfile << layer;
      mapfile << R"(" width=")";
      mapfile << map.getWidth();
      mapfile << R"(" height=")";
      mapfile << map.getHeight();
      mapfile << R"(">)" << std::endl;
      mapfile << R"(    <data encoding="csv">)";

      bool first = true;
      for (const TileUse& tu : map.getItemisedLayers()[layer]) {
        if (first) {
          first = false;
        } else {
          mapfile << ",";
        }

        if (!translatedTileIds.count(tu.id)) {
          Magick::Image renderedTile = map.renderTile(tu.id, false, false);
          renderedTile.magick("png");

          GlobalTile gt(std::move(renderedTile));
          translatedTileIds[tu.id] = globaltiles.add(gt);
        }

        unsigned int outChar = translatedTileIds[tu.id] + 1;
        if (tu.tflipx) outChar |= 0x80000000;
        if (tu.tflipy) outChar |= 0x40000000;
        mapfile << outChar;
      }

      mapfile << R"(</data>)" << std::endl;
      mapfile << R"(  </layer>)" << std::endl;
    }

    mapfile << R"(</map>)" << std::endl;

    // Render map to image.
    /*for (int layer=itemised.size()-1; layer>=0; layer--) {
      for (int mapy = 0; mapy < height; mapy++) {
        for (int mapx = 0; mapx < width; mapx++) {
          const TileUse& tu = itemised[layer][mapx+mapy*width];
          Magick::Image tileRender = renderTile(metatiles.get(tu.id), tu.tflipx, tu.tflipy, palettes, mapTiles, tilesets);
          image.composite(tileRender, mapx << 4, mapy << 4, Magick::OverCompositeOp);
        }
      }
    }*/
  }

  constexpr int TILES_PER_ROW = 10;
  int sheetWidth;
  int sheetHeight;

  if (globaltiles.size() < TILES_PER_ROW) {
    sheetWidth = globaltiles.size() * 16;
    sheetHeight = 16;
  } else {
    sheetWidth = TILES_PER_ROW * 16;
    sheetHeight = (globaltiles.size() / TILES_PER_ROW + 1) * 16;
  }

  std::ofstream tilesetfile("out.tsx");
  tilesetfile << R"(<tileset name="fromRom" tilewidth="16" tileheight="16" tilecount=")";
  tilesetfile << globaltiles.size();
  tilesetfile << R"(" columns=")";
  tilesetfile << TILES_PER_ROW;
  tilesetfile << R"(">)" << std::endl;
  tilesetfile << R"(  <image source="tiles.png" />)" << std::endl;
  tilesetfile << R"(</tileset>)" << std::endl;

  // Render tileset image.
  Magick::Image tilesetImage(Magick::Geometry(sheetWidth, sheetHeight), "transparent");
  for (int i=0; i<globaltiles.size(); i++) {
    const Magick::Image& tileRender = globaltiles.get(i).image;
    tilesetImage.composite(tileRender, (i % TILES_PER_ROW) << 4, (i / TILES_PER_ROW) << 4, Magick::OverCompositeOp);
  }

  tilesetImage.magick("png");
  tilesetImage.write("tiles.png");
}