#include "stdafx.h"
#include "SiaDokanDrive.h"
#include <filesystem>
#include "UploadManager.h"

using namespace Sia::Api;
using namespace Sia::Api::Dokan;

// The general idea is that normal file I/O occurs in a local cache folder and once the file is closed, it is scheduled for upload into Sia.
//	Files requested to be openned that are not cached will be downloaded first. If the file is not found in Sia, it will be treated as new.
//	Keeping cache and Sia in synch will be a bit of a hastle, so it's strongly suggested to treat the cache folder as if it doesn't exist;
//	however, simply deleting files in the cache folder should not be an issue as long as the drive is not mounted.
class DokanImpl
{
private:
	typedef struct
	{
		String SiaPath;
		String CacheFilePath;
		bool ReadOnly;
	} OpenFileInfo;

private:
	static std::mutex _dokanMutex;
	static CSiaApi* _siaApi;
	static CSiaDriveConfig* _siaDriveConfig;
	static std::unique_ptr<CUploadManager> _uploadManager;
	static DOKAN_OPERATIONS _dokanOps;
	static DOKAN_OPTIONS _dokanOptions;
	static String _cacheLocation;
	static std::unordered_map<ULONG64, OpenFileInfo> _openFileMap;
	static std::unique_ptr<std::thread> _fileListThread;
	static bool _fileListStopRequested;
	static CSiaFileTreePtr _siaFileTree;
	static std::mutex _fileTreeMutex;
	static std::unique_ptr<std::thread> _mountThread;
	static NTSTATUS _mountStatus;
	static String _mountPoint;

private:
	inline static const String& GetCacheLocation()
	{
		return _cacheLocation;
	}

	static bool AddFileToCache(const String& siaPath, const String& cacheLocation)
	{
		bool ret = false;
		std::wstring tempPath;
		tempPath.resize(MAX_PATH + 1);
		if (GetTempPath(MAX_PATH + 1, &tempPath[0]))
		{
			ret = ApiSuccess(_siaApi->GetRenter()->DownloadFile(siaPath, tempPath));
			if (ret)
			{

			}
		}

		return ret;
	}

	static void QueueUploadIfChanged(const ULONG64& id, const std::uint64_t& size)
	{
		if (!_openFileMap[id].ReadOnly)
		{
			if (size > 0)
			{
				// TODO Always save for now - need to change to detect modifications
				// TODO Handle error return
				_uploadManager->AddOrUpdate(_openFileMap[id].SiaPath, _openFileMap[id].CacheFilePath);
			}
			else
			{
				// Treat 0 length files as deleted in Sia - cache retains 0-length
				// TODO Retain 0 length in cache?
				// TODO Handle error return
				_uploadManager->Remove(_openFileMap[id].SiaPath);
			}
		}
	}

	static void StartFileListThread()
	{
		if (!_fileListThread)
		{
			_fileListStopRequested = false;
			_fileListThread.reset(new std::thread([]()
			{
				while (!_fileListStopRequested)
				{
					CSiaFileTreePtr siaFileTree;
					_siaApi->GetRenter()->GetFileTree(siaFileTree);
					{
						std::lock_guard<std::mutex> l(_fileTreeMutex);
						_siaFileTree = siaFileTree;
					}
					
					if (!_fileListStopRequested)
					{
						// TODO Change to WaitForSingleObject() for immediate termination
						Sleep(5000);
					}
				}
			}));
		}
	}

	static void StopFileListThread()
	{
		if (_fileListThread)
		{
			_fileListStopRequested = true;
			_fileListThread->join();
			_fileListThread.reset(nullptr);
		}
	}

