#include "UnlockState.h"

#include <algorithm>
#include <functional>

#include "DisplayName.h"
#include "StateChecksum.h"

namespace
{

// Both preview lists read as a list of names, so they are ordered as one (REQ-DEF-SCHEMATIC-DROP).
void sortByDisplayName(std::vector<PreviewedRecipe>& recipes)
{
    std::sort(recipes.begin(), recipes.end(),
              [](const PreviewedRecipe& lhs, const PreviewedRecipe& rhs)
              {
                  return toDisplayName(lhs.recipeId) < toDisplayName(rhs.recipeId);
              });
}

}  // namespace

UnlockState::UnlockState(const GameConfig& config)
    : m_config(config)
{
}

void UnlockState::initializeUnlockState()
{
    // Cache the ids granted by some unlock group (REQ-LOCK-EXPLICIT); an item
    // starts locked iff it is granted by a group.
    m_grantedShipIds.clear();
    m_grantedModuleIds.clear();
    m_grantedBuildingIds.clear();
    m_grantedRecipeIds.clear();
    for (const UnlockGroupDef& group : m_config.unlocks.groups)
    {
        m_grantedShipIds.insert(group.ships.begin(), group.ships.end());
        m_grantedModuleIds.insert(group.modules.begin(), group.modules.end());
        m_grantedBuildingIds.insert(group.buildings.begin(), group.buildings.end());
        m_grantedRecipeIds.insert(group.recipes.begin(), group.recipes.end());
    }

    m_awardedUnlockGroupIds.clear();

    m_schematicLevels.clear();
    for (const ShipDef& def : m_config.ships.ships)
    {
        SchematicState state;
        state.unlocked = (m_grantedShipIds.count(def.id) == 0);
        m_schematicLevels[def.id] = state;
    }

    m_moduleSchematicLevels.clear();
    for (const ModuleDef& def : m_config.modules.modules)
    {
        SchematicState state;
        state.unlocked = (m_grantedModuleIds.count(def.id) == 0);
        m_moduleSchematicLevels[def.id] = state;
    }

    m_buildingLevels.clear();
    for (const BuildingDef& def : m_config.buildings.buildings)
    {
        SchematicState state;
        state.unlocked = (m_grantedBuildingIds.count(def.id) == 0);
        m_buildingLevels[def.id] = state;
    }

    // Gated assembler recipes start locked; unlocked_at_start recipes are handled
    // in the REQ-LOCK-IMPLICIT traversal, not tracked here.
    m_unlockedRecipeSchematicIds.clear();

    recomputeUnlocked();
}

bool UnlockState::isSchematicUnlocked(const std::string& shipId) const
{
    const std::map<std::string, SchematicState>::const_iterator it =
        m_schematicLevels.find(shipId);
    if (it == m_schematicLevels.end())
    {
        return false;
    }
    return it->second.unlocked;
}

bool UnlockState::isModuleSchematicUnlocked(const std::string& moduleId) const
{
    const std::map<std::string, SchematicState>::const_iterator it =
        m_moduleSchematicLevels.find(moduleId);
    if (it == m_moduleSchematicLevels.end())
    {
        return false;
    }
    return it->second.unlocked;
}

bool UnlockState::isRecipeUnlocked(const std::string& recipeId) const
{
    return m_unlockedRecipeIds.count(recipeId) > 0;
}

bool UnlockState::isItemUnlocked(const std::string& itemId) const
{
    return m_unlockedItemIds.count(itemId) > 0;
}

bool UnlockState::isBuildingUnlocked(BuildingType type) const
{
    const BuildingDef* def = m_config.buildings.findBuildingDef(type);
    if (def == nullptr)
    {
        // Types without a config entry (e.g. HQ, defence stations) are unrestricted.
        return true;
    }
    const std::map<std::string, SchematicState>::const_iterator it =
        m_buildingLevels.find(def->id);
    return it == m_buildingLevels.end() ? true : it->second.unlocked;
}

bool UnlockState::isUnlockGroupAwarded(const std::string& groupId) const
{
    return m_awardedUnlockGroupIds.count(groupId) > 0;
}

