#include <algorithm>
#include "stdafx.h"
#include "VirtualFile.h"
#include "HdPackBuilder.h"
#include "HdNesPack.h"
#include "Console.h"
#include "ConsolePauseHelper.h"
#include "VideoDecoder.h"

HdPackBuilder* HdPackBuilder::_instance = nullptr;

enum HDPackOuputTileType
{
	Both = 0,
	BG = 1,
	Sprite = 2
};

enum HdPackRecordFlags
{
	None = 0,
	UseLargeSprites = 1,
	SortByUsageFrequency = 2,
	GroupBlankTiles = 4,
	IgnoreOverscan = 8,
	SaveFrame = 16,
};

HdPackBuilder::HdPackBuilder(shared_ptr<Console> console, string saveFolder, ScaleFilterType filterType, uint32_t scale, uint32_t flags, uint32_t chrRamBankSize, uint32_t outTileType, bool isChrRam)
{
	_console = console;
	_saveFolder = saveFolder;
	_filterType = filterType;
	_chrRamBankSize = chrRamBankSize;
	_flags = flags;
	_isChrRam = isChrRam;
	_outTileType = outTileType;
	_hasNewTile = false;
	_frameID = 0;
	string existingPackDefinition = FolderUtilities::CombinePath(saveFolder, "hires.txt");
	if(ifstream(existingPackDefinition)) {
		HdPackLoader::LoadHdNesPack(existingPackDefinition, _hdData);
		for(unique_ptr<HdPackTileInfo> &tile : _hdData.Tiles) {
			//Mark the tiles in the first PNGs as higher usage (preserves order when adding new tiles to an existing set)
			AddTile(tile.get(), 0xFFFFFFFF - tile->BitmapIndex);
		}

		if(_hdData.Scale != scale) {
			_filterType = ScaleFilterType::Prescale;
		}
	} else {
		_hdData.Scale = scale;
	}

	if (_flags & (HdPackRecordFlags::SaveFrame)) {
		uint32_t foundID;
		vector<string> screenFiles = FolderUtilities::GetFilesInFolder(saveFolder, { ".png", ".csv" }, false);
		string filePrefix = FolderUtilities::CombinePath(_saveFolder, "");
		for (string path : screenFiles) {
			if (path.compare(filePrefix.length(), 7, "screen_") == 0) {
				
				foundID = std::stoi(path.substr(filePrefix.length() + 7, path.length() - filePrefix.length() - 11));
				if (foundID >= _frameID) {
					_frameID = foundID + 1;
				}
			}
		}
		FolderUtilities::CreateFolder(_saveFolder);
	}


	_romName = FolderUtilities::GetFilename(_console->GetRomInfo().RomName, false);
	_instance = this;
}

HdPackBuilder::~HdPackBuilder()
{

	SaveHdPack();
	if(_instance == this) {
		_instance = nullptr;
	}
}

