#include <siacommon.h>
#include <ttmath/ttmath.h>
#include <SQLiteCpp/Exception.h>
#include <sqlite3.h>
#include <bitset>
#include <filepath.h>
#include <siadriveconfig.h>
#include <shlguid.h>
#include <Shlobj.h>
#include <base64.h>

#ifdef _WIN32
#include <Wincrypt.h>

// HACK Fix exports
namespace SQLite
{
// Return the result code (if any, otherwise -1).
int Exception::getErrorCode() const noexcept // nothrow
{
	return mErrcode;
}

// Return the extended numeric result code (if any, otherwise -1).
int Exception::getExtendedErrorCode() const noexcept // nothrow
{
	return mExtendedErrcode;
}

// Return a string, solely based on the error code    
const char* Exception::getErrorStr() const noexcept // nothrow
{
	return sqlite3_errstr(mErrcode);
}
}
#endif

NS_BEGIN(Sia)
NS_BEGIN(Api)

__int64 FileSize(const SString& filePath)
{
  struct _stat64 buf;
  if (_wstat64(&filePath[0], &buf) != 0)
    return -1; // error, could use errno to find out more

  return buf.st_size;
}

SString GenerateSha256(const SString& str)
{
#ifdef _WIN32
	SString ret;
	HCRYPTPROV hCryptProv = 0;
	HCRYPTHASH hHash = 0;
	BOOL ok = ::CryptAcquireContext(&hCryptProv, nullptr, nullptr, PROV_RSA_AES, 0);
	ok = ok && ::CryptCreateHash(hCryptProv, CALG_SHA_256, 0, 0, &hHash);
	ok = ok && ::CryptHashData(hHash, reinterpret_cast<const BYTE*>(&str[0]), static_cast<DWORD>(str.ByteLength()), 0);
	if (ok)
	{
		DWORD dwHashLen;
		DWORD dwCount = sizeof(DWORD);
		if (::CryptGetHashParam(hHash, HP_HASHSIZE, reinterpret_cast<BYTE *>(&dwHashLen), &dwCount, 0))
		{
			std::vector<unsigned char> hash(dwHashLen);
			if (::CryptGetHashParam(hHash, HP_HASHVAL, reinterpret_cast<BYTE *>(&hash[0]), &dwHashLen, 0))
			{
				std::ostringstream ss;
				ss << std::hex << std::uppercase << std::setfill('0');
				for (int c : hash)
				{
					ss << std::setw(2) << c;
				}
				ret = ss.str();
			}
		}
	}
	
	if (hHash) ::CryptDestroyHash(hHash);
	if (hCryptProv) ::CryptReleaseContext(hCryptProv, 0);

	return ret;
#else
  a
#endif
}

