/*
 * Copyright (c) 2020, Andreas Kling <andreas@ladybird.org>
 * Copyright (c) 2025, Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#pragma once

#include <AK/Utf16FlyString.h>
#include <LibJS/Runtime/Completion.h>
#include <LibJS/Runtime/PrimitiveString.h>
#include <LibJS/Runtime/Symbol.h>

namespace JS {

class PropertyKey {
public:
    enum class StringMayBeNumber {
        Yes,
        No,
    };

    static ThrowCompletionOr<PropertyKey> from_value(VM& vm, Value value)
    {
        VERIFY(!value.is_special_empty_value());
        if (value.is_symbol())
            return PropertyKey { value.as_symbol() };
        if (value.is_integral_number() && value.as_double() >= 0 && value.as_double() < NumericLimits<u32>::max())
            return static_cast<u32>(value.as_double());
        return TRY(value.to_utf16_string(vm));
    }

    static constexpr uintptr_t NORMAL_STRING_FLAG = 0;
    static constexpr uintptr_t SHORT_STRING_FLAG = 1;
    static constexpr uintptr_t SYMBOL_FLAG = 2;
    static constexpr uintptr_t NUMBER_FLAG = 3;

    bool is_string() const { return (m_bits & 3) == NORMAL_STRING_FLAG || (m_bits & 3) == SHORT_STRING_FLAG; }
    bool is_number() const { return (m_bits & 3) == NUMBER_FLAG; }
    bool is_symbol() const { return (m_bits & 3) == SYMBOL_FLAG; }

    PropertyKey() = delete;

    PropertyKey(PropertyKey const& other)
    {
        if (other.is_string())
            new (&m_string) Utf16FlyString(other.m_string);
        else
            m_bits = other.m_bits;
    }

    PropertyKey(PropertyKey&& other) noexcept
    {
        if (other.is_string())
            new (&m_string) Utf16FlyString(move(other.m_string));
        else
            m_bits = exchange(other.m_bits, 0);
    }

    template<Integral T>
    PropertyKey(T index)
    {
        // FIXME: Replace this with requires(IsUnsigned<T>)?
        //        Needs changes in various places using `int` (but not actually being in the negative range)
        VERIFY(index >= 0);
        if constexpr (NumericLimits<T>::max() >= NumericLimits<u32>::max()) {
            if (index >= NumericLimits<u32>::max()) {
                new (&m_string) Utf16FlyString { Utf16String::number(index) };
                return;
            }
        }
        m_number = static_cast<u64>(index) << 2 | NUMBER_FLAG;
    }

    PropertyKey(Utf16FlyString string, StringMayBeNumber string_may_be_number = StringMayBeNumber::Yes)
    {
        if (string_may_be_number == StringMayBeNumber::Yes) {
            if (!string.is_empty() && !(string.code_unit_at(0) == '0' && string.length_in_code_units() > 1)) {
                auto property_index = string.to_number<u32>(TrimWhitespace::No);
                if (property_index.has_value() && property_index.value() < NumericLimits<u32>::max()) {
                    m_number = static_cast<u64>(property_index.release_value()) << 2 | NUMBER_FLAG;
                    return;
                }
            }
        }

        new (&m_string) Utf16FlyString(move(string));
    }

    PropertyKey(Utf16String const& string)
        : PropertyKey(Utf16FlyString { string })
    {
    }

    PropertyKey(GC::Ref<Symbol> symbol)
    {
        m_bits = reinterpret_cast<uintptr_t>(symbol.ptr()) | SYMBOL_FLAG;
    }

    PropertyKey& operator=(PropertyKey const& other)
    {
        if (this != &other) {
            if (is_string())
                m_string.~Utf16FlyString();
            new (this) PropertyKey(other);
        }
        return *this;
    }

    PropertyKey& operator=(PropertyKey&& other) noexcept
    {
        if (this != &other) {
            if (is_string())
                m_string.~Utf16FlyString();
            new (this) PropertyKey(move(other));
        }
        return *this;
    }

    ~PropertyKey()
    {
        if (is_string())
            m_string.~Utf16FlyString();
    }

    u32 as_number() const
    {
        VERIFY(is_number());
        return m_number >> 2;
    }

    Utf16FlyString const& as_string() const
    {
        VERIFY(is_string());
        return m_string;
    }

    Symbol const* as_symbol() const
    {
        VERIFY(is_symbol());
        return reinterpret_cast<Symbol const*>(m_bits & ~3ULL);
    }

    Value to_value(VM& vm) const
    {
        if (is_string())
            return Value { PrimitiveString::create(vm, as_string()) };
        if (is_symbol())
            return Value { as_symbol() };
        return Value { PrimitiveString::create_from_unsigned_integer(vm, as_number()) };
    }

    Utf16String to_string() const
    {
        if (is_string())
            return as_string().to_utf16_string();
        if (is_symbol())
            return as_symbol()->descriptive_string();
        return Utf16String::number(as_number());
    }

    void visit_edges(Cell::Visitor& visitor) const
    {
        if (is_symbol())
            visitor.visit(const_cast<Symbol*>(as_symbol()));
    }

    bool operator==(PropertyKey const& other) const
    {
        if (is_string())
            return other.is_string() && m_string == other.m_string;
        if (is_symbol())
            return other.is_symbol() && as_symbol() == other.as_symbol();
        if (other.is_number())
            return as_number() == other.as_number();
        return false;
    }

private:
    friend Traits<PropertyKey>;
    friend struct AK::SentinelOptionalTraits<PropertyKey>;

    enum class ShouldMakeEmptyOptional {
        Indeed,
    };

    explicit PropertyKey(ShouldMakeEmptyOptional)
        : m_bits(0)
    {
    }

    [[nodiscard]] bool is_empty_optional() const { return m_bits == 0; }

    union {
        Utf16FlyString m_string;
        u64 m_number;
        Symbol const* m_symbol;
        uintptr_t m_bits;
    };
};

static_assert(sizeof(PropertyKey) == sizeof(uintptr_t));

}

namespace AK {

template<>
struct Traits<JS::PropertyKey> : public DefaultTraits<JS::PropertyKey> {
    static unsigned hash(JS::PropertyKey const& name)
    {
        if (name.is_string())
            return name.as_string().hash();
        if (name.is_symbol())
            return ptr_hash(name.as_symbol());
        if (name.is_number())
            return u32_hash(name.as_number());
        VERIFY_NOT_REACHED();
    }

    static bool equals(JS::PropertyKey const& a, JS::PropertyKey const& b)
    {
        if (a.is_string())
            return b.is_string() && a.as_string() == b.as_string();
        if (a.is_symbol())
            return b.is_symbol() && a.as_symbol() == b.as_symbol();
        if (a.is_number())
            return b.is_number() && a.as_number() == b.as_number();
        VERIFY_NOT_REACHED();
    }
};

template<>
struct Formatter<JS::PropertyKey> : Formatter<Utf16String> {
    ErrorOr<void> format(FormatBuilder& builder, JS::PropertyKey const& property_key)
    {
        if (property_key.is_number())
            return builder.put_u64(property_key.as_number());
        return Formatter<Utf16String>::format(builder, property_key.to_string());
    }
};

template<>
struct SentinelOptionalTraits<JS::PropertyKey> {
    static JS::PropertyKey sentinel_value() { return JS::PropertyKey { JS::PropertyKey::ShouldMakeEmptyOptional::Indeed }; }
    static bool is_sentinel(JS::PropertyKey const& value) { return value.is_empty_optional(); }
};

template<>
class Optional<JS::PropertyKey> : public SentinelOptional<JS::PropertyKey> {
public:
    using SentinelOptional::SentinelOptional;
};

}