#pragma once

#include "int.hpp"
#include "storage.hpp"
#include "utils.hpp"

namespace mp
{

template <size_t SizeA, size_t SizeB>
constexpr size_t ResultMaxSize = std::max(SizeA, SizeB);

template <AnyConstMpInt TL, AnyConstInt TR>
struct OpResult;

template <AnyConstMpInt TLhs, AnyConstMpInt TRhs>
    requires ElementTypesMatch<TLhs, TRhs>
struct OpResult<TLhs, TRhs>

{
    using type = BasicInt<typename TLhs::ElementType, ResultMaxSize<TLhs::MAX_BYTES, TRhs::MAX_BYTES>>;
};

template <AnyConstMpInt TLhs, AnyRegularInt TRhs>
struct OpResult<TLhs, TRhs>
{
    //using type = BasicInt<typename TLhs::ElementType, ResultMaxSize<TLhs::MAX_BYTES, sizeof(TRhs)>>;
    using type = TLhs; // keep the same size as TLhs when operating with regular ints
};

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
using OpResultType = typename OpResult<TLhs, TRhs>::type;

template <AnyConstMpInt T>
class OverflowError : public std::runtime_error
{
  public:
    OverflowError(T&& value) : std::runtime_error("Overflow"), m_value(std::move(value)) {}

    const T& value() const { return m_value; }

  private:
    T m_value;
};

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
using OverflowErrorOf = OverflowError<OpResultType<TLhs, TRhs>>;

// Regular int to mp int
template <ElementSuitable TElem, AnyRegularInt TR>
BasicInt<TElem, sizeof(TR)> to_mp_int(TR value)
{
    BasicInt<TElem, sizeof(TR)> res;
    if constexpr (sizeof(TR) <= sizeof(TElem))
    {
        res.set(0, static_cast<TElem>(value));
    }
    else
    {
        constexpr TR ELEM_MASK = static_cast<TR>(std::numeric_limits<TElem>::max());
        size_t idx = 0;
        while (value != 0)
        {
            res.set(idx, static_cast<TElem>(value & ELEM_MASK));
            value >>= (sizeof(TElem) * 8);
            idx++;
        }
    }

    return res;
}

// no-op
template <ElementSuitable TElem, AnyConstMpInt TR>
    requires std::is_same_v<typename TR::ElementType, TElem>
const TR& to_mp_int(const TR& value)
{
    return value;
}

template <AnyConstMpInt TLhs, AnyConstMpInt TRhs>
    requires ElementTypesMatch<TLhs, TRhs>
auto operator<=>(const TLhs& lhs, const TRhs& rhs)
{
    //auto size_comp = lhs.size_elems() <=> rhs.size_elems();
    //if (size_comp != std::strong_ordering::equal)
    //{
    //    return size_comp;
    //}

    //for (size_t i = lhs.size_elems(); i-- > 0;)
    //{
    //    auto elem_comp = lhs.get(i) <=> rhs.get(i);
    //    if (elem_comp != std::strong_ordering::equal)
    //    {
    //        return elem_comp;
    //    }
    //}

    size_t max_size = std::max(lhs.size_elems(), rhs.size_elems());
    for (size_t i = max_size; i-- > 0;)
    {
        auto elem_comp = lhs.get(i) <=> rhs.get(i);
        if (elem_comp != std::strong_ordering::equal)
        {
            return elem_comp;
        }
    }

    return std::strong_ordering::equal;
}

template <AnyConstMpInt TLhs, AnyRegularInt TRhs>
auto operator<=>(const TLhs& lhs, TRhs rhs)
{
    return lhs <=> to_mp_int<typename TLhs::ElementType>(rhs);
}

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
bool operator==(const TLhs& lhs, const TRhs& rhs)
{
    return (lhs <=> rhs) == std::strong_ordering::equal;
}

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
bool operator!=(const TLhs& lhs, const TRhs& rhs)
{
    return !(lhs == rhs);
}

struct Addition
{
    template <AnyConstMpInt TLhs, AnyConstMpInt TRhs, AnyMpInt TRes>
        requires ElementTypesMatch<TLhs, TRes> && ElementTypesMatch<TRhs, TRes>
    static void invoke(const TLhs& lhs, const TRhs& rhs, TRes& res)
    {
        using ElementType = typename TRes::ElementType;

        res.zero();

        ElementType carry = 0;
        size_t end = std::max(lhs.size_elems(), rhs.size_elems());

        for (size_t i = 0; i < end; ++i)
        {
            ElementType a = lhs.get(i);
            ElementType b = rhs.get(i);

            ElementType c = carry + a;
            carry = (c < a) ? 1 : 0;
            c += b;
            if (c < b)
                carry = 1;

            res.set(i, c);
        }

        res.set(end, carry);
    }
};

struct Subtraction
{
    template <AnyConstMpInt TLhs, AnyConstMpInt TRhs, AnyMpInt TRes>
        requires ElementTypesMatch<TLhs, TRes> && ElementTypesMatch<TRhs, TRes>
    static void invoke(const TLhs& lhs, const TRhs& rhs, TRes& res)
    {
        using ElementType = typename TRes::ElementType;

        res.zero();

        ElementType borrow = 0;
        size_t end = std::max(lhs.size_elems(), rhs.size_elems());

        for (size_t i = 0; i < end; ++i)
        {
            ElementType a = lhs.get(i);
            ElementType b = rhs.get(i);

            ElementType c = a - b - borrow;
            borrow = (a < b) || (borrow && a == b);
            res.set(i, c);
        }

        if (borrow != 0)
        {
            throw std::overflow_error("Subtraction underflow");
        }
    }
};

struct Multiplication
{
    template <AnyConstMpInt TLhs, AnyConstMpInt TRhs, AnyMpInt TRes>
        requires ElementTypesMatch<TLhs, TRes> && ElementTypesMatch<TRhs, TRes>
    static void invoke(const TLhs& lhs, const TRhs& rhs, TRes& res)
    {
        using ElementType = typename TRes::ElementType;
        using LongerType = DoubleWidthType<ElementType>;

        res.zero();

        bool overflow = false;

        const size_t n = lhs.size_elems();
        const size_t m = rhs.size_elems();

        for (size_t i = 0; i < n; i++)
        {
            LongerType carry = 0;
            ElementType a = lhs[i];
            for (size_t j = 0; j < m; j++)
            {
                ElementType b = rhs[j];
                LongerType t = static_cast<LongerType>(a) * b + res.get(i + j) + carry;

                overflow |= !res.try_set(i + j, static_cast<ElementType>(t));
                carry = t >> (sizeof(ElementType) * 8);
            }
            overflow |= !res.try_set(i + m, static_cast<ElementType>(carry));
        }

        if (overflow)
        {
            throw std::overflow_error("Multiplication overflow");
        }
    }
};

struct Division
{
    template <AnyMpInt T>
    static void shift_left(T& value)
    {
        using ElementType = typename T::ElementType;

        for (size_t i = value.size_elems(); i-- > 0;)
        {
            ElementType new_val = value.get(i - 1);
            value.set(i, new_val);
        }

        value.set(0, ElementType{0});
    }

