// Copyright Epic Games, Inc. All Rights Reserved.

#include <zencore/iobuffer.h>

#include <zencore/except.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iohash.h>
#include <zencore/logging.h>
#include <zencore/memory.h>
#include <zencore/testing.h>
#include <zencore/thread.h>

#include <memory.h>
#include <system_error>

#if ZEN_USE_MIMALLOC
ZEN_THIRD_PARTY_INCLUDES_START
#	include <mimalloc.h>
ZEN_THIRD_PARTY_INCLUDES_END
#endif

#if ZEN_PLATFORM_WINDOWS
#	include <atlfile.h>
#else
#	include <sys/stat.h>
#	include <sys/mman.h>
#endif

#include <gsl/gsl-lite.hpp>

namespace zen {

//////////////////////////////////////////////////////////////////////////

void
IoBufferCore::AllocateBuffer(size_t InSize, size_t Alignment)
{
#if ZEN_PLATFORM_WINDOWS
	if (((InSize & 0xffFF) == 0) && (Alignment == 0x10000))
	{
		m_Flags.fetch_or(kLowLevelAlloc, std::memory_order_relaxed);
		m_DataPtr = VirtualAlloc(nullptr, InSize, MEM_COMMIT, PAGE_READWRITE);

		return;
	}
#endif	// ZEN_PLATFORM_WINDOWS

#if ZEN_USE_MIMALLOC
	void* Ptr = mi_aligned_alloc(Alignment, RoundUp(InSize, Alignment));
	m_Flags.fetch_or(kIoBufferAlloc, std::memory_order_relaxed);
#else
	void* Ptr = Memory::Alloc(InSize, Alignment);
#endif

	ZEN_ASSERT(Ptr);

	m_DataPtr = Ptr;
}

void
IoBufferCore::FreeBuffer()
{
	if (!m_DataPtr)
	{
		return;
	}

	const uint32_t LocalFlags = m_Flags.load(std::memory_order_relaxed);
#if ZEN_PLATFORM_WINDOWS
	if (LocalFlags & kLowLevelAlloc)
	{
		VirtualFree(const_cast<void*>(m_DataPtr), 0, MEM_DECOMMIT);

		return;
	}
#endif	// ZEN_PLATFORM_WINDOWS

#if ZEN_USE_MIMALLOC
	if (LocalFlags & kIoBufferAlloc)
	{
		return mi_free(const_cast<void*>(m_DataPtr));
	}
#endif

	ZEN_UNUSED(LocalFlags);
	return Memory::Free(const_cast<void*>(m_DataPtr));
}

//////////////////////////////////////////////////////////////////////////

static_assert(sizeof(IoBufferCore) == 32);

IoBufferCore::IoBufferCore(size_t InSize)
{
	ZEN_ASSERT(InSize);

	AllocateBuffer(InSize, sizeof(void*));
	m_DataBytes = InSize;

	SetIsOwnedByThis(true);
}

IoBufferCore::IoBufferCore(size_t InSize, size_t Alignment)
{
	ZEN_ASSERT(InSize);

	AllocateBuffer(InSize, Alignment);
	m_DataBytes = InSize;

	SetIsOwnedByThis(true);
}

IoBufferCore::~IoBufferCore()
{
	if (IsOwnedByThis() && m_DataPtr)
	{
		FreeBuffer();
		m_DataPtr = nullptr;
	}
}

void
IoBufferCore::DeleteThis() const
{
	// We do this just to avoid paying for the cost of a vtable
	if (const IoBufferExtendedCore* _ = ExtendedCore())
	{
		delete _;
	}
	else
	{
		delete this;
	}
}

void
IoBufferCore::Materialize() const
{
	if (const IoBufferExtendedCore* _ = ExtendedCore())
	{
		_->Materialize();
	}
}

void
IoBufferCore::MakeOwned(bool Immutable)
{
	if (!IsOwned())
	{
		const void* OldDataPtr = m_DataPtr;
		AllocateBuffer(m_DataBytes, sizeof(void*));
		memcpy(const_cast<void*>(m_DataPtr), OldDataPtr, m_DataBytes);
		SetIsOwnedByThis(true);
	}

	SetIsImmutable(Immutable);
}

void*
IoBufferCore::MutableDataPointer() const
{
	EnsureDataValid();
	ZEN_ASSERT(!IsImmutable());
	return const_cast<void*>(m_DataPtr);
}

//////////////////////////////////////////////////////////////////////////

IoBufferExtendedCore::IoBufferExtendedCore(void* FileHandle, uint64_t Offset, uint64_t Size, bool TransferHandleOwnership)
: IoBufferCore(nullptr, Size)
, m_FileHandle(FileHandle)
, m_FileOffset(Offset)
{
	uint32_t NewFlags = kIsOwnedByThis | kIsExtended;

	if (TransferHandleOwnership)
	{
		NewFlags |= kOwnsFile;
	}
	m_Flags.fetch_or(NewFlags, std::memory_order_relaxed);
}

IoBufferExtendedCore::IoBufferExtendedCore(const IoBufferExtendedCore* Outer, uint64_t Offset, uint64_t Size)
: IoBufferCore(Outer, nullptr, Size)
, m_FileHandle(Outer->m_FileHandle)
, m_FileOffset(Outer->m_FileOffset + Offset)
{
	m_Flags.fetch_or(kIsExtended, std::memory_order_relaxed);
}

IoBufferExtendedCore::~IoBufferExtendedCore()
{
	if (m_MappedPointer)
	{
#if ZEN_PLATFORM_WINDOWS
		UnmapViewOfFile(m_MappedPointer);
#else
		uint64_t MapSize = ~uint64_t(uintptr_t(m_MmapHandle));
		munmap(m_MappedPointer, MapSize);
#endif
	}

	const uint32_t LocalFlags = m_Flags.load(std::memory_order_relaxed);
#if ZEN_PLATFORM_WINDOWS
	if (LocalFlags & kOwnsMmap)
	{
		CloseHandle(m_MmapHandle);
	}
#endif

	if (LocalFlags & kOwnsFile)
	{
		if (m_DeleteOnClose)
		{
#if ZEN_PLATFORM_WINDOWS
			// Mark file for deletion when final handle is closed
			FILE_DISPOSITION_INFO Fdi{.DeleteFile = TRUE};

			SetFileInformationByHandle(m_FileHandle, FileDispositionInfo, &Fdi, sizeof Fdi);
#else
			std::filesystem::path FilePath = zen::PathFromHandle(m_FileHandle);
			unlink(FilePath.c_str());
#endif
		}
#if ZEN_PLATFORM_WINDOWS
		BOOL Success = CloseHandle(m_FileHandle);
#else
		int	 Fd		 = int(uintptr_t(m_FileHandle));
		bool Success = (close(Fd) == 0);
#endif
		if (!Success)
		{
			ZEN_WARN("Error reported on file handle close, reason '{}'", GetLastErrorAsString());
		}
	}

	m_DataPtr = nullptr;
}

static constexpr size_t MappingLockCount = 64;
static_assert(IsPow2(MappingLockCount), "MappingLockCount must be power of two");

static RwLock g_MappingLocks[MappingLockCount];

static RwLock&
MappingLockForInstance(const IoBufferExtendedCore* instance)
{
	intptr_t base		= (intptr_t)instance;
	size_t	 lock_index = ((base >> 8) ^ (base >> 16)) & (MappingLockCount - 1u);
	return g_MappingLocks[lock_index];
}

void
IoBufferExtendedCore::Materialize() const
{
	// The synchronization scheme here is very primitive, if we end up with
	// a lot of contention we can make it more fine-grained

	if (m_Flags.load(std::memory_order_acquire) & kIsMaterialized)
		return;

	RwLock::ExclusiveLockScope _(MappingLockForInstance(this));

	// Someone could have gotten here first
	// We can use memory_order_relaxed on this load because the mutex has already provided the fence
	if (m_Flags.load(std::memory_order_relaxed) & kIsMaterialized)
		return;

	void*	 NewMmapHandle;
	uint32_t NewFlags = kIsMaterialized;

	const uint64_t MapOffset				= m_FileOffset & ~0xffffull;
	const uint64_t MappedOffsetDisplacement = m_FileOffset - MapOffset;
	const uint64_t MapSize					= m_DataBytes + MappedOffsetDisplacement;

	ZEN_ASSERT(MapSize > 0);

#if ZEN_PLATFORM_WINDOWS
	NewMmapHandle = CreateFileMapping(m_FileHandle,
									  /* lpFileMappingAttributes */ nullptr,
									  /* flProtect */ PAGE_READONLY,
									  /* dwMaximumSizeLow */ 0,
									  /* dwMaximumSizeHigh */ 0,
									  /* lpName */ nullptr);

	if (NewMmapHandle == nullptr)
	{
		throw std::system_error(std::error_code(::GetLastError(), std::system_category()),
								fmt::format("CreateFileMapping failed on file '{}'", zen::PathFromHandle(m_FileHandle)));
	}

	NewFlags |= kOwnsMmap;

	void* MappedBase = MapViewOfFile(NewMmapHandle,
									 /* dwDesiredAccess */ FILE_MAP_READ,
									 /* FileOffsetHigh */ uint32_t(MapOffset >> 32),
									 /* FileOffsetLow */ uint32_t(MapOffset & 0xffFFffFFu),
									 /* dwNumberOfBytesToMap */ MapSize);
#else
	NewMmapHandle = (void*)uintptr_t(~MapSize);	 // ~ so it's never null (assuming MapSize >= 0)
	NewFlags |= kOwnsMmap;

	void* MappedBase = mmap(
		/* addr */ nullptr,
		/* length */ MapSize,
		/* prot */ PROT_READ,
		/* flags */ MAP_SHARED | MAP_NORESERVE,
		/* fd */ int(uintptr_t(m_FileHandle)),
		/* offset */ MapOffset);
#endif	// ZEN_PLATFORM_WINDOWS

	if (MappedBase == nullptr)
	{
#if ZEN_PLATFORM_WINDOWS
		CloseHandle(NewMmapHandle);
#endif	// ZEN_PLATFORM_WINDOWS
		throw std::system_error(std::error_code(zen::GetLastError(), std::system_category()),
								fmt::format("MapViewOfFile failed (offset {:#x}, size {:#x}) file: '{}'",
											MapOffset,
											MapSize,
											zen::PathFromHandle(m_FileHandle)));
	}

	m_MappedPointer = MappedBase;
	m_DataPtr		= reinterpret_cast<uint8_t*>(MappedBase) + MappedOffsetDisplacement;
	m_MmapHandle	= NewMmapHandle;

	m_Flags.fetch_or(NewFlags, std::memory_order_release);
}

bool
IoBufferExtendedCore::GetFileReference(IoBufferFileReference& OutRef) const
{
	if (m_FileHandle == nullptr)
	{
		return false;
	}

	OutRef.FileHandle	   = m_FileHandle;
	OutRef.FileChunkOffset = m_FileOffset;
	OutRef.FileChunkSize   = m_DataBytes;

	return true;
}

void
IoBufferExtendedCore::MarkAsDeleteOnClose()
{
	m_DeleteOnClose = true;
}

//////////////////////////////////////////////////////////////////////////

IoBuffer::IoBuffer(size_t InSize) : m_Core(new IoBufferCore(InSize))
{
	m_Core->SetIsImmutable(false);
}

IoBuffer::IoBuffer(size_t InSize, uint64_t InAlignment) : m_Core(new IoBufferCore(InSize, InAlignment))
{
	m_Core->SetIsImmutable(false);
}

IoBuffer::IoBuffer(const IoBuffer& OuterBuffer, size_t Offset, size_t Size)
{
	if (Size == ~(0ull))
	{
		Size = std::clamp<size_t>(Size, 0, OuterBuffer.Size() - Offset);
	}

	ZEN_ASSERT(Offset <= OuterBuffer.Size());
	ZEN_ASSERT((Offset + Size) <= OuterBuffer.Size());

	if (IoBufferExtendedCore* Extended = OuterBuffer.m_Core->ExtendedCore())
	{
		m_Core = new IoBufferExtendedCore(Extended, Offset, Size);
	}
	else
	{
		m_Core = new IoBufferCore(OuterBuffer.m_Core, reinterpret_cast<const uint8_t*>(OuterBuffer.Data()) + Offset, Size);
	}
}

IoBuffer::IoBuffer(EFileTag, void* FileHandle, uint64_t ChunkFileOffset, uint64_t ChunkSize)
: m_Core(new IoBufferExtendedCore(FileHandle, ChunkFileOffset, ChunkSize, /* owned */ true))
{
}

IoBuffer::IoBuffer(EBorrowedFileTag, void* FileHandle, uint64_t ChunkFileOffset, uint64_t ChunkSize)
: m_Core(new IoBufferExtendedCore(FileHandle, ChunkFileOffset, ChunkSize, /* owned */ false))
{
}

bool
IoBuffer::GetFileReference(IoBufferFileReference& OutRef) const
{
	if (IoBufferExtendedCore* ExtCore = m_Core->ExtendedCore())
	{
		if (ExtCore->GetFileReference(OutRef))
		{
			return true;
		}
	}

	// Not a file reference

	OutRef.FileHandle	   = 0;
	OutRef.FileChunkOffset = ~0ull;
	OutRef.FileChunkSize   = 0;

	return false;
}

void
IoBuffer::MarkAsDeleteOnClose()
{
	if (IoBufferExtendedCore* ExtCore = m_Core->ExtendedCore())
	{
		ExtCore->MarkAsDeleteOnClose();
	}
}

//////////////////////////////////////////////////////////////////////////

IoBuffer
IoBufferBuilder::ReadFromFileMaybe(IoBuffer& InBuffer)
{
	IoBufferFileReference FileRef;
	if (InBuffer.GetFileReference(/* out */ FileRef))
	{
		IoBuffer OutBuffer(FileRef.FileChunkSize);

#if ZEN_PLATFORM_WINDOWS
		OVERLAPPED Ovl{};

		const uint64_t	NumberOfBytesToRead = FileRef.FileChunkSize;
		const uint64_t& FileOffset			= FileRef.FileChunkOffset;

		Ovl.Offset	   = DWORD(FileOffset & 0xffff'ffffu);
		Ovl.OffsetHigh = DWORD(FileOffset >> 32);

		DWORD dwNumberOfBytesRead = 0;
		BOOL  Success = ::ReadFile(FileRef.FileHandle, OutBuffer.MutableData(), DWORD(NumberOfBytesToRead), &dwNumberOfBytesRead, &Ovl);
#else
		int	 Fd		 = int(intptr_t(FileRef.FileHandle));
		int	 Result	 = pread(Fd, OutBuffer.MutableData(), size_t(FileRef.FileChunkSize), off_t(FileRef.FileChunkOffset));
		bool Success = (Result < 0);

		uint32_t dwNumberOfBytesRead = uint32_t(Result);
#endif

		if (!Success)
		{
			ThrowLastError("ReadFile failed in IoBufferBuilder::ReadFromFileMaybe");
		}

		ZEN_ASSERT(dwNumberOfBytesRead == FileRef.FileChunkSize);

		return OutBuffer;
	}
	else
	{
		return InBuffer;
	}
}

IoBuffer
IoBufferBuilder::MakeFromFileHandle(void* FileHandle, uint64_t Offset, uint64_t Size)
{
	return IoBuffer(IoBuffer::BorrowedFile, FileHandle, Offset, Size);
}

static IoBuffer
MakeFromFileWithOptions(const std::filesystem::path& FileName, uint64_t Offset, uint64_t Size, bool UseShareDelete)
{
	uint64_t FileSize;

#if ZEN_PLATFORM_WINDOWS
	CAtlFile DataFile;

	DWORD ShareOptions = FILE_SHARE_READ;
	if (UseShareDelete)
	{
		ShareOptions |= FILE_SHARE_DELETE;
	}
	HRESULT hRes = DataFile.Create(FileName.c_str(), GENERIC_READ, FILE_SHARE_READ | ShareOptions, OPEN_EXISTING);

	if (FAILED(hRes))
	{
		return {};
	}

	DataFile.GetSize((ULONGLONG&)FileSize);
#else
	int Flags = O_RDONLY;
	if (UseShareDelete)
	{
		Flags |= O_CLOEXEC;
	}
	int Fd = open(FileName.c_str(), Flags);
	if (Fd < 0)
	{
		return {};
	}

	static_assert(sizeof(decltype(stat::st_size)) == sizeof(uint64_t), "fstat() doesn't support large files");
	struct stat Stat;
	fstat(Fd, &Stat);
	FileSize = Stat.st_size;
#endif	// ZEN_PLATFORM_WINDOWS

	// TODO: should validate that offset is in range

	if (Size == ~0ull)
	{
		Size = FileSize - Offset;
	}
	else
	{
		// Clamp size
		if ((Offset + Size) > FileSize)
		{
			Size = FileSize - Offset;
		}
	}

	if (Size)
	{
#if ZEN_PLATFORM_WINDOWS
		void* Fd = DataFile.Detach();
#endif
		return IoBuffer(IoBuffer::File, (void*)uintptr_t(Fd), Offset, Size);
	}

#if !ZEN_PLATFORM_WINDOWS
	close(Fd);
#endif

	// For an empty file, we may as well just return an empty memory IoBuffer
	return IoBuffer(IoBuffer::Wrap, "", 0);
}

IoBuffer
IoBufferBuilder::MakeFromFileWithSharedDelete(const std::filesystem::path& FileName)
{
	return MakeFromFileWithOptions(FileName, 0, ~0ull, true);
}

IoBuffer
IoBufferBuilder::MakeFromFile(const std::filesystem::path& FileName, uint64_t Offset, uint64_t Size)
{
	return MakeFromFileWithOptions(FileName, Offset, Size, false);
}

IoBuffer
IoBufferBuilder::MakeFromTemporaryFile(const std::filesystem::path& FileName)
{
	uint64_t FileSize;
	void*	 Handle;

#if ZEN_PLATFORM_WINDOWS
	CAtlFile DataFile;

	// We need to open with DELETE since this is used for the case
	// when a file has been written to a staging directory, and is going
	// to be moved in place

	HRESULT hRes = DataFile.Create(FileName.native().c_str(), GENERIC_READ | DELETE, FILE_SHARE_READ | FILE_SHARE_DELETE, OPEN_EXISTING);

	if (FAILED(hRes))
	{
		return {};
	}

	DataFile.GetSize((ULONGLONG&)FileSize);

	Handle = DataFile.Detach();
#else
	int Fd	 = open(FileName.native().c_str(), O_RDONLY);
	if (Fd < 0)
	{
		return {};
	}

	static_assert(sizeof(decltype(stat::st_size)) == sizeof(uint64_t), "fstat() doesn't support large files");
	struct stat Stat;
	fstat(Fd, &Stat);
	FileSize = Stat.st_size;

	Handle = (void*)uintptr_t(Fd);
#endif	// ZEN_PLATFORM_WINDOWS

	IoBuffer Iob(IoBuffer::File, Handle, 0, FileSize);
	Iob.m_Core->SetIsWholeFile(true);

	return Iob;
}

IoHash
HashBuffer(IoBuffer& Buffer)
{
	// TODO: handle disk buffers with special path
	return IoHash::HashBuffer(Buffer.Data(), Buffer.Size());
}

//////////////////////////////////////////////////////////////////////////

#if ZEN_WITH_TESTS

void
iobuffer_forcelink()
{
}

TEST_CASE("IoBuffer")
{
	zen::IoBuffer buffer1;
	zen::IoBuffer buffer2(16384);
	zen::IoBuffer buffer3(buffer2, 0, buffer2.Size());
}

#endif

}  // namespace zen