/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#pragma once

#include <array>
#include <string>
#include <sstream>
#include <regex>
#include <iomanip>
#include <sstream>
#if __has_include("fmt/format.h")
#include <fmt/format.h>
#elif __has_include("spdlog/fmt/fmt.h")
#include <spdlog/fmt/fmt.h>
#else
#warning "Failed to find fmt include"
#define OVKMBS_FMT_MISSING
#endif
#include <magic_enum.hpp>

namespace OpenVulkano
{
	namespace ByteSizeUnitHelper
	{
		constexpr std::array<uint64_t, 13> BuildFactors()
		{
			std::array<uint64_t, 13> factors = { 1 };
			for(uint64_t j = 1, i = 10; i < 70; i += 10, j++)
			{
				factors[j] = 1uLL << i;
			}
			for(uint64_t j = 7, i = 1000; j < 13; i *= 1000, j++)
			{
				factors[j] = i;
			}
			return factors;
		}

		constexpr std::array<double, 13> BuildDivisors()
		{
			std::array<uint64_t, 13> factors = BuildFactors();
			std::array<double, 13> divs{{1}};
			for (size_t i = 1; i < divs.size(); i++)
			{
				divs[i] = 1. / static_cast<double>(factors[i]);
			}
			return divs;
		}
	}

	class ByteSizeUnit final
	{
		static constexpr std::array<uint64_t, 13> FACTORS = ByteSizeUnitHelper::BuildFactors();
		static constexpr std::array<double, 13> DIVISORS = ByteSizeUnitHelper::BuildDivisors();
	public:
		enum Unit : int { B = 0, kiB, MiB, GiB, TiB, PiB, EiB, kB, MB, GB, TB, PB, EB };

		constexpr ByteSizeUnit(Unit unit) : unit(unit) {}

		[[nodiscard]] constexpr uint64_t GetFactor() const { return FACTORS[unit]; }

		[[nodiscard]] constexpr double GetDivisor() const { return DIVISORS[unit]; }

		[[nodiscard]] constexpr std::string_view GetName() const
		{
			return magic_enum::enum_name(unit);
		}

		[[nodiscard]] static constexpr ByteSizeUnit FromName(std::string_view name)
		{
			auto result = magic_enum::enum_cast<Unit>(name);
			if (result) return { result.value() };
			return B;
		}

		[[nodiscard]] static constexpr ByteSizeUnit GetClosestUnit(uint64_t size, bool si = false)
		{
			for (int i = (si ? EB : EiB); i >= (si ? kB : kiB); i--)
			{
				if (size > FACTORS[i]) return { static_cast<Unit>(i) };
			}
			return { B };
		}

	private:
		Unit unit;
	};

	class ByteSize final
	{
		static inline const std::regex REGEX_PATTERN = std::regex(R"(\s*(\d+(\.\d*)?|\.\d+)\s*([kMGTPE]i?B)\s*)");

		uint64_t bytes;

	public:
		constexpr ByteSize(uint64_t size = 0) : bytes(size) {}

		template<typename T>
		constexpr ByteSize(T size, ByteSizeUnit unit)
			: bytes(size * unit.GetFactor()) {}

		ByteSize(const std::string& str)
		{
			std::smatch match;
			if (!std::regex_match(str, match, REGEX_PATTERN)) throw std::runtime_error("Unable to parse '" + str + "' as size.");
			double value = std::stod(match[1].str());
			uint64_t factor = ByteSizeUnit::FromName(match[3].str()).GetFactor();
			bytes = value * factor;
		}

		std::ostream& operator<<(std::ostream& os) const { return FormatStream(os); }

		std::ostream& FormatStream(std::ostream& oss, bool si = false) const
		{
			oss << std::setprecision(3);
			auto unit = ByteSizeUnit::GetClosestUnit(bytes, si);
			oss << static_cast<double>(bytes) * unit.GetDivisor() << ' ' << unit.GetName();
			return oss;
		}

		[[nodiscard]] std::string Format(bool si = false) const
		{
			std::ostringstream oss;
			FormatStream(oss, si);
			return oss.str();
		}

		constexpr operator uint64_t() const { return bytes; }

		operator std::string() const { return Format(); }
		
