IGP-Fan-Reference/Services/Immortal/TechTreeService.cs

using Model.Entity.Data;
using Model.TechTree;
using Model.Types;

namespace Services.Immortal;

public class TechTreeService
{
    private TechTreeGraphModel? _graph;

    public TechTreeGraphModel BuildGraph(string? factionFilter = null)
    {
        _graph = new TechTreeGraphModel();
        var entities = EntityData.Get();
        var factions = new HashSet<string>();

        var factionEntityIds = new HashSet<string>();

        foreach (var kvp in entities)
        {
            var entity = kvp.Value;
            var entityFaction = entity.Faction()?.Faction;

            if (factionFilter != null)
            {
                if (entityFaction == null) continue;
                if (!entityFaction.Equals(factionFilter, StringComparison.OrdinalIgnoreCase)
                    && !entityFaction.Equals(DataType.FACTION_Neutral, StringComparison.OrdinalIgnoreCase))
                    continue;

                factionEntityIds.Add(kvp.Key);
            }

            if (entityFaction != null)
                factions.Add(entityFaction);

            var node = new TechTreeNodeModel
            {
                Id = kvp.Key,
                Name = entity.Info()?.Name ?? kvp.Key,
                EntityType = entity.EntityType,
                Faction = entityFaction ?? "",
                Descriptive = entity.Descriptive
            };

            _graph.Nodes.Add(node);
        }

        foreach (var kvp in entities)
        {
            var entity = kvp.Value;
            var entityId = kvp.Key;

            // ProducedBy -> the building that produces this entity
            var production = entity.Production();
            if (production?.ProducedBy != null && entities.ContainsKey(production.ProducedBy))
                if (factionFilter == null ||
                    (factionEntityIds.Contains(entityId) && factionEntityIds.Contains(production.ProducedBy)))
                    _graph.Edges.Add(new TechTreeEdgeModel
                    {
                        SourceId = production.ProducedBy,
                        TargetId = entityId,
                        EdgeType = TechTreeEdgeType.Produces
                    });

            // Requirements
            foreach (var req in entity.Requirements())
            {
                if (!entities.ContainsKey(req.Id)) continue;

                if (factionFilter != null &&
                    (!factionEntityIds.Contains(entityId) || !factionEntityIds.Contains(req.Id)))
                    continue;

                string edgeType;
                if (req.Requirement == RequirementType.Production_Building)
                    edgeType = TechTreeEdgeType.RequiresProduction;
                else if (req.Requirement == RequirementType.Research_Building)
                    edgeType = TechTreeEdgeType.RequiresResearch;
                else if (req.Requirement == RequirementType.Research_Upgrade)
                    edgeType = TechTreeEdgeType.RequiresResearch;
                else if (req.Requirement == RequirementType.Morph)
                    edgeType = TechTreeEdgeType.Morph;
                else
                    edgeType = TechTreeEdgeType.RequiresProduction;

                _graph.Edges.Add(new TechTreeEdgeModel
                {
                    SourceId = req.Id,
                    TargetId = entityId,
                    EdgeType = edgeType
                });
            }

            // Upgrades
            foreach (var upgrade in entity.IdUpgrades())
            {
                if (!entities.ContainsKey(upgrade.Id)) continue;

                if (factionFilter != null &&
                    (!factionEntityIds.Contains(entityId) || !factionEntityIds.Contains(upgrade.Id)))
                    continue;

                _graph.Edges.Add(new TechTreeEdgeModel
                {
                    SourceId = upgrade.Id,
                    TargetId = entityId,
                    EdgeType = TechTreeEdgeType.Upgrades
                });
            }
        }

        // Build reverse index: what does each entity unlock?
        foreach (var edge in _graph.Edges)
        {
            if (!_graph.Unlocks.ContainsKey(edge.SourceId))
                _graph.Unlocks[edge.SourceId] = new List<string>();

            if (!_graph.Unlocks[edge.SourceId].Contains(edge.TargetId))
                _graph.Unlocks[edge.SourceId].Add(edge.TargetId);
        }

        // Compute layers (BFS from root nodes with no incoming edges)
        ComputeLayers();

        return _graph;
    }

