content: package io.agrest.cayenne.processor.update.stage;
  
  import io.agrest.AgException;
  import io.agrest.EntityUpdate;
  import io.agrest.RelatedResourceEntity;
  import io.agrest.RootResourceEntity;
  import io.agrest.cayenne.path.IPathResolver;
  import io.agrest.cayenne.path.PathOps;
  import io.agrest.cayenne.persister.ICayennePersister;
  import io.agrest.cayenne.processor.CayenneProcessor;
  import io.agrest.meta.AgAttribute;
  import io.agrest.meta.AgEntity;
  import io.agrest.meta.AgIdPart;
  import io.agrest.meta.AgRelationship;
  import io.agrest.processor.ProcessorOutcome;
  import io.agrest.runtime.EntityParent;
  import io.agrest.runtime.constraints.UpdateConstraints;
  import io.agrest.runtime.processor.update.UpdateContext;
  import io.agrest.runtime.processor.update.stage.UpdateApplyServerParamsStage;
  import org.apache.cayenne.di.Inject;
  import org.apache.cayenne.exp.parser.ASTDbPath;
  import org.apache.cayenne.exp.parser.ASTPath;
  import org.apache.cayenne.map.DbJoin;
  import org.apache.cayenne.map.DbRelationship;
  import org.apache.cayenne.map.EntityResolver;
  import org.apache.cayenne.map.ObjAttribute;
  import org.apache.cayenne.map.ObjEntity;
  import org.apache.cayenne.map.ObjRelationship;
  
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  
  /**
   * @since 2.7
   */
  public class CayenneUpdateApplyServerParamsStage extends UpdateApplyServerParamsStage {
  
      private final UpdateConstraints constraints;
      private final EntityResolver entityResolver;
      private final IPathResolver pathResolver;
  
      public CayenneUpdateApplyServerParamsStage(
              @Inject IPathResolver pathResolver,
              @Inject UpdateConstraints constraints,
              @Inject ICayennePersister persister) {
  
          this.pathResolver = pathResolver;
          this.constraints = constraints;
          this.entityResolver = persister.entityResolver();
      }
  
      @Override
      public ProcessorOutcome execute(UpdateContext<?> context) {
          doExecute(context);
          return ProcessorOutcome.CONTINUE;
      }
  
      protected <T> void doExecute(UpdateContext<T> context) {
  
          // this creates a new EntityUpdate if there's no request payload, but the context ID is present,
          // so execute it unconditionally
          fillIdsFromExplicitContextId(context);
  
          if (!context.getUpdates().isEmpty()) {
              fillIdsFromRelatedIds(context);
              fillIdsFromMeaningfulPk(context);
              fillIdsFromParentId(context);
          }
  
          constraints.apply(context);
  
          tagRootEntity(context.getEntity());
      }
  
      private void tagRootEntity(RootResourceEntity<?> entity) {
          if (entityResolver.getObjEntity(entity.getName()) != null) {
              CayenneProcessor.getOrCreateRootEntity(entity);
          }
  
          if (entity.getMapBy() != null) {
              for (RelatedResourceEntity<?> child : entity.getMapBy().getChildren()) {
                  tagRelatedEntity(child);
              }
          }
  
          for (RelatedResourceEntity<?> child : entity.getChildren()) {
              tagRelatedEntity(child);
          }
      }
  
      private void tagRelatedEntity(RelatedResourceEntity<?> entity) {
          if (entityResolver.getObjEntity(entity.getName()) != null) {
              CayenneProcessor.getOrCreateRelatedEntity(entity);
          }
  
          if (entity.getMapBy() != null) {
              for (RelatedResourceEntity<?> child : entity.getMapBy().getChildren()) {
                  tagRelatedEntity(child);
              }
          }
  
          for (RelatedResourceEntity<?> child : entity.getChildren()) {
              tagRelatedEntity(child);
          }
      }
  
      private <T> void fillIdsFromRelatedIds(UpdateContext<T> context) {
  
          // for each relationship that propagates an id from the target entity to this, fill this entity "id" collection
          // with values from related PK...
  
          AgEntity<T> entity = context.getEntity().getAgEntity();
  
          for (AgRelationship r : entity.getRelationships()) {
  
              ObjRelationship objRelationship = objRelationshipForAgRelationship(entity.getName(), r);
  
              if (objRelationship != null) {
                  List<DbRelationship> dbRelationships = objRelationship.getDbRelationships();
                  if (dbRelationships.size() == 1) {
  
                      DbRelationship outgoingDbRelationship = dbRelationships.get(0);
                      DbRelationship incomingDbRelationship = outgoingDbRelationship.getReverseRelationship();
  
                      if (incomingDbRelationship.isToDependentPK()) {
  
                          List<DbJoin> joins = outgoingDbRelationship.getJoins();
  
                          // however unlikely, checking this for completeness
                          if (joins.size() == 0) {
                              throw new IllegalStateException("No joins for relationship " + r.getName());
                          } else if (joins.size() == 1) {
                              fillIdsFromRelatedId(context, r.getName(), joins.get(0));
                          } else {
                              fillIdsFromRelatedIds(context, r.getName(), joins);
                          }
                      }
                  }
              }
          }
      }
  
      private <T> void fillIdsFromMeaningfulPk(UpdateContext<T> context) {
  
          AgEntity<T> entity = context.getEntity().getAgEntity();
          ObjEntity objEntity = entityResolver.getObjEntity(entity.getName());
          if (objEntity != null) {
              for (AgAttribute a : entity.getAttributes()) {
                  ObjAttribute oa = objEntity.getAttribute(a.getName());
                  if (oa != null && oa.isPrimaryKey()) {
                      fillIdsFromMappedPk(context, a.getName());
                  }
              }
          }
      }
  
      private void fillIdsFromMappedPk(UpdateContext<?> context, String propertyName) {
  
          for (EntityUpdate<?> u : context.getUpdates()) {
              Object pk = u.getAttributes().get(propertyName);
  
              // Unlike id parts mapped from the DB layer, this one is tracked by its normal property name
              u.addIdPartIfAbsent(propertyName, pk);
          }
      }
  
      private void fillIdsFromRelatedId(UpdateContext<?> context, String agRelationshipName, DbJoin outgoingJoin) {
  
          for (EntityUpdate<?> u : context.getUpdates()) {
  
              // 'getSourceName' assumes AgEntity's id attribute name is based on DbAttribute name
              Object id = u.getToOneIds().get(agRelationshipName);
              if (id != null) {
                  u.addIdPartIfAbsent(ASTDbPath.DB_PREFIX + outgoingJoin.getSourceName(), id);
              }
          }
      }
  
      private void fillIdsFromRelatedIds(UpdateContext<?> context, String agRelationshipName, List<DbJoin> outgoingJoins) {
  
          for (EntityUpdate<?> u : context.getUpdates()) {
  
              Object pk = u.getToOneIds().get(agRelationshipName);
  
              if (pk != null) {
                  // TODO: should we allow null Map though?
                  if (!(pk instanceof Map)) {
                      throw new IllegalStateException("Expected more than one value in related 'id' for " + agRelationshipName);
                  }
  
                  // TODO: pretty sure this case has no unit tests
                  Map<String, Object> pkMap = (Map<String, Object>) pk;
                  for (DbJoin join : outgoingJoins) {
                      // 'getSourceName' and 'getTargetName' assumes AgEntity's id attribute name is based on DbAttribute name
                      u.addIdPartIfAbsent(ASTDbPath.DB_PREFIX + join.getSourceName(), pkMap.get(join.getTargetName()));
                  }
              }
          }
      }
  
      private <T> void fillIdsFromExplicitContextId(UpdateContext<T> context) {
  
          if (context.isById()) {
  
              // id was specified explicitly ... this means a few things:
              // * we expect zero or one object in the body
              // * if zero, create an empty update that will be attached to the ID.
              // * if more than one - throw...
  
              if (context.getUpdates().isEmpty()) {
                  context.setUpdates(Set.of(new EntityUpdate<>(context.getEntity().getAgEntity())));
              }
  
              AgEntity<T> entity = context.getEntity().getAgEntity();
              EntityUpdate<T> u = context.getFirst();
              u.setIdParts(context.getId().asMap(entity));
          }
      }
  
      private <T> void fillIdsFromParentId(UpdateContext<T> context) {
  
          EntityParent<?> parent = context.getParent();
  
          if (parent != null && parent.getId() != null) {
  
              Map<String, String> parentAgKeysToChildDbPaths = parentAgKeysToChildDbPaths(
                      context.getSchema().getEntity(parent.getType()),
                      parent);
              if (!parentAgKeysToChildDbPaths.isEmpty()) {
  
                  Map<String, Object> idTranslated = new HashMap<>(5);
                  for (Map.Entry<String, String> e : parentAgKeysToChildDbPaths.entrySet()) {
                      Object val = parent.getId().get(e.getKey());
                      if (val == null) {
                          throw AgException.badRequest("Missing parent id value for '%s'", e.getKey());
                      }
                      idTranslated.put(e.getValue(), val);
                  }
  
                  for (EntityUpdate<T> u : context.getUpdates()) {
                      idTranslated.forEach(u::addIdPartIfAbsent);
                  }
              }
          }
      }
  
      private Map<String, String> parentAgKeysToChildDbPaths(AgEntity<?> parentAgEntity, EntityParent<?> parent) {
  
          // TODO: too much Ag to Cayenne metadata translation happens here.
          //  We should cache and reuse the returned map
  
          ObjEntity parentCayenneEntity = entityResolver.getObjEntity(parent.getType());
          ObjRelationship fromParent = parentCayenneEntity.getRelationship(parent.getRelationship());
  
          if (fromParent == null) {
              throw AgException.internalServerError("Invalid parent relationship: '%s'", parent.getRelationship());
          }
  
          // TODO: flattened relationships handling
          DbRelationship dbFromParent = fromParent.getDbRelationships().get(0);
          if (!dbFromParent.isToDependentPK()) {
              return Collections.emptyMap();
          }
  
          Map<String, String> sourceToTargetJoins = new HashMap<>();
          for (DbJoin join : dbFromParent.getJoins()) {
              sourceToTargetJoins.put(join.getSourceName(), join.getTargetName());
          }
  
          Map<String, String> parentAgKeysToChildDbPaths = new HashMap<>();
          for (AgIdPart idPart : parentAgEntity.getIdParts()) {
  
              ASTPath idPath = pathResolver.resolve(parentCayenneEntity.getName(), idPart.getName()).getPathExp();
              ASTDbPath dbPath = PathOps.resolveAsDbPath(idPath);
              String targetPath = sourceToTargetJoins.remove(dbPath.getPath());
              if (targetPath == null) {
                  // not a part of our relationship, ignore
                  continue;
              }
  
              parentAgKeysToChildDbPaths.put(idPart.getName(), ASTDbPath.DB_PREFIX + targetPath);
          }
  
          if (!sourceToTargetJoins.isEmpty()) {
              throw AgException.internalServerError(
                      "One or more parent join keys are not mapped as Ag IDs and can't be propagated over relationship '%s': %s",
                      parent.getRelationship(),
                      sourceToTargetJoins.keySet());
          }
  
          return parentAgKeysToChildDbPaths;
      }
  
      // TODO: copied verbatim from CayenneQueryAssembler... Unify this code?
      protected ObjRelationship objRelationshipForAgRelationship(String sourceEntityName, AgRelationship relationship) {
          return entityResolver.getObjEntity(sourceEntityName).getRelationship(relationship.getName());
      }
  }
  