#include "ThreatCostCalculator.h"

#include <set>

#include "GameConfig.h"

namespace
{

struct RecipeRef
{
    const RecipeDef* recipe;
    std::string outputItem;
    int outputAmount;
    double probability;
};

double computeMaterialThreat(const ThreatCostTable& table,
                             const std::vector<RecipeIngredient>& materials)
{
    double total = 0.0;
    for (const RecipeIngredient& mat : materials)
    {
        std::map<std::string, double>::const_iterator it = table.itemThreat.find(mat.item);
        if (it != table.itemThreat.end())
        {
            total += it->second * mat.amount;
        }
    }
    return total;
}

// Returns true if every input of the recipe has a resolved threat value.
bool allInputsResolved(const RecipeDef& recipe,
                       const std::map<std::string, double>& resolved)
{
    for (const RecipeIngredient& input : recipe.inputs)
    {
        if (resolved.find(input.item) == resolved.end())
        {
            return false;
        }
    }
    return true;
}

// Computes the raw recipe threat (duration + sum of input threats × amounts),
// divided by the output amount to get the per-unit threat.
double computeRecipeThreatPerUnit(const RecipeDef& recipe,
                                  int outputAmount,
                                  const std::map<std::string, double>& resolved)
{
    double threat = recipe.durationSeconds;
    for (const RecipeIngredient& input : recipe.inputs)
    {
        threat += resolved.at(input.item) * input.amount;
    }
    return threat / static_cast<double>(outputAmount);
}

}  // namespace