bool UnlockState::prerequisitesSatisfied(const std::vector<std::string>& requiredGroupIds) const
{
    // A prerequisite is satisfied only once the named unlock group has been
    // awarded (REQ-LOCK-PREREQ).
    for (const std::string& groupId : requiredGroupIds)
    {
        if (m_awardedUnlockGroupIds.count(groupId) == 0) { return false; }
    }
    return true;
}

SchematicChoiceOption UnlockState::makeUnlockOption(const UnlockGroupDef& group) const
{
    SchematicChoiceOption option;
    option.isArtifact    = false;
    option.unlockGroupId = group.id;
    option.displayName   = toDisplayName(group.id);

    for (const std::string& id : group.ships)
    {
        option.grantedItems.push_back({SchematicType::Ship, id, toDisplayName(id)});
    }
    for (const std::string& id : group.modules)
    {
        option.grantedItems.push_back({SchematicType::Module, id, toDisplayName(id)});
    }
    for (const std::string& id : group.buildings)
    {
        option.grantedItems.push_back({SchematicType::Building, id, toDisplayName(id)});
    }
    for (const std::string& id : group.recipes)
    {
        option.grantedItems.push_back({SchematicType::Recipe, id, toDisplayName(id)});
    }

    // REQ-DEF-SCHEMATIC-DROP: preview what awarding this whole group would add. Seed the
    // hypothetical explicit-unlock sets with every grant (ship + module materials via
    // step 1a, recipe outputs via step 1b, and the granted buildings, which no traversal
    // covers), then diff what is usable before against what would be usable after.
    std::set<std::string> hypotheticalShipIds            = getUnlockedShipSchematicIds();
    std::set<std::string> hypotheticalModuleIds          = getUnlockedModuleSchematicIds();
    std::set<std::string> hypotheticalRecipeSchematicIds = m_unlockedRecipeSchematicIds;
    std::set<std::string> hypotheticalBuildingIds        = getUnlockedBuildingIds();
    for (const std::string& id : group.ships)     { hypotheticalShipIds.insert(id); }
    for (const std::string& id : group.modules)   { hypotheticalModuleIds.insert(id); }
    for (const std::string& id : group.recipes)   { hypotheticalRecipeSchematicIds.insert(id); }
    for (const std::string& id : group.buildings) { hypotheticalBuildingIds.insert(id); }

    UnlockedSets current;
    current.itemIds   = m_unlockedItemIds;
    current.recipeIds = m_unlockedRecipeIds;

    const UnlockedSets hypothetical = computeUnlockedSets(
        hypotheticalShipIds, hypotheticalModuleIds, hypotheticalRecipeSchematicIds);

    const std::map<std::string, std::vector<int>> usableBefore =
        computeUsableRecipeGroups(current, getUnlockedBuildingIds());
    const std::map<std::string, std::vector<int>> usableAfter =
        computeUsableRecipeGroups(hypothetical, hypotheticalBuildingIds);

    for (const std::pair<const std::string, std::vector<int>>& entry : usableAfter)
    {
        const std::map<std::string, std::vector<int>>::const_iterator before =
            usableBefore.find(entry.first);
        if (before == usableBefore.end())
        {
            option.newlyUsableRecipes.push_back(PreviewedRecipe{entry.first, entry.second});
        }
        else if (before->second.size() < entry.second.size())
        {
            // Nothing is ever taken away -- unlocking only ever adds -- so a group count
            // that grew means groups were gained, and the after list contains the before
            // list. Comparing counts is therefore the whole test.
            option.upgradedRecipes.push_back(PreviewedRecipe{entry.first, entry.second});
        }
    }
    sortByDisplayName(option.newlyUsableRecipes);
    sortByDisplayName(option.upgradedRecipes);

    return option;
}

void UnlockState::awardUnlockGroup(const SchematicChoiceOption& chosen)
{
    // Award the whole unlock group (REQ-DEF-SCHEMATIC-DROP): unlock every granted
    // ship, module, building, and assembler recipe at once.
    m_awardedUnlockGroupIds.insert(chosen.unlockGroupId);
    for (const GrantedSchematic& grant : chosen.grantedItems)
    {
        switch (grant.type)
        {
        case SchematicType::Ship:     m_schematicLevels.at(grant.id).unlocked = true;       break;
        case SchematicType::Module:   m_moduleSchematicLevels.at(grant.id).unlocked = true; break;
        case SchematicType::Building: m_buildingLevels.at(grant.id).unlocked = true;        break;
        case SchematicType::Recipe:   m_unlockedRecipeSchematicIds.insert(grant.id);        break;
        }
    }

    recomputeUnlocked();
}

