/*
Copyright (c) 2017-2018 Adubbz

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

#include <filesystem>
#include <sstream>
#include <switch.h>
#include "nx/nca_writer.h"
#include "util/error.hpp"
#include <zstd.h>
#include <string.h>
#include "util/crypto.hpp"
#include "util/config.hpp"
#include "util/title_util.hpp"
#include "install/nca.hpp"

//added for debugging messages on screen
#include "util/lang.hpp"
#include "ui/MainApplication.hpp"
//

void append(std::vector<u8>& buffer, const u8* ptr, u64 sz)
{
	u64 offset = buffer.size();
	buffer.resize(offset + sz);
	memcpy(buffer.data() + offset, ptr, sz);
}

NcaBodyWriter::NcaBodyWriter(const NcmContentId& ncaId, u64 offset, std::shared_ptr<nx::ncm::ContentStorage>& contentStorage) : m_contentStorage(contentStorage), m_ncaId(ncaId), m_offset(offset)
{
}

NcaBodyWriter::~NcaBodyWriter()
{
}

u64 NcaBodyWriter::write(const  u8* ptr, u64 sz)
{
	if (isOpen())
	{
		m_contentStorage->WritePlaceholder(*(NcmPlaceHolderId*)&m_ncaId, m_offset, (void*)ptr, sz);
		m_offset += sz;
		return sz;
	}

	return 0;
}

bool NcaBodyWriter::isOpen() const
{
	return m_contentStorage != NULL;
}

//code - https://github.com/nicoboss/nsz
//https://github.com/nicoboss/nsz/blob/master/nsz/NszDecompressor.py
//https://www.w3schools.com/cpp/cpp_classes.asp
//https://github.com/minetest/minetestmapper/blob/master/ZstdDecompressor.cpp
//https://github.com/nicoboss/nsz/blob/master/nsz/BlockDecompressorReader.py
//https://switchbrew.org/wiki/NCA

namespace inst::ui {
	extern MainApplication* mainApp;
}

class NczHeader
{
public:
	static const u64 MAGIC = 0x4E544345535A434E; //NTCESZCN
	static const u64 BLOCK = 0x4E435A424C4F434B; //NCZBLOCK at 0x40D0

	class Section
	{
	public:
		u64 offset;
		u64 size;
		u8 cryptoType;
		u8 padding1[7];
		u64 padding2;
		u8 cryptoKey[0x10];
		u8 cryptoCounter[0x10];
	} PACKED;

	class SectionContext : public Section
	{
	public:
		Crypto::Aes128Ctr crypto;
		
		SectionContext(const Section& s) : Section(s), crypto(s.cryptoKey, Crypto::AesCtr(Crypto::swapEndian(((u64*)&s.cryptoCounter)[0])))
		{
		}

		virtual ~SectionContext()
		{
		}

		void decrypt(void* p, u64 sz, u64 offset)
		{
			if (this->cryptoType == 3 || this->cryptoType == 4)
			{
				crypto.seek(offset);
				crypto.decrypt(p, p, sz);
			}
			else {
				return;
			}
		}

		void encrypt(void* p, u64 sz, u64 offset)
		{
			if (this->cryptoType == 3 || this->cryptoType == 4)
			{
				crypto.seek(offset);
				crypto.encrypt(p, p, sz);
			}
			else {
				return;
			}
		}
	};

	const bool isValid()
	{
		return m_magic == MAGIC && m_sectionCount < 0xFFFF;
	}

	const u64 size() const
	{
		return sizeof(m_magic) + sizeof(m_sectionCount) + sizeof(Section) * m_sectionCount;
	}

	const Section& section(u64 i) const
	{
		return m_sections[i];
	}

	const u64 sectionCount() const
	{
		return m_sectionCount;
	}

protected:
	u64 m_magic;
	u64 m_sectionCount;
	Section m_sections[1];
} PACKED;

class NczBodyWriter : public NcaBodyWriter
{
public:
	NczBodyWriter(const NcmContentId& ncaId, u64 offset, std::shared_ptr<nx::ncm::ContentStorage>& contentStorage) : NcaBodyWriter(ncaId, offset, contentStorage)
	{
		buffIn = malloc(buffInSize);
		buffOut = malloc(buffOutSize);

		dctx = ZSTD_createDCtx();
	}

	virtual ~NczBodyWriter()
	{
		close();

		for (auto& i : sections)
		{
			if (i)
			{
				delete i;
				i = NULL;
			}
		}

		if (dctx)
		{
			ZSTD_freeDCtx(dctx);
			dctx = NULL;
		}
	}

	bool close()
	{
		if (this->m_buffer.size())
		{
			processChunk(m_buffer.data(), m_buffer.size());
		}

		encrypt(m_deflateBuffer.data(), m_deflateBuffer.size(), m_offset);
		flush();

		return true;
	}

	bool flush()
	{
		if (!isOpen())
		{
			return false;
		}

		if (m_deflateBuffer.size())
		{
			m_contentStorage->WritePlaceholder(*(NcmPlaceHolderId*)&m_ncaId, m_offset, m_deflateBuffer.data(), m_deflateBuffer.size());
			m_offset += m_deflateBuffer.size();
			m_deflateBuffer.resize(0);
		}
		return true;
	}

	NczHeader::SectionContext& section(u64 offset)
	{
		for (u64 i = 0; i < sections.size(); i++)
		{
			if (offset >= sections[i]->offset && offset < sections[i]->offset + sections[i]->size)
			{
				return *sections[i];
			}
		}
		return *sections[0];
	}

	bool encrypt(const void* ptr, u64 sz, u64 offset)
	{
		const u8* start = (u8*)ptr;
		const u8* end = start + sz;

		while (start < end)
		{
			auto& s = section(offset);

			u64 sectionEnd = s.offset + s.size;

			u64 chunk = offset + sz > sectionEnd ? sectionEnd - offset : sz;

			s.encrypt((void*)start, chunk, offset);

			offset += chunk;
			start += chunk;
			sz -= chunk;
		}

		return true;
	}

	u64 processChunk(const u8* ptr, u64 sz)
	{
		while (sz > 0)
		{
			const size_t readChunkSz = std::min(sz, buffInSize);
			ZSTD_inBuffer input = { ptr, readChunkSz, 0 };

			while (input.pos < input.size)
			{
				ZSTD_outBuffer output = { buffOut, buffOutSize, 0 };
				size_t const ret = ZSTD_decompressStream(dctx, &output, &input);

				if (ZSTD_isError(ret))
				{
					LOG_DEBUG("%s\n", ZSTD_getErrorName(ret));
					return 0;
				}

				size_t len = output.pos;
				u8* p = (u8*)buffOut;

				while (len)
				{
					const size_t writeChunkSz = std::min(0x1000000 - m_deflateBuffer.size(), len);

					append(m_deflateBuffer, p, writeChunkSz);

					if (m_deflateBuffer.size() >= 0x1000000)
					{
						encrypt(m_deflateBuffer.data(), m_deflateBuffer.size(), m_offset);
						flush();
					}

					p += writeChunkSz;
					len -= writeChunkSz;
				}
			}

			sz -= readChunkSz;
			ptr += readChunkSz;
		}

		return 1;
	}

	u64 write(const  u8* ptr, u64 sz) override
	{
		if (!m_sectionsInitialized)
		{
			if (!m_buffer.size())
			{
				append(m_buffer, ptr, sizeof(u64) * 2);
				ptr += sizeof(u64) * 2;
				sz -= sizeof(u64) * 2;
			}

			auto header = (NczHeader*)m_buffer.data();

			//debug code -- ignore....
			/*
			int x =0;
			if (x == 0) {
				char aByte = m_buffer.data()[0];
				//aByte = 0x10+(aByte*0x40);
				int one = sizeof(NczHeader);
				std::ostringstream ss;
				ss << "0x" << std::hex << one;
				std::string result = ss.str();
				inst::ui::mainApp->CreateShowDialog("Debug info", result, { "common.ok"_lang }, true, "romfs:/images/icons/information.png");
				x++;
			}
			*/

			if (m_buffer.size() + sz > header->size())
			{
				u64 remainder = header->size() - m_buffer.size();
				append(m_buffer, ptr, remainder);
				ptr += remainder;
				sz -= remainder;
			}
			else
			{
				append(m_buffer, ptr, sz);
				ptr += sz;
				sz = 0;
			}

			header = (NczHeader*)m_buffer.data();

			if (m_buffer.size() == header->size())
			{
				for (u64 i = 0; i < header->sectionCount(); i++)
				{
					sections.push_back(new NczHeader::SectionContext(header->section(i)));
				}

				m_sectionsInitialized = true;
				m_buffer.resize(0);
			}
		}

		while (sz)
		{
			if (m_buffer.size() + sz >= 0x1000000)
			{
				u64 chunk = 0x1000000 - m_buffer.size();
				append(m_buffer, ptr, chunk);

				processChunk(m_buffer.data(), m_buffer.size());
				m_buffer.resize(0);

				sz -= chunk;
				ptr += chunk;
			}
			else
			{
				append(m_buffer, ptr, sz);
				sz = 0;
			}

		}

		return sz;
	}

	size_t const buffInSize = ZSTD_DStreamInSize();
	size_t const buffOutSize = ZSTD_DStreamOutSize();

	void* buffIn = NULL;
	void* buffOut = NULL;

	ZSTD_DCtx* dctx = NULL;

	std::vector<u8> m_buffer;
	std::vector<u8> m_deflateBuffer;

	bool m_sectionsInitialized = false;

	std::vector<NczHeader::SectionContext*> sections;
};