void HdPackBuilder::endFrame() {
	if (_hasNewTile && (spritesOnScreen.size() > 0 || bgTilesOnScreen.size() > 0) && _flags & (HdPackRecordFlags::SaveFrame)) {
		//generate file name suffix
		string counterStr = std::to_string(_frameID);
		while (counterStr.length() < 3) {
			counterStr = "0" + counterStr;
		}

		//save screen
		_console->GetVideoDecoder()->TakeScreenshot(FolderUtilities::CombinePath(_saveFolder, "screen_" + counterStr + ".png"));

		//save screen info
		stringstream ss;
		ss << "Type,X,Y,Tile,Palette,Background Priority,Horizontal Mirroring,Vertical Mirroring,Is New," << std::endl;

		stringstream ssidx;
		if (_frameID == 0) {
			ssidx << "Type,Tile,Palette,Screen ID," << std::endl;
		}

		for (HdScreenTileInfo t:spritesOnScreen) {
			ss << "Sprite," << t.ScreenX << "," << t.ScreenY << ",";

			if (t.IsChrRamTile) {
				for (int i = 0; i < 16; i++) {
					ss << HexUtilities::ToHex(t.TileData[i]);
				}
			}
			else {
				ss << HexUtilities::ToHex(t.TileIndex);
			}
			ss << "," << HexUtilities::ToHex(t.PaletteColors, true) << "," <<
				(t.BackgroundPriority ? "Y" : "N") << "," << 
				(t.HorizontalMirroring ? "Y" : "N") << "," << 
				(t.VerticalMirroring ? "Y" : "N") << "," <<
				(t.IsNew ? "Y" : "N") << "," << std::endl;

			if (t.IsNew) {
				ssidx << "Sprite,";
				if (t.IsChrRamTile) {
					for (int i = 0; i < 16; i++) {
						ssidx << HexUtilities::ToHex(t.TileData[i]);
					}
				}
				else {
					ssidx << HexUtilities::ToHex(t.TileIndex);
				}
				ssidx << "," << HexUtilities::ToHex(t.PaletteColors, true) << "," << counterStr << "," << std::endl;
			}
		}

		for (HdScreenTileInfo t : bgTilesOnScreen) {
			ss << "Background," << t.ScreenX << "," << t.ScreenY << ",";

			if (t.IsChrRamTile) {
				for (int i = 0; i < 16; i++) {
					ss << HexUtilities::ToHex(t.TileData[i]);
				}
			}
			else {
				ss << HexUtilities::ToHex(t.TileIndex);
			}
			ss << "," << HexUtilities::ToHex(t.PaletteColors, true) << ",,,," << (t.IsNew ? "Y" : "N") << "," << std::endl;

			if (t.IsNew) {
				ssidx << "Background,";
				if (t.IsChrRamTile) {
					for (int i = 0; i < 16; i++) {
						ssidx << HexUtilities::ToHex(t.TileData[i]);
					}
				}
				else {
					ssidx << HexUtilities::ToHex(t.TileIndex);
				}
				ssidx << "," << HexUtilities::ToHex(t.PaletteColors, true) << "," << counterStr << "," << std::endl;
			}
		}

		ofstream screenInfoFile(FolderUtilities::CombinePath(_saveFolder, "screen_" + counterStr + ".csv"), ios::out);
		screenInfoFile << ss.str();
		screenInfoFile.close();

		ofstream tileIndexFile(FolderUtilities::CombinePath(_saveFolder, "tileIndex.csv"), ios::app);
		tileIndexFile << ssidx.str();
		tileIndexFile.close();

		_frameID++;
	}
	//reset for next frame
	_hasNewTile = false;
	spritesOnScreen.clear();
	bgTilesOnScreen.clear();
}

void HdPackBuilder::AddTile(HdPackTileInfo *tile, uint32_t usageCount)
{
	bool isTileBlank = (_flags & HdPackRecordFlags::GroupBlankTiles) ? tile->Blank : false;

	int chrBankId = isTileBlank ? 0xFFFFFFFF : tile->ChrBankId;
	int palette = isTileBlank ? _blankTilePalette : tile->PaletteColors;

	if(_tilesByChrBankByPalette.find(chrBankId) == _tilesByChrBankByPalette.end()) {
		_tilesByChrBankByPalette[chrBankId] = std::map<uint32_t, vector<HdPackTileInfo*>>();
	}

	std::map<uint32_t, vector<HdPackTileInfo*>> &paletteMap = _tilesByChrBankByPalette[chrBankId];
	if(paletteMap.find(palette) == paletteMap.end()) {
		paletteMap[palette] = vector<HdPackTileInfo*>(256, nullptr);
	}

	if(isTileBlank) {
		paletteMap[palette][_blankTileIndex] = tile;
		_blankTileIndex++;
		if(_blankTileIndex == _chrRamBankSize / 16) {
			_blankTileIndex = 0;
			_blankTilePalette++;
		}
	} else {
		if(tile->TileIndex >= 0) {
			paletteMap[palette][tile->TileIndex % 256] = tile;
		} else {
			//FIXME: This will result in data loss if more than 256 tiles of the same palette exist in the hires.txt file
			//Currently this way to prevent issues when loading a CHR RAM HD pack into the recorder (because TileIndex is -1 in that case)
			for(int i = 0; i < 256; i++) {
				if(paletteMap[palette][i] == nullptr) {
					paletteMap[palette][i] = tile;
					break;
				}
			}
		}
	}

	_tilesByKey[tile->GetKey(false)] = tile;
	_tileUsageCount[tile->GetKey(false)] = usageCount;
	_hasNewTile = true;
}