// ---------------------------------------------------------------------------
// Implicit unlock computation (REQ-LOCK-IMPLICIT)
// ---------------------------------------------------------------------------

void UnlockState::recomputeUnlocked()
{
    const UnlockedSets result = computeUnlockedSets(
        getUnlockedShipSchematicIds(), getUnlockedModuleSchematicIds(), m_unlockedRecipeSchematicIds);
    m_unlockedItemIds   = result.itemIds;
    m_unlockedRecipeIds = result.recipeIds;
}

std::set<std::string> UnlockState::getUnlockedShipSchematicIds() const
{
    std::set<std::string> ids;
    for (const auto& [id, state] : m_schematicLevels)
    {
        if (state.unlocked) { ids.insert(id); }
    }
    return ids;
}

std::set<std::string> UnlockState::getUnlockedModuleSchematicIds() const
{
    std::set<std::string> ids;
    for (const auto& [id, state] : m_moduleSchematicLevels)
    {
        if (state.unlocked) { ids.insert(id); }
    }
    return ids;
}

std::set<std::string> UnlockState::getUnlockedBuildingIds() const
{
    std::set<std::string> ids;
    for (const auto& [id, state] : m_buildingLevels)
    {
        if (state.unlocked) { ids.insert(id); }
    }
    return ids;
}

UnlockState::UnlockedSets UnlockState::computeUnlockedSets(
    const std::set<std::string>& unlockedShipSchematicIds,
    const std::set<std::string>& unlockedModuleSchematicIds,
    const std::set<std::string>& unlockedRecipeSchematicIds) const
{
    UnlockedSets result;

    for (const ShipDef& def : m_config.ships.ships)
    {
        if (unlockedShipSchematicIds.count(def.id) == 0) { continue; }
        for (const RecipeIngredient& mat : def.schematic.materials)
        {
            result.itemIds.insert(mat.item);
        }
    }
    for (const ModuleDef& def : m_config.modules.modules)
    {
        if (unlockedModuleSchematicIds.count(def.id) == 0) { continue; }
        for (const RecipeIngredient& mat : def.materials)
        {
            result.itemIds.insert(mat.item);
        }
    }
    for (const RecipeDef& def : m_config.recipes.recipes)
    {
        // An assembler recipe seeds the base set when it is explicitly available:
        // flagged unlocked_at_start (base recipes the graph can't reach), or a
        // gated recipe whose unlock group has been awarded (REQ-LOCK-EXPLICIT).
        if (def.building == BuildingType::Assembler
            && (def.unlockedAtStart || unlockedRecipeSchematicIds.count(def.id) > 0))
        {
            for (const std::string& item : getProducibleItems(def))
            {
                result.itemIds.insert(item);
            }
        }
    }

    // Skips a gated assembler recipe (granted by an unlock group) whose group has not
    // yet been awarded (REQ-LOCK-IMPLICIT step 2, REQ-LOCK-EXPLICIT). This gating is
    // about permission and says nothing about what the recipe yields, so it applies
    // whether the recipe is propagating demand below or being judged runnable further
    // down.
    const std::function<bool(const RecipeDef&)> isGatedOff =
        [this, &unlockedRecipeSchematicIds](const RecipeDef& recipe)
        {
            return recipe.building == BuildingType::Assembler
                && m_grantedRecipeIds.count(recipe.id) > 0
                && unlockedRecipeSchematicIds.count(recipe.id) == 0;
        };

    // Demand propagation: an item is unlocked when something the player can already
    // reach calls for it, so a recipe producing an unlocked item unlocks its inputs.
    // Reprocessing takes no part in this (REQ-BLD-REPROCESSING): being obtainable from
    // a plant must never be what makes an item or its inputs unlocked.
    bool changed = true;
    while (changed)
    {
        changed = false;
        for (const RecipeDef& recipe : m_config.recipes.recipes)
        {
            if (recipe.building != BuildingType::Miner
                && recipe.building != BuildingType::Smelter
                && recipe.building != BuildingType::Assembler)
            {
                continue;
            }
            if (isGatedOff(recipe)) { continue; }

            bool producesUnlocked = false;
            for (const std::string& item : getProducibleItems(recipe))
            {
                if (result.itemIds.count(item) > 0)
                {
                    producesUnlocked = true;
                    break;
                }
            }
            if (!producesUnlocked) { continue; }

            for (const RecipeIngredient& ing : recipe.inputs)
            {
                if (result.itemIds.insert(ing.item).second)
                {
                    changed = true;
                }
            }
        }
    }

    // Which recipes are unlocked, once the item set has settled. One rule for every
    // building that runs a recipe: the player may run it when it is permitted at all and
    // can currently yield something usable -- some output group all of whose items are
    // unlocked (REQ-LOCK-OUTPUT-POOL). Reprocessing is asked the question here without
    // having taken part in the traversal above: what it can yield is decided by the item
    // set, not the other way round.
    const std::function<bool(const std::string&)> isItemUnlockedHere =
        [&result](const std::string& itemId)
        {
            return result.itemIds.count(itemId) > 0;
        };

    for (const RecipeDef& recipe : m_config.recipes.recipes)
    {
        if (recipe.building != BuildingType::Miner
            && recipe.building != BuildingType::Smelter
            && recipe.building != BuildingType::Assembler
            && recipe.building != BuildingType::ReprocessingPlant)
        {
            continue;
        }
        if (isGatedOff(recipe)) { continue; }

        for (const RecipeOutputGroup& group : recipe.outputGroups)
        {
            if (isOutputGroupUnlocked(group, isItemUnlockedHere))
            {
                result.recipeIds.insert(recipe.id);
                break;
            }
        }
    }

    return result;
}

