package eu.jonahbauer.raytracing.scene.util;

import eu.jonahbauer.raytracing.math.AABB;
import eu.jonahbauer.raytracing.math.Ray;
import eu.jonahbauer.raytracing.scene.Hittable;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Comparator;
import java.util.List;

public final class HittableBinaryTree extends HittableCollection {
    private final @Nullable Hittable left;
    private final @Nullable Hittable right;
    private final @NotNull AABB bbox;

    public HittableBinaryTree(@NotNull List<? extends @NotNull Hittable> objects) {
        bbox = AABB.getBoundingBox(objects).orElse(AABB.EMPTY);
        if (objects.isEmpty()) {
            left = null;
            right = null;
        } else if (objects.size() == 1) {
            left = objects.getFirst();
            right = null;
        } else if (objects.size() == 2) {
            left = objects.getFirst();
            right = objects.getLast();
        } else {
            var x = bbox.x().size();
            var y = bbox.y().size();
            var z = bbox.z().size();
            Comparator<AABB> comparator;
            if (x > y && x > z) {
                comparator = AABB.X_AXIS;
            } else if (y > z) {
                comparator = AABB.Y_AXIS;
            } else {
                comparator = AABB.Z_AXIS;
            }
            var sorted = objects.stream().sorted(Comparator.comparing(Hittable::getBoundingBox, comparator)).toList();
            var size = sorted.size();

            left = new HittableBinaryTree(sorted.subList(0, size / 2));
            right = new HittableBinaryTree(sorted.subList(size / 2, size));
        }
    }

    @Override
    public void hit(@NotNull Ray ray, @NotNull State state) {
        if (!bbox.hit(ray)) return;
        if (left instanceof HittableCollection coll) {
            coll.hit(ray, state);
        } else if (left != null) {
            hit(state, ray, left);
        }
        if (right instanceof HittableCollection coll) {
            coll.hit(ray, state);
        } else if (right != null) {
            hit(state, ray, right);
        }
    }

    @Override
    public @NotNull AABB getBoundingBox() {
        return bbox;
    }
}
