package poussecafe.domain;

import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.function.Supplier;
import poussecafe.apm.ApmSpan;
import poussecafe.apm.ApplicationPerformanceMonitoring;
import poussecafe.environment.EntityFactory;
import poussecafe.environment.NewEntityInstanceSpecification;
import poussecafe.exception.NotFoundException;
import poussecafe.storage.MessageCollection;

import static java.util.stream.Collectors.toList;

public abstract class Repository<A extends AggregateRoot<K, D>, K, D extends EntityAttributes<K>> {

    @SuppressWarnings("unchecked")
    public void setEntityClass(Class<?> entityClass) {
        Objects.requireNonNull(entityClass);
        this.entityClass = (Class<A>) entityClass;
    }

    private Class<A> entityClass;

    public Class<A> entityClass() {
        return entityClass;
    }

    public void setMessageCollectionValidator(MessageCollectionValidator messageCollectionValidator) {
        this.messageCollectionValidator = messageCollectionValidator;
    }

    private MessageCollectionValidator messageCollectionValidator;

    public void setApplicationPerformanceMonitoring(ApplicationPerformanceMonitoring applicationPerformanceMonitoring) {
        this.applicationPerformanceMonitoring = applicationPerformanceMonitoring;
    }

    private ApplicationPerformanceMonitoring applicationPerformanceMonitoring;

    /**
     * @deprecated use getOptional instead.
     */
    @Deprecated(since = "0.8.0")
    public A find(K id) {
        return getOptional(id).orElse(null);
    }

    public Optional<A> getOptional(K id) {
        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();
        span.setName("get(" + entityClass.getSimpleName() + ")");
        try {
            checkId(id);
            return wrapNullable(dataAccess.findData(id));
        } finally {
            span.end();
        }
    }

    private EntityDataAccess<K, D> dataAccess;

    protected Optional<A> getOptional(String queryName, Supplier<D> query) {
        Objects.requireNonNull(query);
        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();
        span.setName(queryName);
        try {
            return wrapNullable(query.get());
        } finally {
            span.end();
        }
    }

    private void checkId(K id) {
        Objects.requireNonNull(id);
    }

    /**
     * @deprecated use wrapNullable instead.
     */
    @Deprecated(since = "0.8.0")
    protected A wrap(D data) {
        if(data == null) {
            return null;
        } else {
            return componentFactory.newEntity(new NewEntityInstanceSpecification.Builder<A>()
                    .entityClass(entityClass)
                    .existingData(data)
                    .build());
        }
    }

    protected Optional<A> wrapNullable(D data) {
        if(data == null) {
            return Optional.empty();
        } else {
            return Optional.of(componentFactory.newEntity(new NewEntityInstanceSpecification.Builder<A>()
                    .entityClass(entityClass)
                    .existingData(data)
                    .build()));
        }
    }

    private EntityFactory componentFactory;

    public A get(K id) {
        return getOptional(id)
                .orElseThrow(() -> new NotFoundException("Entity with id " + id + " not found"));
    }

    public void add(A entity) {
        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();
        span.setName("add(" + entityClass.getSimpleName() + ")");
        try {
            checkEntity(entity);
            entity.onAdd();
            addData(entity);
        } finally {
            span.end();
        }
    }

    protected void addData(A entity) {
        MessageCollection messageCollection = entity.messageCollection();
        if(!entity.dontPersist()) {
            messageCollectionValidator.validate(messageCollection);
            D data = entity.attributes();
            dataAccess.addData(data);
        }
        considerMessageSendingAfterAdd(entity, messageCollection);
    }

    private void considerMessageSendingAfterAdd(A entity,
            MessageCollection messageCollection) {
        entity.storage().getMessageSendingPolicy().considerSending(messageCollection);
    }

    private void checkEntity(A entity) {
        Objects.requireNonNull(entity);
        Objects.requireNonNull(entity.attributes());
    }