ThreatCostTable computeThreatCostTable(const GameConfig& config)
{
    ThreatCostTable table;

    // Scrap threat (REQ-THREAT-SCRAP) is the constant inverse of the scrap-drop
    // conversion (REQ-RES-DEBRIS-DROP): one scrap is worth 1 / scrap_per_threat.
    // Set it up front so reprocessing-only item threats (below) can use it, and so
    // it no longer depends on any ship's threat cost.
    table.scrapThreat = config.world.scrapPerThreat > 0.0
        ? 1.0 / config.world.scrapPerThreat
        : 0.0;

    // -------------------------------------------------------------------------
    // Build per-item recipe lookup tables.
    // -------------------------------------------------------------------------

    // Items that have at least one non-reprocessing recipe that does NOT consume
    // scrap — these items' scrap-consuming recipes are excluded from threat
    // computation (REQ-THREAT-ITEM: scrap-consuming recipes are a fallback only).
    std::set<std::string> scrapFreeItems;

    // nonReprocessingRecipes: item → all eligible non-reprocessing (recipe, output)
    // pairs. Scrap-consuming recipes are collected here temporarily; they are
    // filtered out per item after we know which items have a scrap-free producer.
    struct EligiblePair
    {
        const RecipeDef* recipe;
        int outputAmount;
        bool consumesScrap;
    };
    std::map<std::string, std::vector<EligiblePair>> nonReprocessingCandidates;

    // reprocessingRecipes: item → all reprocessing-recipe refs (probability
    // values are raw from config; we normalize them per-recipe below).
    std::map<std::string, std::vector<RecipeRef>> reprocessingRecipes;

    for (const RecipeDef& recipe : config.recipes.recipes)
    {
        if (recipe.building == BuildingType::ReprocessingPlant)
        {
            // Compute the total weight across all outputs of this reprocessing recipe
            // so we can normalize each output's probability.
            double totalWeight = 0.0;
            for (const RecipeOutput& out : recipe.outputs)
            {
                totalWeight += out.probability.value_or(1.0);
            }
            if (totalWeight <= 0.0)
            {
                continue;
            }

            for (const RecipeOutput& out : recipe.outputs)
            {
                RecipeRef ref;
                ref.recipe = &recipe;
                ref.outputItem = out.item;
                ref.outputAmount = out.amount;
                ref.probability = out.probability.value_or(1.0) / totalWeight;
                reprocessingRecipes[out.item].push_back(ref);
            }
        }
        else
        {
            // Check whether this non-reprocessing recipe consumes scrap.
            bool consumesScrap = false;
            for (const RecipeIngredient& input : recipe.inputs)
            {
                if (input.item == "scrap")
                {
                    consumesScrap = true;
                    break;
                }
            }

            for (const RecipeOutput& out : recipe.outputs)
            {
                if (!consumesScrap)
                {
                    scrapFreeItems.insert(out.item);
                }

                EligiblePair pair;
                pair.recipe = &recipe;
                pair.outputAmount = out.amount;
                pair.consumesScrap = consumesScrap;
                nonReprocessingCandidates[out.item].push_back(pair);
            }
        }
    }

    // Filter nonReprocessingCandidates: for items that have at least one
    // scrap-free producer, drop their scrap-consuming recipes.
    // Build the final per-item list of (recipe, outputAmount) pairs eligible
    // for the max-across-recipes rule (REQ-THREAT-ITEM).
    std::map<std::string, std::vector<std::pair<const RecipeDef*, int>>> eligibleRecipes;
    for (std::map<std::string, std::vector<EligiblePair>>::const_iterator it =
             nonReprocessingCandidates.begin();
         it != nonReprocessingCandidates.end();
         ++it)
    {
        const std::string& item = it->first;
        const std::vector<EligiblePair>& candidates = it->second;
        bool hasScrapFree = (scrapFreeItems.find(item) != scrapFreeItems.end());

        for (const EligiblePair& candidate : candidates)
        {
            if (candidate.consumesScrap && hasScrapFree)
            {
                // Scrap-consuming recipe excluded: item has a scrap-free producer.
                continue;
            }
            eligibleRecipes[item].emplace_back(candidate.recipe, candidate.outputAmount);
        }
    }

    // -------------------------------------------------------------------------
    // Resolution: seed resolved map with scrap, then alternate the
    // non-reprocessing pass and the reprocessing pass to a fixpoint.
    // Fix (8): iterate until neither pass makes progress, rather than running
    // the reprocessing pass once at the end.
    // -------------------------------------------------------------------------

    std::map<std::string, double>& resolved = table.itemThreat;
    resolved["scrap"] = table.scrapThreat;

    // Non-reprocessing resolution pass.
    // Fix (9): commit an item only when EVERY eligible recipe for it is
    // computable, not just the first one that resolves. This ensures a shallow
    // shortcut recipe cannot undercut a deeper base recipe by resolving earlier.
    auto runNonReprocessingPass = [&](bool requireAllRecipes) -> bool
    {
        bool progress = false;
        std::map<std::string, double> newValues;

        for (std::map<std::string, std::vector<std::pair<const RecipeDef*, int>>>::const_iterator
                 it = eligibleRecipes.begin();
             it != eligibleRecipes.end();
             ++it)
        {
            const std::string& item = it->first;
            if (resolved.find(item) != resolved.end())
            {
                continue;
            }

            const std::vector<std::pair<const RecipeDef*, int>>& pairs = it->second;

            bool allComputable = true;
            double maxThreat = -1.0;
            for (const std::pair<const RecipeDef*, int>& pair : pairs)
            {
                if (!allInputsResolved(*pair.first, resolved))
                {
                    allComputable = false;
                    if (requireAllRecipes)
                    {
                        break;
                    }
                    // In fallback mode: skip this recipe but continue gathering
                    // the computable subset.
                    continue;
                }
                double threat = computeRecipeThreatPerUnit(*pair.first, pair.second, resolved);
                if (threat > maxThreat)
                {
                    maxThreat = threat;
                }
            }

            if (requireAllRecipes && !allComputable)
            {
                continue;
            }
            if (maxThreat >= 0.0)
            {
                newValues[item] = maxThreat;
            }
        }

        for (std::map<std::string, double>::const_iterator it = newValues.begin();
             it != newValues.end();
             ++it)
        {
            resolved[it->first] = it->second;
            progress = true;
        }
        return progress;
    };

    // Reprocessing pass: resolve items produced exclusively by reprocessing.
    // Items that also have a non-reprocessing recipe are skipped here (they are
    // covered by the non-reprocessing pass or do not need the reprocessing path).
    auto runReprocessingPass = [&]() -> bool
    {
        bool progress = false;
        for (std::map<std::string, std::vector<RecipeRef>>::const_iterator it =
                 reprocessingRecipes.begin();
             it != reprocessingRecipes.end();
             ++it)
        {
            const std::string& item = it->first;
            if (resolved.find(item) != resolved.end())
            {
                continue;
            }
            // Reprocessing defines an item's threat only when nothing else
            // produces it (REQ-THREAT-ITEM).
            if (scrapFreeItems.find(item) != scrapFreeItems.end())
            {
                continue;
            }
            // Also skip items covered by eligible (non-reprocessing) recipes.
            if (eligibleRecipes.find(item) != eligibleRecipes.end())
            {
                continue;
            }

            for (const RecipeRef& ref : it->second)
            {
                // Sum all scrap inputs for this reprocessing recipe.
                int scrapPerCycle = 0;
                for (const RecipeIngredient& input : ref.recipe->inputs)
                {
                    scrapPerCycle += input.amount;
                }

                double threat = (table.scrapThreat * scrapPerCycle
                                 + ref.recipe->durationSeconds) / ref.probability;

                std::map<std::string, double>::iterator existing = resolved.find(item);
                if (existing == resolved.end() || threat > existing->second)
                {
                    resolved[item] = threat;
                    progress = true;
                }
            }
        }
        return progress;
    };

    // Main fixpoint loop: alternate non-reprocessing and reprocessing passes
    // until neither makes any progress (fix 8).
    bool anyProgress = true;
    while (anyProgress)
    {
        anyProgress = runNonReprocessingPass(true);
        anyProgress = runReprocessingPass() || anyProgress;
    }

    // Deadlock fallback: if any items remain unresolved due to recipe cycles,
    // fall back to committing with the max over the currently computable subset
    // of recipes (fix 9, deadlock guard — same approach as threat_report.py's
    // require_all_recipes=False mode).
    anyProgress = true;
    while (anyProgress)
    {
        anyProgress = runNonReprocessingPass(false);
        anyProgress = runReprocessingPass() || anyProgress;
    }

    // Remove the sentinel scrap entry — scrapThreat is already stored on the
    // table struct; having it in itemThreat would confuse callers iterating items.
    resolved.erase("scrap");

    return table;
}


double calculateShipThreatCost(const ThreatCostTable& table,
                               const GameConfig& config,
                               const std::string& shipId,
                               const std::vector<PlacedModule>& modules)
{
    const ShipDef* shipDef = config.ships.findShipDef(shipId);
    if (shipDef == nullptr)
    {
        return 0.0;
    }

    double threat = shipDef->schematic.productionTimeSeconds;

    // Add material threat for ship base materials.
    threat += computeMaterialThreat(table, shipDef->schematic.materials);

    // Add module production times and material threats.
    for (const PlacedModule& pm : modules)
    {
        const ModuleDef* moduleDef = config.modules.findModuleDef(pm.moduleId);
        if (moduleDef == nullptr)
        {
            continue;
        }

        threat += moduleDef->productionTimeSeconds;
        threat += computeMaterialThreat(table, moduleDef->materials);
    }

    return threat;
}