package eu.jonahbauer.raytracing.scene.util;

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

import java.util.*;
import java.util.function.Predicate;
import java.util.stream.IntStream;

public final class HittableOctree extends HittableCollection {
    private static final int LIST_SIZE_LIMIT = 16;

    private final @Nullable Storage storage;
    private final @NotNull AABB bbox;

    public HittableOctree(@NotNull List<? extends @NotNull Hittable> objects) {
        bbox = AABB.getBoundingBox(objects).orElse(AABB.EMPTY);
        storage = newStorage(bbox, objects);
    }

    private static @NotNull AABB[] getBoundingBoxes(@NotNull AABB aabb, @NotNull Vec3 center) {
        return new AABB[] {
                new AABB(new Range(aabb.x().min(), center.x()), new Range(aabb.y().min(), center.y()), new Range(aabb.z().min(), center.z())),
                new AABB(new Range(center.x(), aabb.x().max()), new Range(aabb.y().min(), center.y()), new Range(aabb.z().min(), center.z())),
                new AABB(new Range(aabb.x().min(), center.x()), new Range(center.y(), aabb.y().max()), new Range(aabb.z().min(), center.z())),
                new AABB(new Range(center.x(), aabb.x().max()), new Range(center.y(), aabb.y().max()), new Range(aabb.z().min(), center.z())),
                new AABB(new Range(aabb.x().min(), center.x()), new Range(aabb.y().min(), center.y()), new Range(center.z(), aabb.z().max())),
                new AABB(new Range(center.x(), aabb.x().max()), new Range(aabb.y().min(), center.y()), new Range(center.z(), aabb.z().max())),
                new AABB(new Range(aabb.x().min(), center.x()), new Range(center.y(), aabb.y().max()), new Range(center.z(), aabb.z().max())),
                new AABB(new Range(center.x(), aabb.x().max()), new Range(center.y(), aabb.y().max()), new Range(center.z(), aabb.z().max())),
        };
    }

    private static @Nullable Storage newStorage(@NotNull AABB aabb, @NotNull List<? extends @NotNull Hittable> objects) {
        if (objects.isEmpty()) return null;
        if (objects.size() < LIST_SIZE_LIMIT) {
            return new ListStorage(aabb, objects);
        } else {
            var center = aabb.center();
            var octants = (List<Hittable>[]) new List<?>[8];
            for (int i = 0; i < 8; i++) octants[i] = new ArrayList<>();
            var bboxes = getBoundingBoxes(aabb, center);
            var list = new ArrayList<Hittable>();

            for (var object : objects) {
                var bbox = object.getBoundingBox();
                var imin = getOctantIndex(center, bbox.min());
                var imax = getOctantIndex(center, bbox.max());
                if (imin == imax) {
                    octants[imin].add(object);
                } else {
                    list.add(object);
                }
            }

            return new NodeStorage(aabb, center, list, IntStream.range(0, 8).mapToObj(i -> newStorage(bboxes[i], octants[i])).toArray(Storage[]::new));
        }
    }

    @Override
    public void hit(@NotNull Ray ray, @NotNull State state) {
        hit(ray, object -> hit(state, ray, object));
    }

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

    /**
     * Use HERO algorithms to find all elements that could possibly be hit by the given ray.
     * @see <a href="https://doi.org/10.1007/978-3-642-76298-7_3">
     *      Agate, M., Grimsdale, R.L., Lister, P.F. (1991).
     *      The HERO Algorithm for Ray-Tracing Octrees.
     *      In: Grimsdale, R.L., Straßer, W. (eds) Advances in Computer Graphics Hardware IV. Eurographic Seminars. Springer, Berlin, Heidelberg.</a>
     */
    private void hit(@NotNull Ray ray, @NotNull Predicate<? super Hittable> action) {
        if (storage != null) storage.hit(ray, action);
    }

    private static int getOctantIndex(@NotNull Vec3 center, @NotNull Vec3 pos) {
        return (pos.x() < center.x() ? 0 : 1)
               | (pos.y() < center.y() ? 0 : 2)
               | (pos.z() < center.z() ? 0 : 4);

    }

    private abstract static sealed class Storage {
        protected final @NotNull AABB bbox;

        public Storage(@NotNull AABB bbox) {
            this.bbox = Objects.requireNonNull(bbox);
        }

        protected boolean hit(@NotNull Ray ray, @NotNull Predicate<? super Hittable> action) {
            var range = bbox.intersect(ray);
            if (range.isEmpty()) return false;

            int vmask = ray.vmask();
            return hit0(ray, vmask, range.get().min(), range.get().max(), action);
        }

        protected abstract boolean hit0(@NotNull Ray ray, int vmask, double tmin, double tmax, @NotNull Predicate<? super Hittable> action);
    }

