#include "UnlockState.h"

#include <algorithm>

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

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 recipes newly implicitly unlocked by
    // awarding this whole group. Seed the hypothetical explicit-unlock sets with
    // every grant (ship + module materials via step 1a, recipe outputs via step
    // 1b), then diff against the current implicit set.
    std::set<std::string> hypotheticalShipIds            = getUnlockedShipSchematicIds();
    std::set<std::string> hypotheticalModuleIds          = getUnlockedModuleSchematicIds();
    std::set<std::string> hypotheticalRecipeSchematicIds = m_unlockedRecipeSchematicIds;
    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); }

    const UnlockedSets hypothetical = computeUnlockedSets(
        hypotheticalShipIds, hypotheticalModuleIds, hypotheticalRecipeSchematicIds);
    option.newlyUnlockedRecipeIds = computeNewlyUnlockedRecipeIds(hypothetical);

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

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 RecipeOutput& out : def.outputs)
            {
                result.itemIds.insert(out.item);
            }
        }
    }

    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;
            }
            // Skip a gated assembler recipe (granted by an unlock group) whose
            // group has not yet been awarded (REQ-LOCK-IMPLICIT step 2).
            if (recipe.building == BuildingType::Assembler
                && m_grantedRecipeIds.count(recipe.id) > 0
                && unlockedRecipeSchematicIds.count(recipe.id) == 0)
            {
                continue;
            }
            bool producesUnlocked = false;
            for (const RecipeOutput& out : recipe.outputs)
            {
                if (result.itemIds.count(out.item) > 0)
                {
                    producesUnlocked = true;
                    break;
                }
            }
            if (!producesUnlocked) { continue; }

            if (recipe.building == BuildingType::Miner
                || recipe.building == BuildingType::Assembler)
            {
                result.recipeIds.insert(recipe.id);
            }
            for (const RecipeIngredient& ing : recipe.inputs)
            {
                if (result.itemIds.insert(ing.item).second)
                {
                    changed = true;
                }
            }
        }
    }

    return result;
}

std::vector<std::string> UnlockState::computeNewlyUnlockedRecipeIds(const UnlockedSets& hypothetical) const
{
    std::vector<std::string> recipeIds;
    for (const std::string& recipeId : hypothetical.recipeIds)
    {
        if (m_unlockedRecipeIds.count(recipeId) > 0) { continue; }
        recipeIds.push_back(recipeId);
    }
    std::sort(recipeIds.begin(), recipeIds.end(),
              [](const std::string& lhs, const std::string& rhs)
              {
                  return toDisplayName(lhs) < toDisplayName(rhs);
              });
    return recipeIds;
}

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