    template <AnyConstMpInt TLhs, AnyConstMpInt TRhs, AnyMpInt TRes>
        requires ElementTypesMatch<TLhs, TRes> && ElementTypesMatch<TRhs, TRes>
    static void invoke(const TLhs& lhs, const TRhs& rhs, TRes& res)
    {
        using ElementType = typename TRes::ElementType;
        using LongerType = DoubleWidthType<ElementType>;

        if (rhs == 0U)
        {
            throw std::runtime_error("Division by zero");
        }

        BasicInt<ElementType, TRes::MAX_BYTES> remainder = lhs;
        res.zero();
        BasicInt<ElementType, TRhs::MAX_BYTES> temp;

        while (remainder >= rhs)
        {
            int shift = remainder.size_elems() - rhs.size_elems();
            
            temp = rhs;
            for (int i = 0; i < shift; ++i)
            {
                shift_left(temp);
            }

            ElementType q_digit = 0;

            if (remainder >= temp)
            {
                remainder -= temp;
                q_digit++;
            }

            res.set(shift, q_digit);
        }        
    };
};

// Binary operation
template <typename TOp, AnyConstMpInt TLhs, AnyConstInt TRhs>
OpResultType<TLhs, TRhs> binary_op(const TLhs& lhs, const TRhs& rhs)
{
    OpResultType<TLhs, TRhs> res{};

    try
    {
        TOp::invoke(lhs, to_mp_int<typename TLhs::ElementType>(rhs), res);
    }
    catch (const std::overflow_error&)
    {
        throw OverflowErrorOf<TLhs, TRhs>(std::move(res));
    }

    return res;
}

// Compound assignment binary operation
template <typename TOp, AnyMpInt TLhs, AnyConstInt TRhs>
TLhs& ca_binary_op(TLhs& lhs, const TRhs& rhs)
{
    TLhs res{};

    try
    {
        TOp::invoke(lhs, to_mp_int<typename TLhs::ElementType>(rhs), res);
    }
    catch (const std::overflow_error&)
    {
        throw OverflowErrorOf<TLhs, TRhs>(std::move(res));
    }

    lhs = std::move(res);
    return lhs;
}

// Addition operators

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
OpResultType<TLhs, TRhs> operator+(const TLhs& lhs, const TRhs& rhs)
{
    return binary_op<Addition>(lhs, rhs);
}

template <AnyMpInt TLhs, AnyConstInt TRhs>
TLhs& operator+=(TLhs& lhs, const TRhs& rhs)
{
    return ca_binary_op<Addition>(lhs, rhs);
}

// Subtraction operators

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
OpResultType<TLhs, TRhs> operator-(const TLhs& lhs, const TRhs& rhs)
{
    return binary_op<Subtraction>(lhs, rhs);
}

template <AnyMpInt TLhs, AnyConstInt TRhs>
TLhs& operator-=(TLhs& lhs, const TRhs& rhs)
{
    return ca_binary_op<Subtraction>(lhs, rhs);
}

// Multiplication operators

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
OpResultType<TLhs, TRhs> operator*(const TLhs& lhs, const TRhs& rhs)
{
    return binary_op<Multiplication>(lhs, rhs);
}

template <AnyMpInt TLhs, AnyConstInt TRhs>
TLhs& operator*=(TLhs& lhs, const TRhs& rhs)
{
    return ca_binary_op<Multiplication>(lhs, rhs);
}

// Division operators

template <AnyConstMpInt TLhs, AnyConstInt TRhs>
OpResultType<TLhs, TRhs> operator/(const TLhs& lhs, const TRhs& rhs)
{
    return binary_op<Division>(lhs, rhs);
}

template <AnyMpInt TLhs, AnyConstInt TRhs>
TLhs& operator/=(TLhs& lhs, const TRhs& rhs)
{
    return ca_binary_op<Division>(lhs, rhs);
}


} // namespace mp