//
// Created by Dan Crosby on 12/29/21.
//

#include "AIWaterCrossingCalculator.hpp"

#include "AIMinimumDistanceAndTarget.hpp"
#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 {

auto UnitIdsRequiringWaterCrossing(
        const GameStateW &gameState,
        const PlayerId pid,
        const CoordsSet &destinations,
        const APDCache &apdCache,
        const BattalionTypeGetter &battalionTypeGetter) -> vector<UnitId> {
    // Put out all the fires, except on bridges
    fb::HexMapW mapCopy = fb::CopyHexMap(gameState->hex_map());
    for (uint32_t index = 0; index < mapCopy->terrain()->size(); index++) {
        if (IsWater(mapCopy->terrain()->Get(index)->type())) continue;
        mapCopy->mutable_terrain()
                ->GetMutableObject(index)
                ->mutable_modifier()
                .mutable_fire()
                .mutate_present(false);
    }

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

    vector<UnitId> unitIdsToReturn{};

    // Set this to -1 for unchecked, 0 for does not need crossing, 1 for does need crossing
    std::vector<int> startingPositionIndices(9, -1);
    for (const auto *unit : *gameState->units()) {
        if (unit->player_id() != pid) continue;

        const auto &battType = battalionTypeGetter(unit->battalion().type());

        if (unit->status() == net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT) {
            for (const Coords &destination : destinations) {
                if (!CanReach(unit->location(), destination, mapCopy, mapId, apdCache, battType)) {
                    unitIdsToReturn.push_back(unit->unit_id());
                    break;
                }
            }
        } else if (unit->status() == net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT) {
            if (startingPositionIndices[unit->starting_position_index()] == -1) {
                startingPositionIndices[unit->starting_position_index()] = 0;
                const auto *startingPositions = gameState->hex_map()
                                                        ->attacker_starting_positions()
                                                        ->Get(unit->starting_position_index())
                                                        ->positions();

                for (const Coords *pos : *startingPositions) {
                    for (const Coords &destination : destinations) {
                        if (!CanReach(*pos, destination, mapCopy, mapId, apdCache, battType)) {
                            startingPositionIndices[unit->starting_position_index()] = 1;
                            break;
                        }
                        if (startingPositionIndices[unit->starting_position_index()] == 1) break;
                    }
                }
            }

            if (startingPositionIndices[unit->starting_position_index()] == 1) {
                unitIdsToReturn.push_back(unit->unit_id());
            }
        }
    }

    return unitIdsToReturn;
}

auto UnitIdsToCreateWaterCrossing(
        const GameStateW &gameState,
        const PlayerId pid,
        const BattalionTypeGetter &battalionTypeGetter) -> vector<UnitId> {
    vector<UnitId> unitIds{};

    for (const auto *unit : *gameState->units()) {
        if (unit->player_id() != pid) continue;
        if (unit->status() != net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT &&
            unit->status() != net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT)
            continue;
        if (!unit->has_attached_hero()) continue;

        const auto profession = unit->attached_hero().profession_info().profession();
        const auto &battalionType = battalionTypeGetter(unit->battalion().type());

        if (profession == net::eagle0::shardok::storage::fb::Profession_ENGINEER ||
            (profession == net::eagle0::shardok::storage::fb::Profession_MAGE &&
             battalionType->allowsCasting)) {
            unitIds.push_back(unit->unit_id());
        }
    }

    return unitIds;
}

auto CanReach(
        const Coords &origin,
        const Coords &destination,
        const HexMap *hexMap,
        const MapId &mapId,
        const APDCache &apdCache,
        const BattalionTypeSPtr &battalionType) -> bool {
    const DIST_T startingDistance =
            apdCache->GetRaw(hexMap, mapId, battalionType, false)->Distance(origin, destination);

    return startingDistance != ActionPointDistances::IMPOSSIBLE;
}