    private static final class ListStorage extends Storage {
        private final @NotNull List<Hittable> list;

        public ListStorage(@NotNull AABB bbox, @NotNull List<? extends @NotNull Hittable> entries) {
            super(bbox);
            this.list = List.copyOf(entries);
        }

        @Override
        protected boolean hit0(@NotNull Ray ray, int vmask, double tmin, double tmax, @NotNull Predicate<? super Hittable> action) {
            var hit = false;
            for (Hittable hittable : list) {
                hit |= action.test(hittable);
            }
            return hit;
        }
    }

    private static final class NodeStorage extends Storage {
        private final @Nullable Storage @NotNull[] octants;
        private final @NotNull Vec3 center;
        private final int degenerate;

        private final @NotNull List<Hittable> list; // track elements spanning multiple octants separately

        public NodeStorage(@NotNull AABB bbox, @NotNull Vec3 center, @NotNull List<? extends @NotNull Hittable> list, @Nullable Storage @NotNull[] octants) {
            super(bbox);
            this.octants = octants;
            this.center = center;

            this.list = List.copyOf(list);

            int count = 0;
            int degenerate = 0;
            for (int i = 0; i < octants.length; i++) {
                if (octants[i] != null) {
                    count++;
                    degenerate = i;
                }
            }
            this.degenerate = count == 1 ? degenerate : -1;
        }

        @Override
        protected boolean hit(@NotNull Ray ray, @NotNull Predicate<? super Hittable> action) {
            if (degenerate >= 0 && list.isEmpty()) {
                return octants[degenerate].hit(ray, action);
            } else {
                return super.hit(ray, action);
            }
        }

        @Override
        protected boolean hit0(@NotNull Ray ray, int vmask, double tmin, double tmax, @NotNull Predicate<? super Hittable> action) {
            if (tmax < 0) return false;

            // check for hit
            var hit = false;

            // process entries spanning multiple children
            for (Hittable object : list) {
                hit |= action.test(object);
            }

            // t values for intersection points of ray with planes through center
            var tmid = AABB.intersect(center, ray);
            // masks of planes in the order of intersection, e.g. [2, 1, 4] for a ray intersection y = center.y() then x = center.x() then z = center.z()
            var masklist = calculateMasklist(tmid);
            // the first child to be hit by the ray assuming a ray with positive x, y and z coordinates
            var childmask = (tmid[0] < tmin ? 1 : 0)
                            | (tmid[1] < tmin ? 2 : 0)
                            | (tmid[2] < tmin ? 4 : 0);
            // the last child to be hit by the ray assuming a ray with positive x, y and z coordinates
            var lastmask = (tmid[0] < tmax ? 1 : 0)
                           | (tmid[1] < tmax ? 2 : 0)
                           | (tmid[2] < tmax ? 4 : 0);

            var childTmin = tmin;

            int i = 0;
            while (true) {
                // use vmask to nullify the assumption of a positive ray made for childmask
                var child = octants[childmask ^ vmask];

                // calculate t value for exit of child
                double childTmax;
                if (childmask == lastmask) {
                    // last child shares tmax
                    childTmax = tmax;
                } else {
                    // determine next child
                    while ((masklist[i] & childmask) != 0) {
                        i++;
                    }
                    childmask = childmask | masklist[i];
                    // tmax of current child is the t value for the intersection with the plane dividing the current and next child
                    childTmax = tmid[Integer.numberOfTrailingZeros(masklist[i])];
                }

                // process child
                var childHit = child != null && child.hit0(ray, vmask, childTmin, childTmax, action);
                hit |= childHit;

                // break after last child has been processed or a hit has been found
                if (childTmax == tmax || childHit) break;

                // tmin of next child is tmax of current child
                childTmin = childTmax;
            }

            return hit;
        }

        private static final int[][] MASKLISTS = new int[][] {
                {1, 2, 4},
                {1, 4, 2},
                {4, 1, 2},
                {2, 1, 4},
                {2, 4, 1},
                {4, 2, 1}
        };

        private static int @NotNull [] calculateMasklist(double @NotNull[] tmid) {
            if (tmid[0] < tmid[1]) {
                if (tmid[1] < tmid[2]) {
                    return MASKLISTS[0]; // {1, 2, 4}
                } else if (tmid[0] < tmid[2]) {
                    return MASKLISTS[1]; // {1, 4, 2}
                } else {
                    return MASKLISTS[2]; // {4, 1, 2}
                }
            } else {
                if (tmid[0] < tmid[2]) {
                    return MASKLISTS[3]; // {2, 1, 4}
                } else if (tmid[1] < tmid[2]) {
                    return MASKLISTS[4]; // {2, 4, 1}
                } else {
                    return MASKLISTS[5]; // {4, 2, 1}
                }
            }
        }
    }
}