    public List<TechTreeNodeModel> GetUpgradePath(string entityId)
    {
        var graph = GetGraph();
        var path = new List<TechTreeNodeModel>();
        var visited = new HashSet<string>();
        var queue = new Queue<string>();
        queue.Enqueue(entityId);

        while (queue.Count > 0)
        {
            var current = queue.Dequeue();
            if (!visited.Add(current)) continue;

            var node = graph.Nodes.FirstOrDefault(n => n.Id == current);
            if (node != null) path.Add(node);

            var upgradeEdges = graph.Edges
                .Where(e => e.SourceId == current && e.EdgeType == TechTreeEdgeType.Upgrades);
            foreach (var edge in upgradeEdges)
                queue.Enqueue(edge.TargetId);
        }

        return path;
    }

    public List<TechTreeNodeModel> GetPrerequisites(string entityId)
    {
        var graph = GetGraph();
        var prereqs = new List<TechTreeNodeModel>();

        var incomingEdges = graph.Edges
            .Where(e => e.TargetId == entityId && e.EdgeType != TechTreeEdgeType.UpgradedBy);

        foreach (var edge in incomingEdges)
        {
            var node = graph.Nodes.FirstOrDefault(n => n.Id == edge.SourceId);
            if (node != null) prereqs.Add(node);
        }

        return prereqs;
    }

    public List<TechTreeNodeModel> GetUnlocks(string entityId)
    {
        var graph = GetGraph();
        var unlocks = new List<TechTreeNodeModel>();

        if (!graph.Unlocks.TryGetValue(entityId, out var unlockIds))
            return unlocks;

        foreach (var id in unlockIds)
        {
            var node = graph.Nodes.FirstOrDefault(n => n.Id == id);
            if (node != null) unlocks.Add(node);
        }

        return unlocks;
    }

    public List<string> GetFactions()
    {
        return GetGraph().Nodes
            .Where(n => !string.IsNullOrEmpty(n.Faction))
            .Select(n => n.Faction)
            .Distinct()
            .ToList();
    }

    public TechTreeGraphModel GetFilteredGraph(string? faction, string? searchText, string? highlightEntity)
    {
        var graph = GetGraph();

        if (faction == null && string.IsNullOrEmpty(searchText) && highlightEntity == null)
            return graph;

        // Build from scratch with faction filter
        return BuildGraph(faction);
    }

    private TechTreeGraphModel GetGraph()
    {
        _graph ??= BuildGraph();
        return _graph;
    }

    private void ComputeLayers()
    {
        var inDegree = new Dictionary<string, int>();
        var adjacency = new Dictionary<string, List<string>>();

        foreach (var node in _graph!.Nodes)
        {
            inDegree[node.Id] = 0;
            adjacency[node.Id] = new List<string>();
        }

        foreach (var edge in _graph.Edges)
        {
            if (adjacency.ContainsKey(edge.SourceId)) adjacency[edge.SourceId].Add(edge.TargetId);

            if (inDegree.ContainsKey(edge.TargetId)) inDegree[edge.TargetId]++;
        }

        var queue = new Queue<string>();
        foreach (var kvp in inDegree)
            if (kvp.Value == 0)
                queue.Enqueue(kvp.Key);

        var layers = new Dictionary<string, int>();

        while (queue.Count > 0)
        {
            var current = queue.Dequeue();
            var currentLayer = layers.TryGetValue(current, out var l) ? l : 0;

            foreach (var next in adjacency.GetValueOrDefault(current, new List<string>()))
            {
                var nextLayer = currentLayer + 1;
                if (!layers.ContainsKey(next) || layers[next] < nextLayer) layers[next] = nextLayer;

                inDegree[next]--;
                if (inDegree[next] == 0) queue.Enqueue(next);
            }
        }

        foreach (var node in _graph.Nodes) node.Layer = layers.TryGetValue(node.Id, out var layer) ? layer : 0;
    }
}