	// Dokan callbacks
private:
	static NTSTATUS DOKAN_CALLBACK SiaDrive_ZwCreateFile(
		LPCWSTR FileName,
		PDOKAN_IO_SECURITY_CONTEXT SecurityContext,
		ACCESS_MASK DesiredAccess,
		ULONG FileAttributes,
		ULONG ShareAccess,
		ULONG CreateDisposition,
		ULONG CreateOptions,
		PDOKAN_FILE_INFO DokanFileInfo)
	{
		std::lock_guard<std::mutex> l(_dokanMutex);
		SECURITY_ATTRIBUTES securityAttrib;
		securityAttrib.nLength = sizeof(securityAttrib);
		securityAttrib.lpSecurityDescriptor = SecurityContext->AccessState.SecurityDescriptor;
		securityAttrib.bInheritHandle = FALSE;

		DWORD fileAttributesAndFlags;
		DWORD creationDisposition;
		DokanMapKernelToUserCreateFileFlags(FileAttributes, CreateOptions, CreateDisposition, &fileAttributesAndFlags, &creationDisposition);
		
		ACCESS_MASK genericDesiredAccess = DokanMapStandardToGenericAccess(DesiredAccess);

		NTSTATUS ret = STATUS_SUCCESS;
		// Probably not going to happen, but just in case
		if (PathIsUNC(FileName))
		{
			ret = STATUS_ILLEGAL_ELEMENT_ADDRESS;
		}
		else
		{
			bool isFile = (FileAttributes & FILE_NON_DIRECTORY_FILE) ? true : false;
			DokanFileInfo->IsDirectory = !isFile;
			if (isFile)
			{
				// Formulate Sia path and cache path
				String siaPath = CSiaApi::FormatToSiaPath(PathSkipRoot(FileName)); // Strip drive letter to get Sia path
				if (siaPath.length())
				{
					String cacheFilePath;
					cacheFilePath.resize(MAX_PATH + 1);
					PathCombine(&cacheFilePath[0], GetCacheLocation().c_str(), siaPath.c_str());

					// If cache file already exists and is a directory, requested file operation isn't valid
					if (GetFileAttributes(cacheFilePath.c_str()) & FILE_ATTRIBUTE_DIRECTORY)
					{
						ret = STATUS_OBJECT_NAME_COLLISION;
					}
					else
					{
						bool exists;
						if (ApiSuccess(_siaApi->GetRenter()->FileExists(siaPath, exists)))
						{
							// Operations on existing files that are requested to be truncated, overwritten or re-created
							//	will first be deleted and then replaced if, after the file operation is done, the resulting file
							//	size is > 0. Sia doesn't support random access to files (upload/download/rename/delete).
							bool isCreateOp = false;
							bool isReplaceOp = false;
							switch (creationDisposition)
							{
							case CREATE_ALWAYS:
							{
								isCreateOp = true;
								isReplaceOp = exists;
							}
							break;

							case CREATE_NEW:
							{
								if (exists)
								{
									ret = STATUS_OBJECT_NAME_EXISTS;
								}
								else
								{
									isCreateOp = true;
								}
							}
							break;

							case OPEN_ALWAYS:
							{
								if (!exists)
								{
									isCreateOp = true;
								}
							}
							break;

							case OPEN_EXISTING:
							{
								if (!exists)
								{
									ret = STATUS_NOT_FOUND;
								}
							}
							break;

							case TRUNCATE_EXISTING:
							{
								if (exists)
								{
									isCreateOp = isReplaceOp = true;
								}
								else
								{
									ret = STATUS_NOT_FOUND;
								}
							}
							break;
							}

							if (ret == STATUS_SUCCESS)
							{
								if (isReplaceOp)
								{
									// Since this is a request to replace an existing file, make sure cache is deleted first.
									//	If file isn't cached, delete from Sia only
									if (!PathFileExists(cacheFilePath.c_str()) || ::DeleteFile(cacheFilePath.c_str()))
									{
										if (!ApiSuccess(_uploadManager->Remove(siaPath)))
										{
											ret = STATUS_INVALID_SERVER_STATE;
										}
									}
									else
									{
										ret = DokanNtStatusFromWin32(GetLastError());
									}
								}

								if (ret == STATUS_SUCCESS)
								{
									// If file must exist, then check for it in cache location. If not found,
									//	it must be downloaded first and placed in cache
									if (!isCreateOp && !PathFileExists(cacheFilePath.c_str()))
									{
										if (!AddFileToCache(siaPath, cacheFilePath))
										{
											ret = STATUS_INVALID_SERVER_STATE;
										}
									}

									if (ret == STATUS_SUCCESS)
									{
										// Create file as specified
										HANDLE handle = CreateFile(
											cacheFilePath.c_str(),
											genericDesiredAccess,
											ShareAccess,
											&securityAttrib,
											creationDisposition,
											fileAttributesAndFlags,
											nullptr);
										if (handle == INVALID_HANDLE_VALUE)
										{
											ret = DokanNtStatusFromWin32(GetLastError());
										}
										else
										{
											DokanFileInfo->Context = reinterpret_cast<ULONG64>(handle); // save the file handle in Context
											if (isFile)
											{
												OpenFileInfo ofi;
												ofi.SiaPath = siaPath;
												ofi.CacheFilePath = cacheFilePath;
												// TODO Detect if file is read-only
												// TODO Quick hash to detect changes
												ofi.ReadOnly = false;
												_openFileMap.insert({ DokanFileInfo->Context, ofi });
											}

											/*if (creationDisposition == OPEN_ALWAYS ||
												  creationDisposition == CREATE_ALWAYS) {
												error = GetLastError();
												if (error == ERROR_ALREADY_EXISTS) {
												  DbgPrint(L"\tOpen an already existing file\n");
												  // Open succeed but we need to inform the driver
												  // that the file open and not created by returning STATUS_OBJECT_NAME_COLLISION
												  return STATUS_OBJECT_NAME_COLLISION;
												}
											  }*/
										}
									}
								}
							}
						}
						else
						{
							ret = STATUS_INVALID_SERVER_STATE;
						}
					}
				}
				else
				{
					ret = STATUS_OBJECT_NAME_INVALID;
				}
			}
			else // Folder Operation (cache operation only)
			{
				ret = STATUS_NOT_IMPLEMENTED;
			}
		}
		
		return ret;
	}