BOOL RecurDeleteFilesByExtentsion(const SString& folder, const SString& extensionWithDot)
{
#ifdef _WIN32
	BOOL ret = FALSE;
	WIN32_FIND_DATA fd = { 0 };	
	
	SString search;
	search.Resize(MAX_PATH + 1);
	::PathCombine(&search[0], folder.str().c_str(), L"*.*");

	HANDLE find = ::FindFirstFile(search.str().c_str(), &fd);
	if (find != INVALID_HANDLE_VALUE)
	{
		ret = TRUE;
		do
		{
			if (fd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
			{
				if ((SString(fd.cFileName) != ".") && (SString(fd.cFileName) != ".."))
				{
					SString nextFolder;
					nextFolder.Resize(MAX_PATH + 1);
					::PathCombine(&nextFolder[0], folder.str().c_str(), fd.cFileName);
					ret = RecurDeleteFilesByExtentsion(nextFolder, extensionWithDot);
				}
			}
			else
			{
				SString ext = ::PathFindExtension(fd.cFileName);
				if (ext == extensionWithDot)
				{
					SString filePath;
					filePath.Resize(MAX_PATH + 1);
					::PathCombine(&filePath[0], folder.str().c_str(), fd.cFileName);

					ret = RetryDeleteFileIfExists(filePath.str().c_str());
				}
			}
		} while (ret && (::FindNextFile(find, &fd) != 0));
	}

	return ret;
#else
  a
#endif
}

BOOL RetryDeleteFileIfExists(const SString& filePath)
{
#ifdef _WIN32
	BOOL ret = TRUE;
	if (::PathFileExists(filePath.str().c_str()))
	{
		ret = RetryableAction(::DeleteFile(filePath.str().c_str()), DEFAULT_RETRY_COUNT, DEFAULT_RETRY_DELAY_MS);
	}

	return ret;
#else
  a
#endif
}

BOOL RetryAction(std::function<BOOL()> func, std::uint16_t retryCount, const DWORD& retryDelay)
{
#ifdef _WIN32
	BOOL ret = FALSE;
	while (retryCount-- && !((ret = func())))
	{
		::Sleep(retryDelay);
	}

	return ret;
#else
  a
#endif
}

SString BytesToFriendlyDisplay(const SiaCurrency& bytes)
{
  SString units[] = { "B", "KB", "MB", "GB", "TB", "PB" };
  SString readableunit = "B";
  SiaCurrency readablesize = bytes;
  for (const auto& unit : units)
  {
    if (readablesize < 1000)
    {
      readableunit = unit;
      break;
    }
    readablesize /= 1000;
  }
  ttmath::Conv conv;
  conv.scient_from = 256;
  conv.base = 10;
  conv.round = 2;
  return readablesize.ToWString(conv) + ' ' + readableunit;
}

#ifdef _WIN32
std::vector<SString> GetAvailableDrives()
{
  static const std::vector<char> alpha = { 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z' };
  std::bitset<26> drives(~GetLogicalDrives() & 0xFFFFFFFF);

  std::vector<SString> avail;
  for (size_t i = 0; i < alpha.size(); i++)
  {
    if (drives[i] && (alpha[i] != 'A') && (alpha[i] != 'B'))
    {
      avail.push_back(alpha[i]);
    }
  }

  return std::move(avail);
}

std::int32_t GetRegistry(const SString& name, const std::int32_t& defaultValue, const bool& user)
{
  std::int32_t ret = defaultValue;

  DWORD dataLen = 0;
  if (::RegGetValue((user ? HKEY_CURRENT_USER : HKEY_LOCAL_MACHINE), L"Software\\SiaExtensions\\SiaDrive", &name[0], RRF_RT_REG_SZ, nullptr, nullptr, &dataLen) == ERROR_SUCCESS)
  {
    SString data;
    data.Resize(dataLen + 1);
    if (::RegGetValue((user ? HKEY_CURRENT_USER : HKEY_LOCAL_MACHINE), L"Software\\SiaExtensions\\SiaDrive", &name[0], RRF_RT_REG_SZ, nullptr, &data[0], &dataLen) == ERROR_SUCCESS)
    {
      ret = SString::ToInt32(data);
    }
  }

  return ret;
}

SString GetRegistry(const SString& name, const SString& defaultValue, const bool& user)
{
  SString ret = defaultValue;

  DWORD dataLen = 0;
  if (::RegGetValue((user ? HKEY_CURRENT_USER : HKEY_LOCAL_MACHINE), L"Software\\SiaExtensions\\SiaDrive", &name[0], RRF_RT_REG_SZ, nullptr, nullptr, &dataLen) == ERROR_SUCCESS)
  {
    SString data;
    data.Resize(dataLen + 1);
    if (::RegGetValue((user ? HKEY_CURRENT_USER : HKEY_LOCAL_MACHINE), L"Software\\SiaExtensions\\SiaDrive", &name[0], RRF_RT_REG_SZ, nullptr, &data[0], &dataLen) == ERROR_SUCCESS)
    {
      ret = data;
    }
  }

  return ret;
}

void SetRegistry(const SString& name, const SString& value)
{
  ::RegSetKeyValue(HKEY_CURRENT_USER, L"Software\\SiaExtensions\\SiaDrive", &name[0], REG_SZ, &value[0], value.ByteLength() + sizeof(SString::SChar));
}

bool ExecuteProcess(CSiaDriveConfig* siaDriveConfig, FilePath process, FilePath workingDir, const bool& waitForExit)
{
  process.MakeAbsolute();
  workingDir.MakeAbsolute();

  SString command = process;
  if (static_cast<SString>(process).EndsWith("siad"))
  {
    SString apiAdr = siaDriveConfig->GetHostNameOrIp() + ":" + SString::FromUInt32(siaDriveConfig->GetApiPort());
    SString hostAdr = ":" + SString::FromUInt32(siaDriveConfig->GetHostPort());
    SString rpcAddr = ":" + SString::FromUInt32(siaDriveConfig->GetRpcPort());
    command = "cmd.exe /c \"" + static_cast<SString>(process) + "\" --api-addr=" + apiAdr + " --host-addr=" + hostAdr + " --rpc-addr=" + rpcAddr;
  }
  STARTUPINFO si = { 0 };
  si.cb = sizeof(si);

  PROCESS_INFORMATION pi;
  if (::CreateProcess(nullptr, &command[0], nullptr, nullptr, FALSE, 0, nullptr, &workingDir[0], &si, &pi))
  {
    return true;
  }

  return false;
}

HRESULT CreateShortcut(const SString& execPath, const SString& description, const SString& shortcutPath, const bool& minimized) 
{
  ComInitWrapper cw;

  IShellLink* shellLink(nullptr);
  HRESULT hr = ::CoCreateInstance(CLSID_ShellLink, nullptr, CLSCTX_INPROC_SERVER, IID_IShellLink, reinterpret_cast<LPVOID*>(&shellLink));
  if (SUCCEEDED(hr))
  {
    shellLink->SetPath(&execPath[0]);
    shellLink->SetWorkingDirectory(&FilePath(execPath).RemoveFileName()[0]);
    shellLink->SetDescription(&description[0]);
    shellLink->SetShowCmd(minimized ? SW_MINIMIZE : SW_SHOWDEFAULT);

    IPersistFile* persistFile(nullptr);
    hr = shellLink->QueryInterface(IID_IPersistFile, reinterpret_cast<LPVOID*>(&persistFile));
    if (SUCCEEDED(hr))
    {
      hr = persistFile->Save(&shortcutPath[0], TRUE);

      persistFile->Release();
    }
    shellLink->Release();
  }

  return hr;
}

SString SecureEncryptString(const SString& name, const SString& data)
{
  DATA_BLOB inputData = {0};
  inputData.pbData = (BYTE*)data.str().c_str();
  inputData.cbData = data.ByteLength();

  DATA_BLOB outputData = {0};
  if (::CryptProtectData(&inputData, &name[0], nullptr, nullptr, nullptr, 0, &outputData))
  {
    std::vector<BYTE> encoded;
    Base64::Encode(outputData.pbData, outputData.cbData, encoded);
    return std::string(reinterpret_cast<const char*>(&encoded[0]), encoded.size());
  }

  return "";
}

SString SecureDecryptString(const SString& name, const SString& data)
{
  SString ret;
  std::string d = SString::ToUtf8(data);
  std::vector<BYTE> decoded;
  Base64::Decode((BYTE*)&d[0], d.length(), decoded);

  DATA_BLOB inputData = {0};
  inputData.pbData = &decoded[0];
  inputData.cbData = decoded.size();

  DATA_BLOB outputData = {0};
  LPWSTR n = nullptr;
  if (::CryptUnprotectData(&inputData, &n, nullptr, nullptr, nullptr, 0, &outputData))
  {
    if (name == n)
    {
      ret = SString::String(reinterpret_cast<const SString::SChar*>(outputData.pbData), outputData.cbData / sizeof(SString::SChar));
    }
  }

  return ret;
}

class SIWrapper
{
public:
  SIWrapper(HANDLE& mtx) :
    _mtx(mtx)
  {
  }

  ~SIWrapper()
  {
    ::CloseHandle(_mtx);
  }

private:
  const HANDLE _mtx;
};

static std::unordered_map<SString, std::unique_ptr<SIWrapper>> g_inst;

bool CheckSingleInstance(const SString& name)
{
  bool ret = false;
  auto mtx = ::CreateMutex(nullptr, FALSE, &name[0]);
  if (::WaitForSingleObject(mtx, 5000) == WAIT_OBJECT_0)
  {
    g_inst.insert({ name, std::make_unique<SIWrapper>(mtx) });
    ret = true;
  }
  else
  {
    ::CloseHandle(mtx);
  }

  return ret;
}

ComInitWrapper::ComInitWrapper() :
  _uninit(SUCCEEDED(::CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED)))
{
  
}

ComInitWrapper::~ComInitWrapper()
{
  if (_uninit)
  {
    ::CoUninitialize();
  }
}

#endif

NS_END(2)