NcaWriter::NcaWriter(const NcmContentId& ncaId, std::shared_ptr<nx::ncm::ContentStorage>& contentStorage) : m_ncaId(ncaId), m_contentStorage(contentStorage), m_writer(NULL)
{
}

NcaWriter::~NcaWriter()
{
	close();
}

bool NcaWriter::close()
{
	if (m_writer)
	{
		m_writer = NULL;
	}
	else if (m_buffer.size())
	{
		if (isOpen())
		{
			flushHeader();
		}

		m_buffer.resize(0);
	}
	m_contentStorage = NULL;
	return true;
}

bool NcaWriter::isOpen() const
{
	return (bool)m_contentStorage;
}

u64 NcaWriter::write(const  u8* ptr, u64 sz)
{
	if (m_buffer.size() < NCA_HEADER_SIZE)
	{
		if (m_buffer.size() + sz > NCA_HEADER_SIZE)
		{
			u64 remainder = NCA_HEADER_SIZE - m_buffer.size();
			append(m_buffer, ptr, remainder);

			ptr += remainder;
			sz -= remainder;
		}
		else
		{
			append(m_buffer, ptr, sz);
			ptr += sz;
			sz = 0;
		}

		if (m_buffer.size() == NCA_HEADER_SIZE)
		{
			flushHeader();
		}
	}

	if (sz)
	{
		if (!m_writer)
		{
			if (sz >= sizeof(NczHeader::MAGIC))
			{
				if (*(u64*)ptr == NczHeader::MAGIC)
				{
					m_writer = std::shared_ptr<NcaBodyWriter>(new NczBodyWriter(m_ncaId, m_buffer.size(), m_contentStorage));
				}
				else
				{
					m_writer = std::shared_ptr<NcaBodyWriter>(new NcaBodyWriter(m_ncaId, m_buffer.size(), m_contentStorage));
				}
			}
			else
			{
				THROW_FORMAT("not enough data to read ncz header");
			}
		}

		if (m_writer)
		{
			m_writer->write(ptr, sz);
		}
		else
		{
			THROW_FORMAT("null writer");
		}
	}

	return sz;
}

void NcaWriter::flushHeader()
{
	tin::install::NcaHeader header;
	memcpy(&header, m_buffer.data(), sizeof(header));
	Crypto::AesXtr decryptor(Crypto::Keys().headerKey, false);
	Crypto::AesXtr encryptor(Crypto::Keys().headerKey, true);
	decryptor.decrypt(&header, &header, sizeof(header), 0, 0x200);

	if (header.magic == MAGIC_NCA3)
	{
		if (isOpen())
		{
			m_contentStorage->CreatePlaceholder(m_ncaId, *(NcmPlaceHolderId*)&m_ncaId, header.nca_size);
		}
	}
	else
	{
		THROW_FORMAT("Invalid NCA magic");
	}

	if (header.distribution == 1)
	{
		header.distribution = 0;
	}
	encryptor.encrypt(m_buffer.data(), &header, sizeof(header), 0, 0x200);

	if (isOpen())
	{
		m_contentStorage->WritePlaceholder(*(NcmPlaceHolderId*)&m_ncaId, 0, m_buffer.data(), m_buffer.size());
	}
}