// Copyright (C) 2024-2025  ilobilo

#pragma once

namespace std
{
    template<bool B>
    struct mutex_err
    {
        [[noreturn]]
        constexpr void err() { static_assert(B, "std::mutex is not implemented, use lib::mutex"); }

        constexpr void lock() { err(); }
        constexpr bool try_lock() { err(); }
        constexpr bool try_lock_until(auto) { err(); }
        constexpr void unlock() { err(); }
    };
    using mutex = mutex_err<false>;

    struct defer_lock_t { explicit defer_lock_t() = default; };
    struct try_to_lock_t { explicit try_to_lock_t() = default; };
    struct adopt_lock_t { explicit adopt_lock_t() = default; };

    inline constexpr defer_lock_t  defer_lock { };
    inline constexpr try_to_lock_t try_to_lock { };
    inline constexpr adopt_lock_t  adopt_lock { };

    template<typename Type>
    class unique_lock
    {
        private:
        Type *_mutex;
        bool _locked;

        public:
        using mutex_type = Type;

        unique_lock() : _mutex { nullptr }, _locked { false } { }
        unique_lock(unique_lock &&other) noexcept : _mutex { nullptr }, _locked { false } { this->swap(other); }

        explicit unique_lock(mutex_type &m) : _mutex { addressof(m) }, _locked { true } { m.lock(); }
        unique_lock(mutex_type &m, defer_lock_t) noexcept : _mutex { addressof(m) }, _locked { false } { }
        unique_lock(mutex_type &m, try_to_lock_t) : _mutex { addressof(m) }, _locked { true } { m.try_lock(); }
        unique_lock(mutex_type &m, adopt_lock_t) : _mutex { addressof(m) }, _locked { true } { }

        ~unique_lock()
        {
            if (_locked)
                _mutex->unlock();
        }

        unique_lock &operator=(unique_lock &&other)
        {
            if (_locked)
                unlock();

            this->swap(other);
            return *this;
        }

        void lock()
        {
            if (_mutex)
            {
                _mutex->lock();
                _locked = true;
            }
        }

        bool try_lock()
        {
            if (_mutex)
            {
                auto ret = _mutex->try_lock();
                if (ret == true)
                    _locked = true;
                return ret;
            }
            return false;
        }

        void unlock()
        {
            if (_mutex)
            {
                _mutex->unlock();
                _locked = false;
            }
        }

        void swap(unique_lock &other) noexcept
        {
            using std::swap;
            swap(_mutex, other._mutex);
            swap(_locked, other._locked);
        }

        mutex_type *release() noexcept
        {
            auto ret = _mutex;
            _mutex = nullptr;
            _locked = false;
            return ret;
        }

        mutex_type *mutex() const noexcept
        {
            return _mutex;
        }

        bool owns_lock() const noexcept
        {
            return _locked;
        }

        explicit operator bool() const
        {
            return owns_lock();
        }

        friend void swap(unique_lock &lhs, unique_lock &rhs) noexcept
        {
            return lhs.swap(rhs);
        }
    };

    template<typename L0, typename L1>
    void lock(L0 &l0, L1 &l1)
    {
        while (true)
        {
            {
                unique_lock<L0> u0(l0);
                if (l1.try_lock())
                {
                    u0.release();
                    break;
                }
            }
            {
                unique_lock<L1> u1(l1);
                if (l0.try_lock())
                {
                    u1.release();
                    break;
                }
            }
        }
    }
} // namespace std