#include "factors_domain.h"
#include <kernel/factor_storage/factor_view_base.h>
#include <util/generic/typetraits.h>

static_assert(std::is_same<decltype(NRTYFactors::TIndexWebBase::Offset), NFactorSlices::TFactorIndex>::value);

namespace NRTYFactors {
    TBaseDomain::TBaseDomain(const TSet<NFactorSlices::EFactorSlice>& sliceIds, NFactorSlices::EFactorUniverse universe) {
        // we need a local slicesInfo object here, due to the fact that TConfig may be created multiple times (with different sliceIds)
        NFactorSlices::TSlicesMetaInfo slicesInfo = NFactorSlices::TGlobalSlicesMetaInfo::Instance();
        for (NFactorSlices::EFactorSlice sliceId : sliceIds) {
            NFactorSlices::EnableSlices(slicesInfo, sliceId);
        }

        for (NFactorSlices::EFactorSlice sliceId : sliceIds) {
            Y_ENSURE(slicesInfo.IsSliceInitialized(sliceId), "slice " << sliceId << " has not been initialized. Check PEERDIRs (is the codegen library linked?)");
            Y_ENSURE(slicesInfo.GetNumFactors(sliceId) > 0, "slice has no factors: " << sliceId);
        }

        FactorsDomain_ = NFactorSlices::TFactorDomain(slicesInfo, universe);
        for (auto it = FactorsDomain_.Begin(); it != FactorsDomain_.End(); ++it) {
            Lookup_.insert({it.GetFactorInfo().GetFactorName(), {it.GetIndex(), it.GetLeaf()}});
        }
    }

    TIndexWebBase TBaseDomain::GetBaseIdx(const TString& factorName, TMaybe<NFactorSlices::EFactorSlice> preferredSlice) const {
        const size_t cnt = Lookup_.count(factorName);
        if (Y_UNLIKELY(cnt == 0)) {
            ythrow yexception() << "dynamic factor not found: " << factorName;
        }

        auto&& [begin, end] = Lookup_.equal_range(factorName);
        if (Y_LIKELY(cnt == 1)) {
            if (preferredSlice.Defined() && preferredSlice != begin->second.SliceId) {
                ythrow yexception() << "dynamic factor not found: " << *preferredSlice << "/" << factorName;
            }
            return TIndexWebBase{begin->second.BaseIdx};
        } else {
            i32 bestPrio = Min<i32>();
            i32 lastNonUnique = Min<i32>();
            NFactorSlices::TFactorIndex baseIdx{};
            for (auto p = begin; p != end; ++p) {
                i32 prio = GetLookupPriority(p->second.SliceId, preferredSlice);
                if (prio < bestPrio)
                    continue;
                else if (prio == bestPrio) {
                    lastNonUnique = prio;
                } else {
                    baseIdx = p->second.BaseIdx;
                    bestPrio = prio;
                }
            }
            if (bestPrio == lastNonUnique) {
                ythrow yexception() << "dynamic factor exists in multiple slices, cannot decide: " << FormatNonUnique(begin, end);
            }
            return TIndexWebBase{baseIdx};
        }
    }

    TIndexWebBase TBaseDomain::GetBaseIdx(const ui32 webProductionFactorId) const {
        TFactorIndex baseIdx = FactorsDomain_.GetRelativeIndex(EFactorSlice::ALL, NFactorSlices::TFullFactorIndex(EFactorSlice::WEB_PRODUCTION, webProductionFactorId));
        return TIndexWebBase{baseIdx};
    }

    TMaybe<TFactorIndex> TBaseDomain::GetWebProductionFactorId(TIndexWebBase baseIdx) const {
        // GetWebProductionFactorId maps TIndexWebBase back to TFactorIndex, so GetWebProductionFactorId(GetBaseIdx(i)) always equals i
        const TSliceOffsets& borders = FactorsDomain_[EFactorSlice::WEB_PRODUCTION];
        if (baseIdx.Defined() && borders.Contains(baseIdx.Offset)) {
            TFactorIndex result = borders.GetRelativeIndex(baseIdx.Offset);
            Y_ASSERT(result >= 0 && GetBaseIdx(static_cast<ui32>(result)) == baseIdx);
            return result;
        } else {
            return Nothing();
        }
    }

    Y_FORCE_INLINE i32 TBaseDomain::GetLookupPriority(NFactorSlices::EFactorSlice slice, TMaybe<NFactorSlices::EFactorSlice> preferredSlice) {
        if (preferredSlice.Defined() && slice == *preferredSlice)
            return Max<i32>();
        else if (slice == NFactorSlices::EFactorSlice::WEB_PRODUCTION)
            return 1;
        else if (slice == NFactorSlices::EFactorSlice::BEGEMOT_QUERY_FACTORS)
            return 0; // WEB_PRODUCTION overrides BEGEMOT_QUERY_FACTORS
        else
            return 2; // all other slices override WEB_PRODUCTION
    }

    TString TBaseDomain::FormatNonUnique(TLookup::const_iterator begin, TLookup::const_iterator end) {
        TStringStream ss;
        for (auto p = begin; p != end; ++p) {
            if (!ss.Empty())
                ss << ", ";
            ss << p->second.SliceId << "/" << p->first;
        }
        return ss.Str();
    }

    void TBaseDomain::SetTrait(EFactorSlice slice, TFactorIndex factorId, ui32 category) {
        NFactorSlices::TFullFactorIndex ffi(slice, factorId);
        TIndexWebBase baseIdx{FactorsDomain_.GetRelativeIndex(EFactorSlice::ALL, ffi)};
        Traits_.insert({baseIdx.Offset, category});
    }

    void TBaseDomain::GetTrait(TSet<ui32>& result, const TIndexWebBase baseIdx) const {
        if (!baseIdx.Defined())
            return;
        Y_ASSERT(baseIdx.Offset < Max<TFactorIndex>());
        auto begin = Traits_.lower_bound({baseIdx.Offset, 0});
        auto end = Traits_.lower_bound({baseIdx.Offset + 1, 0});
        for (auto p = begin; p != end; ++p) {
            Y_ASSERT(p->first == baseIdx.Offset);
            result.insert(p->second);
        }
    }

}