		ByteSize operator +(const ByteSize other) const { return bytes + other.bytes; }
		ByteSize operator -(const ByteSize other) const { return bytes - other.bytes; }
		ByteSize operator +(const size_t other) const { return bytes + other; }
		ByteSize operator -(const size_t other) const { return bytes - other; }
		ByteSize& operator =(const size_t other) { bytes = other; return *this; }
		ByteSize& operator =(const ByteSize other) { bytes = other.bytes; return *this; }
		ByteSize& operator +=(const size_t other) { bytes += other; return *this; }
		ByteSize& operator +=(const ByteSize other) { bytes += other.bytes; return *this; }
		ByteSize& operator -=(const size_t other) { bytes -= other; return *this; }
		ByteSize& operator -=(const ByteSize other) { bytes -= other.bytes; return *this; }
	};

	inline constexpr ByteSize operator"" _kiB(long double num) { return { num, ByteSizeUnit::kiB }; }
	inline constexpr ByteSize operator"" _MiB(long double num) { return { num, ByteSizeUnit::MiB }; }
	inline constexpr ByteSize operator"" _GiB(long double num) { return { num, ByteSizeUnit::GiB }; }
	inline constexpr ByteSize operator"" _TiB(long double num) { return { num, ByteSizeUnit::TiB }; }
	inline constexpr ByteSize operator"" _PiB(long double num) { return { num, ByteSizeUnit::PiB }; }
	inline constexpr ByteSize operator"" _EiB(long double num) { return { num, ByteSizeUnit::EiB }; }
	inline constexpr ByteSize operator""  _kB(long double num) { return { num, ByteSizeUnit::kB }; }
	inline constexpr ByteSize operator""  _MB(long double num) { return { num, ByteSizeUnit::MB }; }
	inline constexpr ByteSize operator""  _GB(long double num) { return { num, ByteSizeUnit::GB }; }
	inline constexpr ByteSize operator""  _TB(long double num) { return { num, ByteSizeUnit::TB }; }
	inline constexpr ByteSize operator""  _PB(long double num) { return { num, ByteSizeUnit::PB }; }
	inline constexpr ByteSize operator""  _EB(long double num) { return { num, ByteSizeUnit::EB }; }
	inline constexpr ByteSize operator""   _B(unsigned long long int num) { return { num }; }
	inline constexpr ByteSize operator"" _kiB(unsigned long long int num) { return { num << 10 }; }
	inline constexpr ByteSize operator"" _MiB(unsigned long long int num) { return { num << 20 }; }
	inline constexpr ByteSize operator"" _GiB(unsigned long long int num) { return { num << 30 }; }
	inline constexpr ByteSize operator"" _TiB(unsigned long long int num) { return { num << 40 }; }
	inline constexpr ByteSize operator"" _PiB(unsigned long long int num) { return { num << 50 }; }
	inline constexpr ByteSize operator"" _EiB(unsigned long long int num) { return { num << 60 }; }
	inline constexpr ByteSize operator""  _kB(unsigned long long int num) { return { num, ByteSizeUnit::kB }; }
	inline constexpr ByteSize operator""  _MB(unsigned long long int num) { return { num, ByteSizeUnit::MB }; }
	inline constexpr ByteSize operator""  _GB(unsigned long long int num) { return { num, ByteSizeUnit::GB }; }
	inline constexpr ByteSize operator""  _TB(unsigned long long int num) { return { num, ByteSizeUnit::TB }; }
	inline constexpr ByteSize operator""  _PB(unsigned long long int num) { return { num, ByteSizeUnit::PB }; }
	inline constexpr ByteSize operator""  _EB(unsigned long long int num) { return { num, ByteSizeUnit::EB }; }
}

#ifndef OVKMBS_FMT_MISSING
template<> struct fmt::formatter<OpenVulkano::ByteSize>
{
	template<typename ParseContext> constexpr auto parse(ParseContext& ctx)
	{
		return ctx.begin();
	}

	template<typename FormatContext> auto format(const OpenVulkano::ByteSize& bs, FormatContext& ctx)
	{
		return fmt::format_to(ctx.out(), "{}", bs.Format());
	}
};
#endif