void HdPackBuilder::ProcessTile(uint32_t x, uint32_t y, uint16_t tileAddr, HdPpuTileInfo &tile, BaseMapper *mapper, uint32_t chrBankHash, bool transparencyRequired)
{
	if(_flags & HdPackRecordFlags::IgnoreOverscan) {
		OverscanDimensions overscan = _console->GetSettings()->GetOverscanDimensions();
		if(x < overscan.Left || y < overscan.Top || (PPU::ScreenWidth - x - 1) < overscan.Right || (PPU::ScreenHeight - y - 1) < overscan.Bottom) {
			//Ignore tiles inside overscan
			return;
		}
	}

	if(_outTileType == HDPackOuputTileType::BG && tile.IsSpriteTile()) {
		return;
	}

	if (_outTileType == HDPackOuputTileType::Sprite && !tile.IsSpriteTile()) {
		return;
	}

	auto result = _tileUsageCount.find(tile.GetKey(false));
	if(result == _tileUsageCount.end()) {
		//Check to see if a default tile matches
		result = _tileUsageCount.find(tile.GetKey(true));
	}
	bool isNew = false;
	if(result == _tileUsageCount.end()) {
		//First time seeing this tile/palette combination, store it
		HdPackTileInfo* hdTile = new HdPackTileInfo();
		hdTile->PaletteColors = tile.PaletteColors;
		hdTile->TileIndex = tile.TileIndex;
		hdTile->DefaultTile = false;
		hdTile->IsChrRamTile = _isChrRam;
		hdTile->Brightness = 255;
		hdTile->ChrBankId = _isChrRam ? chrBankHash : (tileAddr / 16 / 256);
		hdTile->TransparencyRequired = transparencyRequired;

		memcpy(hdTile->TileData, tile.TileData, 16);

		_hdData.Tiles.push_back(unique_ptr<HdPackTileInfo>(hdTile));
		AddTile(hdTile, 1);
		isNew = true;
	} else {
		if(transparencyRequired) {
			auto existingTile = _tilesByKey.find(tile.GetKey(false));
			if(existingTile != _tilesByKey.end()) {
				existingTile->second->TransparencyRequired = true;
			}
		}
		
		if(result->second < 0x7FFFFFFF) {
			//Increase usage count
			result->second++;
		}
	}

	if ((x == 0 || ((tile.OffsetX & 0x07) == 0)) && (y == 0 || (((tile.OffsetY & 0x07) == 0) && !tile.VerticalMirroring) || (((tile.OffsetY & 0x07) == 0x07) && tile.VerticalMirroring))) {
		HdScreenTileInfo t;
		t.IsChrRamTile = tile.IsChrRamTile;
		t.PaletteColors = tile.PaletteColors;
		memcpy(t.TileData, tile.TileData, 16);
		t.TileIndex = tile.TileIndex;
		t.IsNew = isNew;
		if (tile.IsSpriteTile()) {
			t.ScreenX = x;
			t.ScreenY = y;
			t.BackgroundPriority = tile.BackgroundPriority;
			t.HorizontalMirroring = tile.HorizontalMirroring;
			t.VerticalMirroring = tile.VerticalMirroring;
			spritesOnScreen.push_back(t);
		}
		else {
			t.ScreenX = x - tile.OffsetX;
			t.ScreenY = y - tile.OffsetY;
			bgTilesOnScreen.push_back(t);
		}
	}
}

