//
// Created by Dan Crosby on 9/6/20.
//

#include "src/main/cpp/net/eagle0/shardok/library/action_point_distances/ActionPointDistances.hpp"

#include <queue>
#include <utility>

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

namespace shardok {
using std::begin;
using std::end;
using std::vector;

struct ActionPointDistances::BraveableTileInfo {
    struct BraveableFromInfo {
        Coords from;
        CoordsSet to;
    };

    CoordsSet cs;
    vector<BraveableFromInfo> details;

    explicit BraveableTileInfo(const HexMap* hexMap) : cs(hexMap) {}
};

auto ActionPointDistances::BraveWaterPossibleCoords(const HexMap* hexMap) const
        -> std::shared_ptr<BraveableTileInfo> {
    std::shared_ptr<BraveableTileInfo> info = std::make_shared<BraveableTileInfo>(hexMap);
    CoordsSet& braveWaterPossibleCoords = info->cs;
    const int indexCount = hexMap->row_count() * hexMap->column_count();
    const auto indexToCoords = CreateIndexToCoords(hexMap);
    for (int i = 0; i < indexCount; i++) {
        if (const Terrain* terrain = hexMap->terrain()->Get(i);
            IsWater(terrain->type()) && !terrain->modifier().bridge().present() &&
            !terrain->modifier().ice().present()) {
            const Coords& currentCoords = indexToCoords[i];

            const CoordsSet adjacentToWaterCoords =
                    HexMapUtils::GetAdjacentCoords(hexMap, currentCoords);
            CoordsSet braveableTo(hexMap);
            for (const Coords& adjacentToWater : adjacentToWaterCoords) {
                for (auto acrossWaterCoords = GetTilesAcrossWater(hexMap, adjacentToWater);
                     const auto& braveCoords : acrossWaterCoords) {
                    const auto braveIndex = ToIndex(braveCoords);
                    if (const auto* swimTerrain = hexMap->terrain()->Get(braveIndex);
                        IsWater(swimTerrain->type()) ||
                        swimTerrain->type() ==
                                net::eagle0::shardok::storage::fb::Terrain_::Type_MOUNTAIN) {
                        continue;
                    }

                    braveableTo.Add(braveCoords);
                }

                if (!braveableTo.empty()) {
                    info->details.push_back(
                            BraveableTileInfo::BraveableFromInfo{adjacentToWater, braveableTo});
                    braveWaterPossibleCoords.Add(adjacentToWater);
                }
            }
        }
    }

    return info;
}

auto ActionPointDistances::CreateIndexToCoords(const HexMap* hexMap) -> vector<Coords> {
    const int8_t columnCount = hexMap->column_count();
    const int indexCount = hexMap->row_count() * columnCount;
    vector<Coords> indexToCoords;
    indexToCoords.reserve(indexCount);
    for (int i = 0; i < indexCount; i++) {
        indexToCoords.emplace_back(
                static_cast<int8_t>(i / columnCount),
                static_cast<int8_t>(i % columnCount));
    }
    return indexToCoords;
}

auto ActionPointDistances::CreateAdjacencyTable(const HexMap* hexMap)
        -> vector<std::array<int, 6>> {
    const int8_t columnCount = hexMap->column_count();
    const int indexCount = hexMap->row_count() * columnCount;
    vector<std::array<int, 6>> adjacencyTable;
    adjacencyTable.reserve(indexCount);
    for (int i = 0; i < indexCount; i++) {
        const Coords coords(
                static_cast<int8_t>(i / columnCount),
                static_cast<int8_t>(i % columnCount));
        const CoordsSet adjacentCoords = HexMapUtils::GetAdjacentCoords(hexMap, coords);
        std::array<int, 6> neighbors{};
        neighbors.fill(-1);  // -1 indicates invalid/no neighbor
        int neighborIdx = 0;
        for (const auto adjacentIndex : adjacentCoords.indexIterator()) {
            if (neighborIdx < 6) { neighbors[neighborIdx++] = static_cast<int>(adjacentIndex); }
        }
        adjacencyTable.push_back(neighbors);
    }

    return adjacencyTable;
}

void ActionPointDistances::PopulateOne(
        vector<DIST_T>& ds,
        const HexMap* hexMap,
        const bool includeBravingWater,
        const int braveWaterCost,
        const BattalionTypeSPtr& battalionType,
        const std::shared_ptr<BraveableTileInfo>& braveWaterPossibleCoords,
        const vector<Coords>& indexToCoords,
        const vector<std::array<int, 6>>& adjacencyTable) {
    vector<uint8_t> visited(ds.size());

    // Priority queue for efficient minimum selection: {distance, index}
    std::priority_queue<std::pair<DIST_T, int>, std::vector<std::pair<DIST_T, int>>, std::greater<>>
            pq;

    // Find starting index (the one with distance 0)
    for (int i = 0; i < static_cast<int>(ds.size()); i++) {
        if (ds[i] == 0) {
            pq.emplace(0, i);
            break;
        }
    }

    while (!pq.empty()) {
        auto [currentIndexDistance, currentIndex] = pq.top();
        pq.pop();

        // Skip if already visited (can happen due to multiple insertions)
        if (visited[currentIndex]) continue;

        // Skip if we found a better path since insertion
        if (currentIndexDistance > ds[currentIndex]) continue;

        const Coords& currentCoords = indexToCoords[currentIndex];
        visited[currentIndex] = true;

        // Prefetch terrain data for all neighbors to reduce memory stalls
        const std::array<int, 6>& neighbors = adjacencyTable[currentIndex];
        for (int i = 0; i < 6 && neighbors[i] != -1; i++) {
            __builtin_prefetch(hexMap->terrain()->Get(neighbors[i]), 0, 3);
        }

        for (int adjacentIndex : neighbors) {
            if (adjacentIndex == -1) break;  // End of valid neighbors
            if (visited[adjacentIndex]) continue;

            const auto adjacentTerrain = hexMap->terrain()->Get(adjacentIndex);
            const auto adjacentCost = battalionType->GetCostToEnterTerrain(adjacentTerrain);

            if (adjacentCost.type == ActionCost::impossible) continue;

            const auto costThroughCurrentTile = currentIndexDistance + adjacentCost.points;
            if (const int currentBestDistance = ds[adjacentIndex];
                currentBestDistance > costThroughCurrentTile) {
                ds[adjacentIndex] = static_cast<DIST_T>(costThroughCurrentTile);
                // Add to priority queue for future processing
                pq.emplace(ds[adjacentIndex], adjacentIndex);
            }
        }

        // check for swimmable tiles
        if (includeBravingWater && battalionType->allowsBraveWater &&
            braveWaterPossibleCoords->cs.Contains(currentCoords)) {
            auto entry = std::ranges::find_if(
                    braveWaterPossibleCoords->details,
                    [currentCoords](const BraveableTileInfo::BraveableFromInfo& from) {
                        return from.from == currentCoords;
                    });
            if (entry != braveWaterPossibleCoords->details.end()) {
                // Prefetch terrain data for water braving targets
                for (const auto braveIndex : entry->to.indexIterator()) {
                    __builtin_prefetch(hexMap->terrain()->Get(braveIndex), 0, 3);
                }

                for (const auto braveIndex : entry->to.indexIterator()) {
                    if (visited[braveIndex]) continue;

                    const auto* swimTerrain = hexMap->terrain()->Get(braveIndex);

                    if (const auto adjacentCost = battalionType->GetCostToEnterTerrain(swimTerrain);
                        adjacentCost.type == ActionCost::impossible)
                        continue;

                    const auto costThroughCurrentTile = currentIndexDistance + braveWaterCost;
                    if (const int currentBestDistance = ds[braveIndex];
                        currentBestDistance > costThroughCurrentTile) {
                        ds[braveIndex] = static_cast<DIST_T>(costThroughCurrentTile);
                        // Add to priority queue for future processing
                        pq.push({ds[braveIndex], braveIndex});
                    }
                }
            }
        }
    }
}

auto ActionPointDistances::GenerateDistances(
        const int fromIndex,
        const HexMap* hexMap,
        const bool includeBravingWater,
        const int braveWaterCost,
        const BattalionTypeSPtr& battalionType,
        const std::shared_ptr<BraveableTileInfo>& braveWaterPossibleCoords) -> vector<DIST_T> {
    const int indexCount = hexMap->row_count() * hexMap->column_count();
    auto ds = vector(indexCount, IMPOSSIBLE);

    ds[fromIndex] = 0;

    if (!includeBravingWater || battalionType->allowsBraveWater) {
        // Create lookup tables once per distance calculation
        const auto indexToCoords = CreateIndexToCoords(hexMap);
        const auto adjacencyTable = CreateAdjacencyTable(hexMap);

        PopulateOne(
                ds,
                hexMap,
                includeBravingWater,
                braveWaterCost,
                battalionType,
                braveWaterPossibleCoords,
                indexToCoords,
                adjacencyTable);
    }

    return ds;
}

OnDemandActionPointDistances::OnDemandActionPointDistances(
        const HexMap* map,
        BattalionTypeSPtr battTp,
        bool includeBravingWater,
        int braveWaterActionPointCost)
    : ActionPointDistances(map->column_count()),
      hexMap(fb::CopyHexMap(map)),
      battalionType(std::move(battTp)),
      distances(static_cast<size_t>(map->row_count() * map->column_count())) {
    const int indexCount = map->row_count() * map->column_count();
    distances.resize(indexCount);

    auto braveWaterPossibleCoords = includeBravingWater ? BraveWaterPossibleCoords(map) : nullptr;

    for (int fromIndex = 0; fromIndex < indexCount; fromIndex++) {
        distances[fromIndex] = std::async(
                                       std::launch::deferred,
                                       &OnDemandActionPointDistances::GenerateDistances,
                                       fromIndex,
                                       hexMap,
                                       includeBravingWater,
                                       braveWaterActionPointCost,
                                       battalionType,
                                       braveWaterPossibleCoords)
                                       .share();
    }
}

}  // namespace shardok