//
// Created by Dan Crosby on 6/19/23.
//

#include "AIAttackGroups.hpp"

#include <cstdlib>
#include <iterator>
#include <ranges>
#include <unordered_map>

#include "src/main/cpp/net/eagle0/shardok/library/util/HexMapUtils.hpp"

namespace shardok {

constexpr double kOverpowerRatio = 2.0;
constexpr double kBraveWaterCostMultiplier = 1.2;

DIST_T NormalizedCostWhenBraving(const DIST_T cost) {
    if (cost >= static_cast<double>(ActionPointDistances::IMPOSSIBLE) / kBraveWaterCostMultiplier)
        return ActionPointDistances::IMPOSSIBLE;

    return static_cast<DIST_T>(kBraveWaterCostMultiplier * cost);
}

struct TargetAndDistance {
    Coords target;
    size_t targetIndex;
    CoordsSet attackLocations;
    double targetPower;
    DIST_T distance;

    TargetAndDistance(
            const Coords& t,
            const size_t ti,
            const CoordsSet& al,
            const double tp,
            const DIST_T d)
        : target(t),
          targetIndex(ti),
          attackLocations(al),
          targetPower(tp),
          distance(d) {}
};

auto CoordsIndex(const Coords& coords, const int columnCount) {
    return coords.row() * columnCount + coords.column();
}

auto CoordsFromIndex(const size_t index, const int columnCount) -> Coords {
    return Coords(
            static_cast<int8_t>(index / static_cast<size_t>(columnCount)),
            static_cast<int8_t>(index % static_cast<size_t>(columnCount)));
}

auto MinDistance(
        const Coords& start,
        const CoordsSet& destinations,
        const ActionPointDistances* apd) -> DIST_T {
    DIST_T minDistance = ActionPointDistances::IMPOSSIBLE;
    const int startIndex = CoordsIndex(start, destinations.ColumnCount());

    for (const auto destIndex : destinations.indexIterator()) {
        if (const DIST_T distance = apd->Distance(startIndex, static_cast<int>(destIndex));
            distance < minDistance) {
            minDistance = distance;
        }
    }

    return minDistance;
}

auto MinDistanceIncludingBraving(
        const Coords& start,
        const CoordsSet& destinations,
        const ActionPointDistances* notBravingApd,
        const ActionPointDistances* bravingApd) {
    // First try to get there without braving
    if (const DIST_T notBravingDistance = MinDistance(start, destinations, notBravingApd);
        notBravingDistance < ActionPointDistances::IMPOSSIBLE) {
        return notBravingDistance;
    }

    // Then normalize the cost with braving
    if (bravingApd) {
        return NormalizedCostWhenBraving(MinDistance(start, destinations, bravingApd));
    }
    return ActionPointDistances::IMPOSSIBLE;
}

auto EffectiveDistance(
        const Unit* unit,
        const HexMap* map,
        const AttackLocations& attackLocations,
        const APDCache& apdCache,
        const BattalionTypeGetter& battalionTypeGetter,
        ActionPoints braveWaterCost) -> DIST_T {
    return EffectiveDistance(
            unit,
            map,
            attackLocations.LocationsWithEnemyInRange(unit),
            apdCache,
            battalionTypeGetter,
            braveWaterCost);
}

auto EffectiveDistance(
        const Unit* unit,
        const HexMap* map,
        const CoordsSet& locations,
        const APDCache& apdCache,
        const BattalionTypeGetter& battalionTypeGetter,
        ActionPoints braveWaterCost) -> DIST_T {
    const auto mapId = ActionPointDistancesCache::GetMapId(map);
    const auto& battType = battalionTypeGetter(unit->battalion().type());
    const auto* notBravingApd = apdCache->GetRaw(map, mapId, battType, false);
    const ActionPointDistances* bravingApd = nullptr;
    if (battType->allowsBraveWater) {
        bravingApd = apdCache->GetRaw(map, mapId, battType, true, braveWaterCost);
    }

    return MinDistanceIncludingBraving(unit->location(), locations, notBravingApd, bravingApd);
}

auto EffectiveDistance(
        const Unit* unit,
        const ActionPointDistances* notBravingApd,
        const ActionPointDistances* bravingApd,
        const CoordsSet& locations) -> DIST_T {
    return MinDistanceIncludingBraving(unit->location(), locations, notBravingApd, bravingApd);
}

auto Power(const Unit* unit) -> double { return unit->battalion().size(); }

// Chooses a list of targets in priority order for each unit.
auto GenerateTargetPriorities(
        const vector<const Unit*>& occupants,
        const HexMap* map,
        const CoordsSet& targets,
        const PlayerId attackerId,
        const vector<const Unit*>& remainingUnits,
        const APDCache& apdCache,
        const ALCache& alCache,
        const BattalionTypeGetter& battalionTypeGetter,
        ActionPoints braveWaterCost,
        const bool isLateGame) -> vector<TargetPriorityList> {
    auto cc = map->column_count();

    const auto mapId = ActionPointDistancesCache::GetMapId(map);

    vector<TargetPriorityList> allTargetsUnitsAndDistances{};
    allTargetsUnitsAndDistances.reserve(remainingUnits.size());

    vector<double> powerAttackingEachTarget(occupants.size());

    vector<const Unit*> sortedAttackers = remainingUnits;
    // Handle stronger units first
    std::ranges::sort(sortedAttackers, [](const Unit* left, const Unit* right) {
        return Power(left) > Power(right);
    });

    // APDCache now has built-in thread-local caching - no need for local apdByBattType map
    // For each unit, sort the targets by distance from the unit to an attack location for the
    // target
    for (const Unit* unit : sortedAttackers) {
        auto& tpl = allTargetsUnitsAndDistances.emplace_back(
                unit->unit_id(),
                vector<TargetAndAttackLocations>{});

        vector<TargetAndDistance> targetsWithDistance;
        targetsWithDistance.reserve(targets.size());

        // Get APDs directly from cache (now with built-in thread-local optimization)
        const auto& battType = battalionTypeGetter(unit->battalion().type());
        const auto* notBravingApd = apdCache->GetRaw(map, mapId, battType, false);
        const ActionPointDistances* bravingApd = nullptr;
        if (battType->allowsBraveWater) {
            bravingApd = apdCache->GetRaw(map, mapId, battType, true, braveWaterCost);
        }

        for (const auto targetIndex : targets.indexIterator()) {
            const auto& occupant = occupants[targetIndex];

            const Coords targetLocation = CoordsFromIndex(targetIndex, cc);
            const auto& attackLocations = alCache->CachedLocations(targetLocation, isLateGame);

            if (occupant && occupant->player_id() != attackerId) {
                double occupantPower = Power(occupant);

                if (unit->location().row() >= 0) {
                    const auto& attackLocsForUnit = attackLocations.LocationsWithEnemyInRange(unit);
                    auto distance =
                            EffectiveDistance(unit, notBravingApd, bravingApd, attackLocsForUnit);

                    targetsWithDistance.emplace_back(
                            targetLocation,
                            targetIndex,
                            attackLocsForUnit,
                            occupantPower,
                            distance);
                } else {
                    targetsWithDistance.emplace_back(
                            targetLocation,
                            targetIndex,
                            attackLocations.LocationsWithEnemyInRange(unit),
                            occupantPower,
                            ActionPointDistances::IMPOSSIBLE);
                }
            }
        }

        // Sort by distance
        std::ranges::sort(
                targetsWithDistance,
                [&powerAttackingEachTarget](
                        const TargetAndDistance& left,
                        const TargetAndDistance& right) {
                    const double powerAttackingLeft = powerAttackingEachTarget[left.targetIndex];
                    const double powerAttackingRight = powerAttackingEachTarget[right.targetIndex];

                    const bool leftAlreadyOverwhelmed =
                            powerAttackingLeft >= kOverpowerRatio * left.targetPower;
                    const bool rightAlreadyOverwhelmed =
                            powerAttackingRight >= kOverpowerRatio * right.targetPower;

                    if (leftAlreadyOverwhelmed == rightAlreadyOverwhelmed)
                        return left.distance < right.distance;

                    return leftAlreadyOverwhelmed < rightAlreadyOverwhelmed;
                });

        if (!targetsWithDistance.empty()) {
            powerAttackingEachTarget[targetsWithDistance.front().targetIndex] += Power(unit);
        }

        tpl.priorityOrder.reserve(targetsWithDistance.size());
        std::ranges::transform(
                targetsWithDistance,
                std::back_inserter(tpl.priorityOrder),
                [](const TargetAndDistance& tad) {
                    return TargetAndAttackLocations{
                            .target = tad.target,
                            .attackLocations = tad.attackLocations};
                });
    }

    return allTargetsUnitsAndDistances;
}

}  // namespace shardok