void HdPackBuilder::GenerateHdTile(HdPackTileInfo *tile)
{
	uint32_t hdScale = _hdData.Scale;

	vector<uint32_t> originalTile = tile->ToRgb(_console->GetSettings()->GetRgbPalette());
	vector<uint32_t> hdTile(8 * 8 * hdScale*hdScale, 0);

	switch(_filterType) {
		case ScaleFilterType::HQX: 
			hqx(hdScale, originalTile.data(), hdTile.data(), 8, 8);
			break;

		case ScaleFilterType::Prescale:
			hdTile.clear();
			for(uint8_t i = 0; i < 8 * hdScale; i++) {
				for(uint8_t j = 0; j < 8 * hdScale; j++) {
					hdTile.push_back(originalTile[i/hdScale*8+j/hdScale]);
				}
			}
			break;

		case ScaleFilterType::Scale2x:
			scale(hdScale, hdTile.data(), 8 * sizeof(uint32_t) * hdScale, originalTile.data(), 8 * sizeof(uint32_t), 4, 8, 8);
			break;

		case ScaleFilterType::_2xSai:
			twoxsai_generic_xrgb8888(8, 8, originalTile.data(), 8, hdTile.data(), 8 * hdScale);
			break;

		case ScaleFilterType::Super2xSai:
			supertwoxsai_generic_xrgb8888(8, 8, originalTile.data(), 8, hdTile.data(), 8 * hdScale);
			break;

		case ScaleFilterType::SuperEagle:
			supereagle_generic_xrgb8888(8, 8, originalTile.data(), 8, hdTile.data(), 8 * hdScale);
			break;

		case ScaleFilterType::xBRZ:
			xbrz::scale(hdScale, originalTile.data(), hdTile.data(), 8, 8, xbrz::ColorFormat::ARGB);
			break;
	}

	tile->HdTileData = hdTile;
}

void HdPackBuilder::DrawTile(HdPackTileInfo *tile, int tileNumber, uint32_t *pngBuffer, int pageNumber, bool containsSpritesOnly)
{
	if(tile->HdTileData.empty()) {
		GenerateHdTile(tile);
		tile->UpdateFlags();
	}

	if(containsSpritesOnly && (_flags & HdPackRecordFlags::UseLargeSprites)) {
		int row = tileNumber / 16;
		int column = tileNumber % 16;

		int newColumn = column / 2 + ((row & 1) ? 8 : 0);
		int newRow = (row & 0xFE) + ((column & 1) ? 1 : 0);

		tileNumber = newRow * 16 + newColumn;
	}

	tileNumber += pageNumber * (256 / (0x1000 / _chrRamBankSize));

	int tileDimension = 8 * _hdData.Scale;
	int x = tileNumber % 16 * tileDimension;
	int y = tileNumber / 16 * tileDimension;

	tile->X = x;
	tile->Y = y;

	int pngWidth = 128 * _hdData.Scale;
	int pngPos = y * pngWidth + x;
	int tilePos = 0;
	for(uint8_t i = 0; i < tileDimension; i++) {
		for(uint8_t j = 0; j < tileDimension; j++) {
			pngBuffer[pngPos] = tile->HdTileData[tilePos++];
			pngPos++;
		}
		pngPos += pngWidth - tileDimension;
	}
}

