From d4110d5cebfeb1812411569e84cda98e1061e5d4 Mon Sep 17 00:00:00 2001
From: Benjamin Davies <bentendavies@gmail.com>
Date: Sun, 27 Aug 2023 00:03:40 +1200
Subject: [PATCH] Add associations between entities

---
 .../command/CommandAssociateRelationship.java | 29 ++++++++++++++++--
 .../command/CommandCreateAttribute.java       | 30 ++++++++++++++-----
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/src/net/sourceforge/plantuml/cheneer/command/CommandAssociateRelationship.java b/src/net/sourceforge/plantuml/cheneer/command/CommandAssociateRelationship.java
index 01d2b1e24..9272b0d0b 100644
--- a/src/net/sourceforge/plantuml/cheneer/command/CommandAssociateRelationship.java
+++ b/src/net/sourceforge/plantuml/cheneer/command/CommandAssociateRelationship.java
@@ -35,10 +35,17 @@
  */
 package net.sourceforge.plantuml.cheneer.command;
 
+import net.sourceforge.plantuml.abel.Entity;
+import net.sourceforge.plantuml.abel.Link;
+import net.sourceforge.plantuml.abel.LinkArg;
 import net.sourceforge.plantuml.cheneer.ChenEerDiagram;
 import net.sourceforge.plantuml.command.CommandExecutionResult;
 import net.sourceforge.plantuml.command.SingleLineCommand2;
+import net.sourceforge.plantuml.decoration.LinkDecor;
+import net.sourceforge.plantuml.decoration.LinkType;
 import net.sourceforge.plantuml.klimt.color.NoSuchColorException;
+import net.sourceforge.plantuml.klimt.creole.Display;
+import net.sourceforge.plantuml.plasma.Quark;
 import net.sourceforge.plantuml.regex.IRegex;
 import net.sourceforge.plantuml.regex.RegexConcat;
 import net.sourceforge.plantuml.regex.RegexLeaf;
@@ -64,11 +71,27 @@ public class CommandAssociateRelationship extends SingleLineCommand2<ChenEerDiag
   }
 
   @Override
-  protected CommandExecutionResult executeArg(ChenEerDiagram system, LineLocation location, RegexResult arg)
+  protected CommandExecutionResult executeArg(ChenEerDiagram diagram, LineLocation location, RegexResult arg)
       throws NoSuchColorException {
-    final String name = arg.get("NAME", 0);
+    final Entity relationship = diagram.peekOwner();
+    if (relationship == null) {
+      return CommandExecutionResult.error("Can only associate from a relationship");
+    }
 
-    System.out.println("- " + name);
+    final String entityName = diagram.cleanId(arg.get("NAME", 0));
+
+    final Quark<Entity> entityQuark = diagram.quarkInContext(true, entityName);
+    final Entity entity = entityQuark.getData();
+    if (entity == null) {
+      return CommandExecutionResult.error("No such entity: " + entityName);
+    }
+
+    final LinkType linkType = new LinkType(LinkDecor.NONE, LinkDecor.NONE);
+    final Link link = new Link(diagram.getEntityFactory(), diagram.getCurrentStyleBuilder(), relationship, entity,
+        linkType,
+        // TODO: Cardinality
+        LinkArg.build(Display.NULL, 1));
+    diagram.addLink(link);
 
     return CommandExecutionResult.ok();
   }
diff --git a/src/net/sourceforge/plantuml/cheneer/command/CommandCreateAttribute.java b/src/net/sourceforge/plantuml/cheneer/command/CommandCreateAttribute.java
index 6db4cb0bb..9981fcf1e 100644
--- a/src/net/sourceforge/plantuml/cheneer/command/CommandCreateAttribute.java
+++ b/src/net/sourceforge/plantuml/cheneer/command/CommandCreateAttribute.java
@@ -37,9 +37,13 @@ package net.sourceforge.plantuml.cheneer.command;
 
 import net.sourceforge.plantuml.abel.Entity;
 import net.sourceforge.plantuml.abel.LeafType;
+import net.sourceforge.plantuml.abel.Link;
+import net.sourceforge.plantuml.abel.LinkArg;
 import net.sourceforge.plantuml.cheneer.ChenEerDiagram;
 import net.sourceforge.plantuml.command.CommandExecutionResult;
 import net.sourceforge.plantuml.command.SingleLineCommand2;
+import net.sourceforge.plantuml.decoration.LinkDecor;
+import net.sourceforge.plantuml.decoration.LinkType;
 import net.sourceforge.plantuml.klimt.color.NoSuchColorException;
 import net.sourceforge.plantuml.klimt.creole.Display;
 import net.sourceforge.plantuml.plasma.Quark;
@@ -83,28 +87,38 @@ public class CommandCreateAttribute extends SingleLineCommand2<ChenEerDiagram> {
         new RegexOptional(//
             new RegexConcat( //
                 RegexLeaf.spaceZeroOrMore(), //
-                new RegexLeaf("COMPOSITE", "\\{"))), //
+                new RegexLeaf("COMPOSITE", "(\\{)"))), //
         RegexLeaf.end());
   }
 
   @Override
   protected CommandExecutionResult executeArg(ChenEerDiagram diagram, LineLocation location, RegexResult arg)
       throws NoSuchColorException {
-    final LeafType type = LeafType.OBJECT;
+    final Entity owner = diagram.peekOwner();
+    if (owner == null) {
+      return CommandExecutionResult.error("Attribute must be inside an entity, relationship or another attribute");
+    }
+
+    final LeafType type = LeafType.USECASE;
     final String name = diagram.cleanId(arg.get("NAME", 0));
-    final boolean composite = arg.get("COMPOSITE") != null;
+    final String id = owner.getName() + "/" + name;
+    final boolean composite = arg.get("COMPOSITE", 0) != null;
+
+    final Quark<Entity> quark = diagram.quarkInContext(true, id);
 
-    final Quark<Entity> quark = diagram.quarkInContext(true, name);
     Entity entity = quark.getData();
-
     if (entity == null) {
-      Display display = Display.getWithNewlines(name);
+      final Display display = Display.getWithNewlines(name);
       entity = diagram.reallyCreateLeaf(quark, display, type, null);
     } else {
-      if (entity.muteToType(type, null) == false)
-        return CommandExecutionResult.error("Bad name");
+      return CommandExecutionResult.error("Attribute already exists");
     }
 
+    final LinkType linkType = new LinkType(LinkDecor.NONE, LinkDecor.NONE);
+    final Link link = new Link(diagram.getEntityFactory(), diagram.getCurrentStyleBuilder(), entity, owner, linkType,
+        LinkArg.build(Display.NULL, 1));
+    diagram.addLink(link);
+
     if (composite) {
       diagram.pushOwner(entity);
     }