package de.bwaldvogel.mongo.backend.aggregation.stage;

import static de.bwaldvogel.mongo.backend.Constants.ID_FIELD;

import java.util.Map.Entry;
import java.util.stream.Stream;

import de.bwaldvogel.mongo.backend.Missing;
import de.bwaldvogel.mongo.backend.Utils;
import de.bwaldvogel.mongo.backend.aggregation.Expression;
import de.bwaldvogel.mongo.bson.Document;
import de.bwaldvogel.mongo.exception.MongoServerError;

public class ProjectStage implements AggregationStage {

    private final Document projection;
    private final boolean hasInclusions;

    public ProjectStage(Document projection) {
        if (projection.isEmpty()) {
            throw new MongoServerError(40177, "Invalid $project :: caused by :: specification must have at least one field");
        }
        this.projection = projection;
        this.hasInclusions = hasInclusions(projection);
    }

    private static boolean hasInclusions(Document projection) {
        return projection.values().stream().anyMatch(Utils::isTrue);
    }

    @Override
    public Stream<Document> apply(Stream<Document> stream) {
        return stream.map(this::projectDocument);
    }

    Document projectDocument(Document document) {
        final Document result;
        if (hasInclusions) {
            result = new Document();
            if (!projection.containsKey(ID_FIELD)) {
                putIfContainsField(document, result, ID_FIELD);
            }

            boolean nonIdExclusion = projection.entrySet().stream()
                .filter(entry -> !entry.getKey().equals(ID_FIELD))
                .map(Entry::getValue)
                .filter(ProjectStage::isNumberOrBoolean)
                .anyMatch(entry -> !Utils.isTrue(entry));
            if (nonIdExclusion) {
                throw new MongoServerError(40178,
                    "Bad projection specification, cannot exclude fields other than '_id' in an inclusion projection: "
                        + projection.toString(true, "{ ", " }"));
            }

        } else {
            result = document.cloneDeeply();
        }

        for (Entry<String, Object> entry : projection.entrySet()) {
            String field = entry.getKey();
            Object projectionValue = entry.getValue();
            if (isNumberOrBoolean(projectionValue)) {
                if (Utils.isTrue(projectionValue)) {
                    putIfContainsField(document, result, field);
                } else {
                    Utils.removeSubdocumentValue(result, field);
                }
            } else if (projectionValue == null) {
                result.put(field, null);
            } else {
                Object value = Expression.evaluateDocument(projectionValue, document);
                if (!(value instanceof Missing)) {
                    result.put(field, value);
                }
            }
        }
        return result;
    }

    private static boolean isNumberOrBoolean(Object projectionValue) {
        return projectionValue instanceof Number || projectionValue instanceof Boolean;
    }

    private static void putIfContainsField(Document input, Document result, String key) {
        Object value = Utils.getSubdocumentValue(input, key);
        if (!(value instanceof Missing)) {
            Utils.changeSubdocumentValue(result, key, value);
        }
    }
}

@@ -867,22 +867,52 @@
 ction);%0A
+        validateProjection();%0A
     %7D%0A
-
 %0A    pri
@@ -1458,582 +1458,8 @@
   %7D%0A
-%0A            boolean nonIdExclusion = projection.entrySet().stream()%0A                .filter(entry -%3E !entry.getKey().equals(ID_FIELD))%0A                .map(Entry::getValue)%0A                .filter(ProjectStage::isNumberOrBoolean)%0A                .anyMatch(entry -%3E !Utils.isTrue(entry));%0A            if (nonIdExclusion) %7B%0A                throw new MongoServerError(40178,%0A                    %22Bad projection specification, cannot exclude fields other than '_id' in an inclusion projection: %22%0A                        + projection.toString(true, %22%7B %22, %22 %7D%22));%0A            %7D%0A%0A
     
@@ -2303,32 +2303,32 @@
     %7D%0A        %7D%0A
-
         return r
@@ -2333,32 +2333,690 @@
  result;%0A    %7D%0A%0A
+    private void validateProjection() %7B%0A        if (hasInclusions) %7B%0A            boolean nonIdExclusion = projection.entrySet().stream()%0A                .filter(entry -%3E !entry.getKey().equals(ID_FIELD))%0A                .map(Entry::getValue)%0A                .filter(ProjectStage::isNumberOrBoolean)%0A                .anyMatch(entry -%3E !Utils.isTrue(entry));%0A            if (nonIdExclusion) %7B%0A                throw new MongoServerError(40178,%0A                    %22Bad projection specification, cannot exclude fields other than '_id' in an inclusion projection: %22%0A                        + projection.toString(true, %22%7B %22, %22 %7D%22));%0A            %7D%0A        %7D%0A    %7D%0A%0A
     private stat

Validate projection early in ProjectStage