void HdPackBuilder::SaveHdPack()
{
	FolderUtilities::CreateFolder(_saveFolder);

	stringstream pngRows;
	stringstream tileRows;
	stringstream ss;
	int pngIndex = 0;
	ss << "<ver>" << std::to_string(HdNesPack::CurrentVersion) << std::endl;
	ss << "<scale>" << _hdData.Scale << std::endl;
	ss << "<supportedRom>" << _console->GetRomInfo().Hash.Sha1 << std::endl;
	if(_flags & HdPackRecordFlags::IgnoreOverscan) {
		OverscanDimensions overscan = _console->GetSettings()->GetOverscanDimensions();
		ss << "<overscan>" << overscan.Top << "," << overscan.Right << "," << overscan.Bottom << "," << overscan.Left << std::endl;
	}

	int tileDimension = 8 * _hdData.Scale;
	int pngDimension = 16 * tileDimension;
	int pngBufferSize = pngDimension * pngDimension;
	uint32_t* pngBuffer = new uint32_t[pngBufferSize];

	int maxPageNumber = 0x1000 / _chrRamBankSize;
	int pageNumber = 0;
	bool pngEmpty = true;
	int pngNumber = 0;

	for(int i = 0; i < pngBufferSize; i++) {
		pngBuffer[i] = 0xFFFF00FF;
	}

	auto savePng = [&tileRows, &pngRows, &ss, &pngBuffer, &pngDimension, &pngIndex, &pngBufferSize, &pngEmpty, &pngNumber, this](uint32_t chrBankId) {
		if(!pngEmpty) {
			string pngName;
			if(_isChrRam) {
				pngName = "Chr_" + std::to_string(pngNumber) + ".png";
			} else {
				pngName = "Chr_" + HexUtilities::ToHex(chrBankId) + "_" + std::to_string(pngNumber) + ".png";
			}

			tileRows << std::endl << "#" << pngName << std::endl;
			tileRows << pngRows.str();
			pngRows = stringstream();

			ss << "<img>" << pngName << std::endl;
			PNGHelper::WritePNG(FolderUtilities::CombinePath(_saveFolder, pngName), pngBuffer, pngDimension, pngDimension, 32);
			pngNumber++;
			pngIndex++;

			for(int i = 0; i < pngBufferSize; i++) {
				pngBuffer[i] = 0xFFFF00FF;
			}
			pngEmpty = true;
		}
	};

	for(std::pair<const uint32_t, std::map<uint32_t, vector<HdPackTileInfo*>>> &kvp : _tilesByChrBankByPalette) {
		if(_flags & HdPackRecordFlags::SortByUsageFrequency) {
			for(int i = 0; i < 256; i++) {
				vector<std::pair<uint32_t, HdPackTileInfo*>> tiles;
				for(std::pair<const uint32_t, vector<HdPackTileInfo*>> &paletteMap : kvp.second) {
					if(paletteMap.second[i]) {
						tiles.push_back({ _tileUsageCount[paletteMap.second[i]->GetKey(false)], paletteMap.second[i] });
					}
				}
				std::sort(tiles.begin(), tiles.end(), [=](std::pair<uint32_t, HdPackTileInfo*> &a, std::pair<uint32_t, HdPackTileInfo*> &b) {
					return a.first > b.first;
				});

				size_t j = 0;
				for(std::pair<const uint32_t, vector<HdPackTileInfo*>> &paletteMap : kvp.second) {
					if(j < tiles.size()) {
						paletteMap.second[i] = tiles[j].second;
						j++;
					} else {
						paletteMap.second[i] = nullptr;
					}
				}
			}
		}

		if(!_isChrRam) {
			pngNumber = 0;
		}

		for(std::pair<const uint32_t, vector<HdPackTileInfo*>> &tileKvp : kvp.second) {
			bool pageEmpty = true;
			bool spritesOnly = true;
			for(HdPackTileInfo* tileInfo : tileKvp.second) {
				if(tileInfo && !tileInfo->IsSpriteTile()) {
					spritesOnly = false;
				}
			}

			for(int i = 0; i < 256; i++) {
				HdPackTileInfo* tileInfo = tileKvp.second[i];
				if(tileInfo) {
					DrawTile(tileInfo, i, pngBuffer, pageNumber, spritesOnly);

					pngRows << tileInfo->ToString(pngIndex) << std::endl;

					pageEmpty = false;
					pngEmpty = false;
				}
			}

			if(!pageEmpty) {
				pageNumber++;

				if(pageNumber == maxPageNumber) {
					savePng(kvp.first);
					pageNumber = 0;
				}
			}
		}
	}
	savePng(-1);

	for(unique_ptr<HdPackCondition> &condition : _hdData.Conditions) {
		if(!condition->IsExcludedFromFile()) {
			ss << condition->ToString() << std::endl;
		}
	}

	for(HdBackgroundInfo &bgInfo : _hdData.Backgrounds) {
		ss << bgInfo.ToString() << std::endl;
	}

	for(auto &bgmInfo : _hdData.BgmFilesById) {
		ss << "<bgm>" << std::to_string(bgmInfo.first >> 8) << "," << std::to_string(bgmInfo.first & 0xFF) << "," << VirtualFile(bgmInfo.second.File).GetFileName() << "," << bgmInfo.second.LoopPoint << "," << bgmInfo.second.PlaybackOptions << "," << bgmInfo.second.Volume << std::endl;
	}

	for(auto &sfxInfo : _hdData.SfxFilesById) {
		ss << "<sfx>" << std::to_string(sfxInfo.first >> 8) << "," << std::to_string(sfxInfo.first & 0xFF) << "," << VirtualFile(sfxInfo.second).GetFileName() << std::endl;
	}

	for(auto &patchInfo : _hdData.PatchesByHash) {
		ss << "<patch>" << VirtualFile(patchInfo.second).GetFileName() << "," << patchInfo.first << std::endl;
	}

	if(_hdData.OptionFlags != 0) {
		ss << "<options>";
		if(_hdData.OptionFlags & (int)HdPackOptions::NoSpriteLimit) {
			ss << "disableSpriteLimit,";
		}
		if(_hdData.OptionFlags & (int)HdPackOptions::AlternateRegisterRange) {
			ss << "alternateRegisterRange,";
		}
		if(_hdData.OptionFlags & (int)HdPackOptions::DisableCache) {
			ss << "disableCache,";
		}
	}

	ss << tileRows.str();

	ofstream hiresFile(FolderUtilities::CombinePath(_saveFolder, "hires.txt"), ios::out);
	hiresFile << ss.str();
	hiresFile.close();

	delete[] pngBuffer;
}