	static NTSTATUS DOKAN_CALLBACK SiaDrive_FindFiles(
		LPCWSTR FileName,
		PFillFindData FillFindData,
		PDOKAN_FILE_INFO DokanFileInfo)
	{
		std::lock_guard<std::mutex> l(_dokanMutex);
		auto siaFileTree = _siaFileTree;
		if (siaFileTree)
		{
			String siaQuery = CSiaApi::FormatToSiaPath(PathSkipRoot(FileName));
			auto fileList = siaFileTree->Query(siaQuery);
			for (auto& file : fileList)
			{
				WIN32_FIND_DATA fd = { 0 };
				wcscpy_s(fd.cFileName, PathFindFileName(file->GetSiaPath().c_str()));
				
				FillFindData(&fd, DokanFileInfo);
			}
		}

		return STATUS_SUCCESS;
	}

	static void DOKAN_CALLBACK Sia_CloseFile(LPCWSTR FileName, PDOKAN_FILE_INFO DokanFileInfo)
	{
		std::lock_guard<std::mutex> l(_dokanMutex);
		ULONG64 id = DokanFileInfo->Context;
		if (id)
		{
			HANDLE handle = reinterpret_cast<HANDLE>(DokanFileInfo->Context);
			// Ignore directories in Sia - reside in cache only.
			if (DokanFileInfo->IsDirectory)
			{
				::CloseHandle(handle);
			}
			else
			{
				LARGE_INTEGER li = { 0 };
				BOOL sizeOk = GetFileSizeEx(handle, &li);

				::CloseHandle(handle);

				// TODO If it's not ok, why and what to do?
				if (sizeOk)
				{
					QueueUploadIfChanged(id, li.QuadPart);
				}
				_openFileMap.erase(id);
			}

			DokanFileInfo->Context = 0;
		}
	}

	static NTSTATUS DOKAN_CALLBACK Sia_Mounted(PDOKAN_FILE_INFO DokanFileInfo)
	{
		std::lock_guard<std::mutex> l(_dokanMutex);
		StartFileListThread();
		return STATUS_SUCCESS;
	}

	static NTSTATUS DOKAN_CALLBACK Sia_Unmounted(PDOKAN_FILE_INFO DokanFileInfo)
	{
		std::lock_guard<std::mutex> l(_dokanMutex);
		StopFileListThread();
		return STATUS_SUCCESS;
	}

public:
	static void Initialize(CSiaApi* siaApi, CSiaDriveConfig* siaDriveConfig)
	{
		_siaApi = siaApi;
		_siaDriveConfig = siaDriveConfig;
		_uploadManager.reset(new CUploadManager(CSiaCurl(siaApi->GetHostConfig()), siaDriveConfig));
		_dokanOps.Cleanup = nullptr;
		_dokanOps.CloseFile = Sia_CloseFile;
		_dokanOps.DeleteDirectory = nullptr;
		_dokanOps.DeleteFileW = nullptr;
		_dokanOps.FindFiles = SiaDrive_FindFiles;
		_dokanOps.FindFilesWithPattern = nullptr;
		_dokanOps.FindStreams = nullptr;
		_dokanOps.FlushFileBuffers = nullptr;
		_dokanOps.GetDiskFreeSpaceW = nullptr;
		_dokanOps.GetFileInformation = nullptr;
		_dokanOps.GetFileSecurityW = nullptr;
		_dokanOps.GetVolumeInformationW = nullptr;
		_dokanOps.LockFile = nullptr;
		_dokanOps.Mounted = Sia_Mounted;
		_dokanOps.MoveFileW = nullptr;
		_dokanOps.ReadFile = nullptr;
		_dokanOps.SetAllocationSize = nullptr;
		_dokanOps.SetEndOfFile = nullptr;
		_dokanOps.SetFileAttributesW = nullptr;
		_dokanOps.SetFileSecurityW = nullptr;
		_dokanOps.SetFileTime = nullptr;
		_dokanOps.UnlockFile = nullptr;
		_dokanOps.Unmounted = Sia_Unmounted;
		_dokanOps.WriteFile = nullptr;
		_dokanOps.ZwCreateFile = SiaDrive_ZwCreateFile;

		ZeroMemory(&_dokanOptions, sizeof(DOKAN_OPTIONS));
		_dokanOptions.Version = DOKAN_VERSION;
		_dokanOptions.ThreadCount = 0; // use default
		_dokanOptions.Options |= DOKAN_OPTION_CURRENT_SESSION;
	}

