#include "ThreatCostCalculator.h"

#include <limits>
#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;
}

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;
}

double computeRecipeThreat(const RecipeDef& recipe,
                           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;
}

}  // namespace


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

    // Build lookup: output item → non-reprocessing recipes and reprocessing recipes.
    std::map<std::string, std::vector<RecipeRef>> nonReprocessingRecipes;
    std::map<std::string, std::vector<RecipeRef>> reprocessingRecipes;

    for (const RecipeDef& recipe : config.recipes.recipes)
    {
        if (recipe.building == BuildingType::ReprocessingPlant)
        {
            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);
                reprocessingRecipes[out.item].push_back(ref);
            }
        }
        else
        {
            for (const RecipeOutput& out : recipe.outputs)
            {
                RecipeRef ref;
                ref.recipe = &recipe;
                ref.outputItem = out.item;
                ref.outputAmount = out.amount;
                ref.probability = 1.0;
                nonReprocessingRecipes[out.item].push_back(ref);
            }
        }
    }

    // Collect all item names that need resolving.
    std::set<std::string> unresolved;
    for (const std::pair<const std::string, std::vector<RecipeRef>>& entry : nonReprocessingRecipes)
    {
        unresolved.insert(entry.first);
    }
    for (const std::pair<const std::string, std::vector<RecipeRef>>& entry : reprocessingRecipes)
    {
        unresolved.insert(entry.first);
    }

    // Iteratively resolve non-reprocessing items.
    bool progress = true;
    while (progress)
    {
        progress = false;
        std::set<std::string> newlyResolved;
        for (const std::string& item : unresolved)
        {
            std::map<std::string, std::vector<RecipeRef>>::const_iterator it =
                nonReprocessingRecipes.find(item);
            if (it == nonReprocessingRecipes.end())
            {
                continue;
            }

            double maxThreat = -1.0;
            for (const RecipeRef& ref : it->second)
            {
                if (allInputsResolved(*ref.recipe, table.itemThreat))
                {
                    double threat = computeRecipeThreat(*ref.recipe, table.itemThreat);
                    if (threat > maxThreat)
                    {
                        maxThreat = threat;
                    }
                }
            }

            if (maxThreat >= 0.0)
            {
                table.itemThreat[item] = maxThreat;
                newlyResolved.insert(item);
                progress = true;
            }
        }
        for (const std::string& item : newlyResolved)
        {
            unresolved.erase(item);
        }
    }

    // Compute scrap threat (REQ-THREAT-SCRAP): find the ship with the smallest
    // scrap_drop and use its threat cost.
    int minScrapDrop = std::numeric_limits<int>::max();
    const ShipDef* cheapestScrapShip = nullptr;
    for (const ShipDef& def : config.ships.ships)
    {
        if (def.loot.scrapDrop > 0 && def.loot.scrapDrop < minScrapDrop)
        {
            minScrapDrop = def.loot.scrapDrop;
            cheapestScrapShip = &def;
        }
    }

    if (cheapestScrapShip != nullptr)
    {
        double shipThreat = calculateShipThreatCost(table, config,
            cheapestScrapShip->id, cheapestScrapShip->defaultModules);
        table.scrapThreat = shipThreat / minScrapDrop;
    }

    // Resolve reprocessing-only items.
    for (const std::string& item : unresolved)
    {
        std::map<std::string, std::vector<RecipeRef>>::const_iterator it =
            reprocessingRecipes.find(item);
        if (it == reprocessingRecipes.end())
        {
            continue;
        }

        for (const RecipeRef& ref : it->second)
        {
            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 = table.itemThreat.find(item);
            if (existing == table.itemThreat.end() || threat > existing->second)
            {
                table.itemThreat[item] = threat;
            }
        }
    }

    return table;
}


double calculateShipThreatCost(const ThreatCostTable& table,
                               const GameConfig& config,
                               const std::string& shipId,
                               const std::vector<PlacedModule>& modules)
{
    const ShipDef* shipDef = nullptr;
    for (const ShipDef& d : config.ships.ships)
    {
        if (d.id == shipId)
        {
            shipDef = &d;
            break;
        }
    }
    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 = nullptr;
        for (const ModuleDef& d : config.modules.modules)
        {
            if (d.id == pm.moduleId)
            {
                moduleDef = &d;
                break;
            }
        }
        if (moduleDef == nullptr)
        {
            continue;
        }

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

    return threat;
}