void HdPackBuilder::GetChrBankList(uint32_t *banks)
{
	ConsolePauseHelper helper(_instance->_console.get());
	for(std::pair<const uint32_t, std::map<uint32_t, vector<HdPackTileInfo*>>> &kvp : _instance->_tilesByChrBankByPalette) {
		*banks = kvp.first;
		banks++;
	}
	*banks = -1;
}

void HdPackBuilder::GetBankPreview(uint32_t bankNumber, uint32_t pageNumber, uint32_t *rgbBuffer)
{
	ConsolePauseHelper helper(_instance->_console.get());

	for(uint32_t i = 0; i < 128 * 128 * _instance->_hdData.Scale*_instance->_hdData.Scale; i++) {
		rgbBuffer[i] = 0xFF666666;
	}

	auto result = _instance->_tilesByChrBankByPalette.find(bankNumber);
	if(result != _instance->_tilesByChrBankByPalette.end()) {
		std::map<uint32_t, vector<HdPackTileInfo*>> bankData = result->second;

		if(_instance->_flags & HdPackRecordFlags::SortByUsageFrequency) {
			for(int i = 0; i < 256; i++) {
				vector<std::pair<uint32_t, HdPackTileInfo*>> tiles;
				for(std::pair<const uint32_t, vector<HdPackTileInfo*>> &pageData : bankData) {
					if(pageData.second[i]) {
						tiles.push_back({ _instance->_tileUsageCount[pageData.second[i]->GetKey(false)], pageData.second[i] });
					}
				}

				std::sort(tiles.begin(), tiles.end(), [=](std::pair<uint32_t, HdPackTileInfo*> &a, std::pair<uint32_t, HdPackTileInfo*> &b) {
					return a.first > b.first;
				});

				size_t j = 0;
				for(std::pair<const uint32_t, vector<HdPackTileInfo*>> &pageData : bankData) {
					if(j < tiles.size()) {
						pageData.second[i] = tiles[j].second;
						j++;
					} else {
						pageData.second[i] = nullptr;
					}
				}
			}
		}

		bool spritesOnly = true;
		for(HdPackTileInfo* tileInfo : (*bankData.begin()).second) {
			if(tileInfo && !tileInfo->IsSpriteTile()) {
				spritesOnly = false;
			}
		}

		for(int i = 0; i < 256; i++) {
			HdPackTileInfo* tileInfo = (*bankData.begin()).second[i];
			if(tileInfo) {
				_instance->DrawTile(tileInfo, i, (uint32_t*)rgbBuffer, 0, spritesOnly);
			}
		}
	}
}