// Returns the *non*-water tiles from which you could bridge/freeze water to allow a unit to cross
// from origin to destination
auto CrossingStartLocations(const HexMap *hexMap, const CoordsSet &waterCrossingTiles)
        -> CoordsSet {
    CoordsSet adjacentToWaterCrossingTiles(hexMap);

    for (const Coords &waterTile : waterCrossingTiles) {
        AssertValid(waterTile, hexMap);
        auto adjacentCoords = HexMapUtils::GetAdjacentCoords(hexMap, waterTile);
        for (const auto &adj : adjacentCoords) {
            AssertValid(adj, hexMap);
            if (adjacentToWaterCrossingTiles.Contains(adj)) continue;

            const auto &terr = GetTerrain(hexMap, adj);
            if (IsWater(terr->type())) continue;
            if (terr->type() == net::eagle0::shardok::storage::fb::Terrain_::Type_MOUNTAIN)
                continue;

            adjacentToWaterCrossingTiles.Add(adj);
        }
    }

    return adjacentToWaterCrossingTiles;
}

auto WaterCrossingTiles(
        const Coords &origin,
        const Coords &destination,
        const HexMap *hexMap,
        const APDCache &apdCache,
        const BattalionTypeSPtr &battalionType) -> CoordsSet {
    AssertValid(origin, hexMap);
    AssertValid(destination, hexMap);
    auto returnCoords = CoordsSet(hexMap);

    fb::HexMapW mapCopy = fb::CopyHexMap(hexMap);
    const int mapSize = hexMap->row_count() * hexMap->column_count();

    // For each tile on the map, check if and only if:
    //  1) the tile is water
    //  2) the tile does not already have a bridge (that is not on fire)
    //  3) the tile is not already frozen
    for (int index = 0; index < mapSize; index++) {
        const auto startTerrain = *hexMap->terrain()->Get(index);

        if (!IsWater(startTerrain.type())) continue;

        const auto modifier = startTerrain.modifier();
        if (modifier.bridge().present() && !modifier.fire().present()) continue;
        if (modifier.ice().present() && !modifier.fire().present()) continue;

        // Now try adding a bridge to the tile to see if it helps
        mapCopy->mutable_terrain()
                ->GetMutableObject(index)
                ->mutable_modifier()
                .mutable_bridge()
                .mutate_present(true);
        mapCopy->mutable_terrain()
                ->GetMutableObject(index)
                ->mutable_modifier()
                .mutable_fire()
                .mutate_present(false);

        auto hash = ActionPointDistancesCache::GetMapId(mapCopy);

        if (const auto *distances = apdCache->GetRaw(mapCopy, hash, battalionType, false);
            distances->Distance(origin, destination) != ActionPointDistances::IMPOSSIBLE) {
            returnCoords.Add(index / hexMap->column_count(), index % hexMap->column_count());
        }

        // Undo the new bridge for the next iteration of the loop
        mapCopy->mutable_terrain()
                ->GetMutableObject(index)
                ->mutable_modifier()
                .mutable_bridge()
                .mutate_present(false);
    }

    return returnCoords;
}

// Returns the set of tiles that the attacker should try to approach in order to bridge/freeze
auto IntendedCrossingStarts(
        const GameStateW &gameState,
        const vector<UnitId> &unitIdsCreatingCrossing,
        const CoordsSet &tilesToStartCrossingFrom,
        const APDCache &apdCache,
        const BattalionTypeGetter &battalionTypeGetter) -> CoordsSet {
    CoordsSet intendedCrossingStarts(gameState->hex_map());
    const MapId mapId = apdCache->GetMapId(gameState->hex_map());
    for (const UnitId uid : unitIdsCreatingCrossing) {
        const Unit *unit = gameState->units()->Get(uid);
        const Coords &location = unit->location();
        const auto &battalionType = battalionTypeGetter(unit->battalion().type());
        const auto *apd = apdCache->GetRaw(gameState->hex_map(), mapId, battalionType, false);

        if (location.row() >= 0) {
            Coords intended =
                    MinimumDistanceAndTarget(apd, location, tilesToStartCrossingFrom).coords;
            if (intended.row() >= 0) { intendedCrossingStarts.Add(intended); }
        }
    }

    return intendedCrossingStarts;
}