    public void update(A entity) {
        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();
        span.setName("update(" + entityClass.getSimpleName() + ")");
        try {
            checkEntity(entity);
            entity.onUpdate();
            updateData(entity);
        } finally {
            span.end();
        }
    }

    protected void updateData(A entity) {
        MessageCollection messageCollection = entity.messageCollection();
        if(!entity.dontPersist()) {
            messageCollectionValidator.validate(messageCollection);
            D data = entity.attributes();
            dataAccess.updateData(data);
        }
        considerMessageSendingAfterUpdate(entity, messageCollection);
    }

    private void considerMessageSendingAfterUpdate(A entity,
            MessageCollection messageCollection) {
        entity.storage().getMessageSendingPolicy().considerSending(messageCollection);
    }

    public void delete(K id) {
        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();
        span.setName("deleteByKey(" + entityClass.getSimpleName() + ")");
        try {
            checkId(id);
            Optional<A> entity = getOptional(id);
            if (entity.isPresent()) {
                deleteWithoutSpan(entity.get());
            }
        } finally {
            span.end();
        }
    }

    private void deleteWithoutSpan(A entity) {
        entity.onDelete();
        MessageCollection messageCollection = entity.messageCollection();
        messageCollectionValidator.validate(messageCollection);
        dataAccess.deleteData(entity.attributes().identifier().value());
        considerMessageSendingAfterDelete(entity, messageCollection);
    }

    public void delete(A entity) {
        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();
        span.setName("delete(" + entityClass.getSimpleName() + ")");
        try {
            deleteWithoutSpan(entity);
        } finally {
            span.end();
        }
    }

    private void considerMessageSendingAfterDelete(A entity,
            MessageCollection messageCollection) {
        entity.storage().getMessageSendingPolicy().considerSending(messageCollection);
    }

    protected List<A> find(String queryName, Supplier<List<D>> query) {
        Objects.requireNonNull(query);
        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();
        span.setName(queryName);
        try {
            return wrap(query.get());
        } finally {
            span.end();
        }
    }

    protected List<A> wrap(List<D> data) {
        return data.stream().map(this::wrapNullable).filter(Optional::isPresent).map(Optional::get).collect(toList());
    }

    @SuppressWarnings("unchecked")
    public void setDataAccess(Object dataAccess) {
        Objects.requireNonNull(dataAccess);
        this.dataAccess = (EntityDataAccess<K, D>) dataAccess;
    }

    public EntityDataAccess<K, D> dataAccess() {
        return dataAccess;
    }
}

@@ -4089,24 +4089,16 @@
 eSending
-AfterAdd
 (entity,
@@ -4168,16 +4168,8 @@
 ding
-AfterAdd
 (A e
@@ -4218,32 +4218,171 @@
 geCollection) %7B%0A
+        ApmSpan span = applicationPerformanceMonitoring.currentSpan().startSpan();%0A        span.setName(%22sendMessages%22);%0A        try %7B%0A    
         entity.s
@@ -4444,32 +4444,86 @@
 ageCollection);%0A
+        %7D finally %7B%0A            span.end();%0A        %7D%0A
     %7D%0A%0A    priva
@@ -5359,234 +5359,17 @@
 ding
-AfterUpdate(entity, messageCollection);%0A    %7D%0A%0A    private void considerMessageSendingAfterUpdate(A entity,%0A            MessageCollection messageCollection) %7B%0A        entity.storage().getMessageSendingPolicy().considerSending(
+(entity, 
 mess
@@ -6145,27 +6145,16 @@
 eSending
-AfterDelete
 (entity,
@@ -6486,214 +6486,8 @@
  %7D%0A%0A
-    private void considerMessageSendingAfterDelete(A entity,%0A            MessageCollection messageCollection) %7B%0A        entity.storage().getMessageSendingPolicy().considerSending(messageCollection);%0A    %7D%0A%0A
     

Add span for message sending.