From e460b55eb4df8fe91baec2800dc3b96851de8cce Mon Sep 17 00:00:00 2001
From: Soar Qin <soarchin@gmail.com>
Date: Tue, 13 May 2025 22:51:07 +0800
Subject: [PATCH] UXAssist: fix tech tree

---
 UXAssist/Patches/TechPatch.cs | 38 ++++++++++++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/UXAssist/Patches/TechPatch.cs b/UXAssist/Patches/TechPatch.cs
index 223d961..998c170 100644
--- a/UXAssist/Patches/TechPatch.cs
+++ b/UXAssist/Patches/TechPatch.cs
@@ -133,9 +133,9 @@ public static class TechPatch
             {
                 (int, int)[] techListToDisable =
                 [
-                    // Combustible Unit, Explosive Unit, Crystal Explosive Unit
-                    // 燃烧单元，爆破单元，晶石爆破单元
-                    (1802, 1804),
+                    // Explosive Unit, Crystal Explosive Unit
+                    // 爆破单元，晶石爆破单元
+                    (1803, 1804),
                     // Implosion Cannon
                     // 聚爆加农炮
                     (1807, 1807),
@@ -244,16 +244,44 @@ public static class TechPatch
                 _originTechPosts.Clear();
                 _protoPatched = false;
             }
-            foreach (var item in UIRoot.instance.uiGame.techTree.nodes)
+            var nodes = UIRoot.instance.uiGame.techTree.nodes;
+            var history = GameMain.history;
+            foreach (var item in nodes)
             {
                 var node = item.Value;
                 foreach (var arrow in node.arrows)
                 {
                     arrow.gameObject.SetActive(false);
                 }
-                var active = node.techProto.postTechArray.Length > 0;
+                var postTechArray = node.techProto.postTechArray;
+                var active = postTechArray.Length > 0;
                 node.connGroup.gameObject.SetActive(active);
                 node.connImages = active ? node.connGroup.GetComponentsInChildren<Image>(true) : [];
+                if (!active) continue;
+                bool onlyOneOutput = postTechArray.Length == 1;
+                for (var i = 0; i < postTechArray.Length; i++)
+                {
+                    var postTech = postTechArray[i];
+                    if (!nodes.ContainsKey(postTech.ID)) continue;
+                    bool itemUnlocked = true;
+                    for (var n = 0; n < postTech.PreItem.Length; n++)
+                    {
+                        if (history.ItemUnlocked(postTech.PreItem[n])) continue;
+                        itemUnlocked = false;
+                        break;
+                    }
+                    bool techUnlocked = history.TechUnlocked(postTech.ID);
+                    node.arrows[i].gameObject.SetActive(itemUnlocked || techUnlocked);
+                    if (onlyOneOutput && !itemUnlocked && !techUnlocked)
+                    {
+                        node.outputArrow.enabled = false;
+                        node.outputLine.gameObject.SetActive(false);
+                    }
+                    else
+                    {
+                        node.outputLine.gameObject.SetActive(true);
+                    }
+                }
             }
         }
     }