std::map<std::string, std::vector<int>> UnlockState::computeUsableRecipeGroups(
    const UnlockedSets& sets, const std::set<std::string>& unlockedBuildingIds) const
{
    const std::function<bool(const std::string&)> isItemUnlockedHere =
        [&sets](const std::string& itemId) { return sets.itemIds.count(itemId) > 0; };

    std::map<std::string, std::vector<int>> usable;
    for (const RecipeDef& recipe : m_config.recipes.recipes)
    {
        if (sets.recipeIds.count(recipe.id) == 0) { continue; }

        // Building types with no config entry are unrestricted, as isBuildingUnlocked
        // treats them (REQ-LOCK-BUILDING).
        const BuildingDef* def = m_config.buildings.findBuildingDef(recipe.building);
        if (def != nullptr && unlockedBuildingIds.count(def->id) == 0) { continue; }

        std::vector<int> groupIndices;
        for (std::size_t i = 0; i < recipe.outputGroups.size(); ++i)
        {
            if (isOutputGroupUnlocked(recipe.outputGroups[i], isItemUnlockedHere))
            {
                groupIndices.push_back(static_cast<int>(i));
            }
        }
        // An unlocked recipe has at least one yieldable group by REQ-LOCK-IMPLICIT step
        // 4, so this is a consistency check on that rather than a case to handle.
        if (groupIndices.empty()) { continue; }

        usable[recipe.id] = std::move(groupIndices);
    }
    return usable;
}

// ---------------------------------------------------------------------------
// Determinism (see docs/replay_design.md)
// ---------------------------------------------------------------------------

void UnlockState::appendSchematicMap(Hasher& hasher,
                                     const std::map<std::string, SchematicState>& levels)
{
    hasher.append(levels.size());
    for (const std::pair<const std::string, SchematicState>& entry : levels)
    {
        hasher.append(entry.first);
        hasher.append(entry.second.unlocked);
    }
}

void UnlockState::appendStringSet(Hasher& hasher, const std::set<std::string>& ids)
{
    hasher.append(ids.size());
    for (const std::string& id : ids)
    {
        hasher.append(id);
    }
}

void UnlockState::appendChecksum(Hasher& hasher) const
{
    appendSchematicMap(hasher, m_schematicLevels);
    appendSchematicMap(hasher, m_moduleSchematicLevels);
    appendSchematicMap(hasher, m_buildingLevels);
    appendStringSet(hasher, m_awardedUnlockGroupIds);
    appendStringSet(hasher, m_unlockedRecipeSchematicIds);
    appendStringSet(hasher, m_unlockedRecipeIds);
    appendStringSet(hasher, m_unlockedItemIds);
}