using Unit = net::eagle0::shardok::storage::fb::Unit;

constexpr double kNoRequiredCrossingScore = std::numeric_limits<double>::max();
constexpr double kNoCrossingCreatorsScore = std::numeric_limits<double>::min();

auto WaterCrossingScore(
        const PlayerId playerId,
        const BattalionTypeGetter &battalionTypeGetter,
        const GameStateW &gameState,
        const CoordsSet &castleCoords,
        const CoordsSet &startCrossingFrom,
        const APDCache &apdCache) -> double {
    uint32_t castleClaimCount = 0;
    for (const auto *unit : *gameState->units()) {
        if (unit->player_id() != playerId) continue;
        const auto status = unit->status();
        if (status != net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT &&
            status != net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT)
            continue;
        if (!unit->has_attached_hero()) continue;

        ++castleClaimCount;
    }

    CoordsSet destinations = castleCoords;
    if (castleClaimCount < castleCoords.size()) {
        destinations = CoordsSet(gameState->hex_map());
        for (const auto *enemyUnit : *gameState->units()) {
            if (enemyUnit->player_id() == playerId) continue;
            const auto status = enemyUnit->status();
            if (status != net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT) continue;
            AssertValid(enemyUnit->location(), gameState->hex_map());
            destinations.Add(enemyUnit->location());
        }
    }

    const auto unitIdsRequiringCrossing = UnitIdsRequiringWaterCrossing(
            gameState,
            playerId,
            castleCoords,
            apdCache,
            battalionTypeGetter);
    if (unitIdsRequiringCrossing.empty()) return kNoRequiredCrossingScore;

    const auto unitIdsCreatingCrossing =
            UnitIdsToCreateWaterCrossing(gameState, playerId, battalionTypeGetter);
    if (unitIdsCreatingCrossing.empty()) return kNoCrossingCreatorsScore;

    double totalScore = 0;

    const auto mapId = ActionPointDistancesCache::GetMapId(gameState->hex_map());

    // First put a big penalty on the distance for units that can create a crossing
    for (const UnitId uid : unitIdsCreatingCrossing) {
        const Unit *unit = gameState->units()->Get(uid);
        const auto &battalionType = battalionTypeGetter(unit->battalion().type());
        Coords location = unit->location();

        int thisDistance;
        if (location.row() < 0) thisDistance = 1000;
        else {
            const auto *apd = apdCache->GetRaw(gameState->hex_map(), mapId, battalionType, false);

            thisDistance = MinimumDistance(apd, location, startCrossingFrom);
        }

        totalScore -= thisDistance * 100.0;
    }

    // Now a smaller penalty for distance for units that need to cross, except if they block -- then
    // a large penalty
    for (const UnitId uid : unitIdsRequiringCrossing) {
        // If this unit ID can also create a crossing, we already handled it
        if (std::ranges::contains(unitIdsCreatingCrossing, uid)) continue;

        const Unit *unit = gameState->units()->Get(uid);
        const auto &battalionType = battalionTypeGetter(unit->battalion().type());
        Coords location = unit->location();
        const auto *apd = apdCache->GetRaw(gameState->hex_map(), mapId, battalionType, false);

        int thisDistance;
        if (location.row() < 0) thisDistance = 1000;
        else { thisDistance = MinimumDistance(apd, location, startCrossingFrom); }

        bool targetBlocks = false;
        // If we're not capable of creating a crossing, don't get in the way of somebody that is.
        for (const UnitId crossingUid : unitIdsCreatingCrossing) {
            const auto *crossingCapableUnit = gameState->units()->Get(crossingUid);

            // Don't check for units that aren't yet placed
            if (crossingCapableUnit->location().row() < 0) continue;
            AssertValid(crossingCapableUnit->location(), gameState->hex_map());

            if (thisDistance <
                MinimumDistance(apd, crossingCapableUnit->location(), startCrossingFrom)) {
                targetBlocks = true;
                break;
            }
        }
        if (targetBlocks) continue;

        totalScore -= thisDistance;
    }

    return totalScore;
}

}  // namespace shardok