	static void Mount(const wchar_t& driveLetter, const String& cacheLocation)
	{
		if (_siaApi && !_mountThread)
		{
			_cacheLocation = cacheLocation;
			wchar_t tmp[] = { driveLetter, ':', '\\', 0 };
			_mountPoint = tmp;
			_mountThread.reset(new std::thread([&]()
			{
				_dokanOptions.MountPoint = _mountPoint.c_str();
				_mountStatus = DokanMain(&_dokanOptions, &_dokanOps);
			}));
		}
	}

	static void Unmount()
	{
		if (_mountThread)
		{
			DokanRemoveMountPoint(_mountPoint.c_str());
			_mountThread->join();
			_mountThread.reset(nullptr);
			_mountPoint.clear();
		}
	}

	static void Shutdown()
	{
		Unmount();
		_uploadManager.reset(nullptr);
		_siaApi = nullptr;
		_siaDriveConfig = nullptr;
		ZeroMemory(&_dokanOps, sizeof(_dokanOps));
		ZeroMemory(&_dokanOptions, sizeof(_dokanOptions));
	}

	static std::mutex& GetMutex()
	{
		return _dokanMutex;
	}

	static bool IsInitialized()
	{
		return _siaApi != nullptr;
	}
};
// Static member variables
std::mutex DokanImpl::_dokanMutex;
CSiaApi* DokanImpl::_siaApi = nullptr;
CSiaDriveConfig* DokanImpl::_siaDriveConfig = nullptr;
std::unique_ptr<CUploadManager> DokanImpl::_uploadManager;
DOKAN_OPERATIONS DokanImpl::_dokanOps;
DOKAN_OPTIONS DokanImpl::_dokanOptions;
String DokanImpl::_cacheLocation;
bool DokanImpl::_fileListStopRequested;
CSiaFileTreePtr DokanImpl::_siaFileTree;
std::mutex DokanImpl::_fileTreeMutex;
std::unique_ptr<std::thread> DokanImpl::_fileListThread;
std::unordered_map<ULONG64, DokanImpl::OpenFileInfo> DokanImpl::_openFileMap;
std::unique_ptr<std::thread> DokanImpl::_mountThread;
NTSTATUS DokanImpl::_mountStatus = STATUS_SUCCESS;
String DokanImpl::_mountPoint;


CSiaDokanDrive::CSiaDokanDrive(CSiaApi& siaApi, CSiaDriveConfig* siaDriveConfig) :
	_siaApi(siaApi),
	_siaDriveConfig(siaDriveConfig),
	_Mounted(false)
{
	std::lock_guard<std::mutex> l(DokanImpl::GetMutex());
	if (DokanImpl::IsInitialized())
		throw SiaDokanDriveException("Sia drive has already been activated");
	
	DokanImpl::Initialize(&_siaApi, _siaDriveConfig);
}

CSiaDokanDrive::~CSiaDokanDrive()
{
	std::lock_guard<std::mutex> l(DokanImpl::GetMutex());
	Unmount();
	DokanImpl::Shutdown();
}

void CSiaDokanDrive::Mount(const wchar_t& driveLetter, const String& cacheLocation, const std::uint64_t& maxCacheSizeBytes)
{
	std::lock_guard<std::mutex> l(DokanImpl::GetMutex());
	DokanImpl::Mount(driveLetter, cacheLocation);
}

void CSiaDokanDrive::Unmount(const bool& clearCache)
{
	std::lock_guard<std::mutex> l(DokanImpl::GetMutex());
	// TODO When files are open, need to prompt and prevent shutdown?
	DokanImpl::Unmount();
}

void CSiaDokanDrive::ClearCache()
{
	std::lock_guard<std::mutex> l(DokanImpl::GetMutex());
}