#include <maps/wikimap/mapspro/libs/flat_range/include/distribution.h>
#include <maps/wikimap/mapspro/libs/flat_range/include/range.h>
#include <maps/wikimap/mapspro/libs/flat_range/include/validation.h>

#include <yandex/maps/wiki/common/natural_sort.h>

#include <algorithm>
#include <optional>
#include <unordered_map>

namespace maps::wiki::flat_range {

struct FlatRange {
    int first;
    int last;
};

namespace {

std::optional<FlatRange>
getFlatRangeOnLevel(int flatFirst, int flatLast, int levelFirst, int levelLast, int level)
{
    ASSERT(flatFirst <= flatLast && levelFirst <= levelLast);
    ASSERT(level <= levelLast && level >= levelFirst);
    const int numLevels = (levelLast - levelFirst + 1);
    const int numFlats = (flatLast - flatFirst + 1);
    int excessFlats = numFlats % numLevels;
    int flatsPerLevel = numFlats / numLevels;
    int levelsWithFullFlats = numLevels;
    if (!excessFlats) {
        return FlatRange {
            (level - levelFirst) * flatsPerLevel,
            (level + 1 - levelFirst) * flatsPerLevel - 1,
        };
    }
    flatsPerLevel = flatsPerLevel + 1;
    levelsWithFullFlats = numFlats / flatsPerLevel;
    int numFlatsOnPartialLevel = (numFlats % flatsPerLevel);
    if (level < (levelLast - levelsWithFullFlats) ||
        (level == (levelLast - levelsWithFullFlats) && !numFlatsOnPartialLevel)) {
        return std::nullopt;
    }
    int lastFlatOnPartialLevel = flatFirst + numFlatsOnPartialLevel - 1;
    if (level == (levelLast - levelsWithFullFlats)) {
        return FlatRange {
            flatFirst,
            lastFlatOnPartialLevel
        };
    }
    return FlatRange {
        (level - (levelLast - levelsWithFullFlats) - 1) * flatsPerLevel + lastFlatOnPartialLevel + 1,
        (level - (levelLast - levelsWithFullFlats)) * flatsPerLevel + lastFlatOnPartialLevel
    };
}

struct LevelWithFlatsLesser
{
    bool operator()(const LevelWithFlats& lhs, const LevelWithFlats& rhs) {
        return common::natural_sort()(rhs.levelName, lhs.levelName);
    }
};

} // namespace

DistributionResult
distributeFlatsByLevels(const std::vector<FlatLevelRange>& flatLevelRanges)
{
    std::unordered_map<std::string, Ranges> levelNameToFlatRanges;
    auto isComplete = true;

    for (const auto& flatLevelRange : flatLevelRanges) {
        if (flatLevelRange.levels.empty()) {
            isComplete = false;
            continue;
        }

        Ranges levelRanges;
        try {
            levelRanges = parse(flatLevelRange.levels);
        } catch (ParseException&) {
            isComplete = false;
            continue;
        }
        if (levelRanges.size() != 1) {
            isComplete = false;
            continue;
        }

        Ranges flatRanges;
        try {
            flatRanges = parse(flatLevelRange.flats);
        } catch (ParseException&) {
            isComplete = false;
            continue;
        }
        const auto& levelRange = levelRanges[0];
        if (flatRanges.size() > 1 && levelRange.size() > 1) {
            isComplete = false;
            continue;
        }
        if (levelRange.size() == 1) {
            auto levelName = levelRange.first();
            levelNameToFlatRanges[levelName].insert(
                levelNameToFlatRanges[levelName].end(),
                flatRanges.begin(),
                flatRanges.end());
            continue;
        }

        const auto& flatRange = flatRanges[0];
        if (flatRange.size() > MAX_FLATS_PER_ENTRANCE) {
            isComplete = false;
            continue;
        }
        for (size_t levelNum = 0; levelNum < levelRange.size(); ++levelNum) {
            auto flatRangeOnLevel = getFlatRangeOnLevel(
                0, flatRange.size() - 1, 0, levelRange.size() - 1, levelNum);
            if (!flatRangeOnLevel) {
                continue;
            }

            auto levelName = levelRange.value(levelNum);
            levelNameToFlatRanges[levelName].insert(
                levelNameToFlatRanges[levelName].end(),
                Range(
                    flatRange.value(flatRangeOnLevel->first),
                    flatRange.value(flatRangeOnLevel->last)));
        }
    }

    DistributionResult result {
        .isComplete = isComplete
    };

    for (const auto& [levelName, flatRanges] : levelNameToFlatRanges) {
        result.levelsWithFlats.push_back({levelName, flatRanges});
    }

    LevelWithFlatsLesser levelWithFlatsLesser;
    std::sort(
        result.levelsWithFlats.begin(),
        result.levelsWithFlats.end(),
        levelWithFlatsLesser);
    return result;
}

} // namespace maps::wiki::flat_range
