usingnamespace @import("xrvk.zig");
const std = @import("std");
const glm = @import("glm");
const gfx = @import("graphics.zig");
const math = @import("math.zig");
const scene = @import("scene.zig");
const renderdoc = @import("renderdoc.zig");
const mem = std.mem;
const Allocator = mem.Allocator;

pub fn main() !void {
    try renderdoc.load_api();
    var session = try Session.init(std.heap.page_allocator);
    try session.loop();
    session.deinit();
}

const HandState = struct {
    subaction_path: xr.Path,
    handedness: []const u8,
    space: xr.Space,

    active: bool,
    grab: ?f32,
    pose: ?glm.Mat4,
    quit: ?bool,
    vibrate: ?xr.HapticVibration,

    pub fn init(path: xr.Path, handedness: []const u8, space: xr.Space) HandState {
        return .{
            .subaction_path = path,
            .space = space,
            .handedness =  handedness,

            .active = false,
            .grab = null,
            .pose = null,
            .quit = null,
            .vibrate = null,
        };
    }

    pub fn sync(self: *HandState, actions: *const Actions) !void {
        const xri = actions.xri;
        const session = actions.session;

        var grab = xr.ActionStateFloat.empty();
        var pose = xr.ActionStatePose.empty();
        var quit = xr.ActionStateBoolean.empty();

        try xri.getActionStateFloat(session, .{
            .action = actions.grab_object,
            .subaction_path = self.subaction_path,
        }, &grab);
        self.grab = if (grab.is_active == xr.TRUE) grab.current_state else null;

        try xri.getActionStatePose(session, .{
            .action = actions.grip_pose,
            .subaction_path = self.subaction_path,
        }, &pose);
        self.active = pose.is_active == xr.TRUE;

        try xri.getActionStateBoolean(session, .{
            .action = actions.quit_session,
            .subaction_path = self.subaction_path,
        }, &quit);
        self.quit = if (quit.is_active == xr.TRUE) quit.current_state == xr.TRUE else null;
    }

    pub fn locate(
        self: *HandState,
        xri: xr.InstanceDispatch,
        reference_space: xr.Space,
        query_time: xr.Time,
    ) !void {
        if (!self.active) {
            self.pose = null;
            return;
        }

        var location = xr.SpaceLocation.empty();
        try xri.locateSpace(self.space, reference_space, query_time, &location);

        if (!location.location_flags.contains(.{
            .position_valid_bit = true,
            .orientation_valid_bit = true,
        })) {
            self.pose = null;
            return;
        }

        self.pose = math.pose2mat(location.pose);
    }

    pub fn applyOutput(self: *const HandState, actions: *const Actions) !void {
        if (self.vibrate) |vibration| {
            try xri.applyHapticFeedback(
                session,
                .{
                    .action = self.actions.vibrate,
                    .subaction_path = self.subaction_path,
                },
                @ptrCast(*const xr.HapticBaseHeader, &vibration),
            );
        }
    }
};

const Actions = struct {
    xri: xr.InstanceDispatch,
    session: xr.Session,

    set: xr.ActionSet,
    grab_object: xr.Action,
    grip_pose: xr.Action,
    quit_session: xr.Action,
    vibrate: xr.Action,

    hands: [2]HandState,

    pub fn init(
        xri: xr.InstanceDispatch,
        instance: xr.Instance,
        session: xr.Session,
    ) !Actions {
        var self: Actions = undefined;
        self.xri = xri;
        self.session = session;

        const subaction_strs = [_][:0]const u8{ "/user/hand/right", "/user/hand/left" };
        const handedness_strs = [_][:0]const u8{ "right", "left" };
        var subaction_paths: [subaction_strs.len]xr.Path = undefined;
        for (subaction_strs) |str, i| {
            subaction_paths[i] = try xri.stringToPath(instance, str);
        }

        var set_ci = xr.ActionSetCreateInfo{
            .action_set_name = undefined,
            .localized_action_set_name = undefined,
            .priority = 0,
        };
        mem.copy(u8, set_ci.action_set_name[0..], "gameplay\x00");
        mem.copy(u8, set_ci.localized_action_set_name[0..], "Gameplay\x00");
        self.set = try xri.createActionSet(instance, set_ci);

        {
            // create actions
            var action_ci = xr.ActionCreateInfo{
                .action_type = undefined,
                .action_name = undefined,
                .localized_action_name = undefined,
                .count_subaction_paths = subaction_paths.len,
                .subaction_paths = &subaction_paths,
            };

            action_ci.action_type = .float_input;
            mem.copy(u8, action_ci.action_name[0..], "grab_object\x00");
            mem.copy(u8, action_ci.localized_action_name[0..], "Grab Object\x00");
            self.grab_object = try xri.createAction(self.set, action_ci);

            action_ci.action_type = .pose_input;
            mem.copy(u8, action_ci.action_name[0..], "grip_pose\x00");
            mem.copy(u8, action_ci.localized_action_name[0..], "Grip Pose\x00");
            self.grip_pose = try xri.createAction(self.set, action_ci);

            action_ci.action_type = .boolean_input;
            mem.copy(u8, action_ci.action_name[0..], "quit_session\x00");
            mem.copy(u8, action_ci.localized_action_name[0..], "Quit Session\x00");
            self.quit_session = try xri.createAction(self.set, action_ci);

            action_ci.action_type = .vibration_output;
            mem.copy(u8, action_ci.action_name[0..], "vibrate_hand\x00");
            mem.copy(u8, action_ci.localized_action_name[0..], "Vibrate Hand\x00");
            self.vibrate = try xri.createAction(self.set, action_ci);
        }

        {
            // suggested bindings
            const PathStruct = struct {
                squeeze_force: xr.Path,
                pose: xr.Path,
                b_click: xr.Path,
                haptic: xr.Path,
            };

            var paths: [subaction_strs.len]PathStruct = undefined;
            for (subaction_strs) |subaction, i| {
                var path_buffer = [_]u8{0} ** 2048;
                var string: [:0]const u8 = undefined;

                string = try std.fmt.bufPrintZ(path_buffer[0..], "{}/input/squeeze/force", .{subaction});
                const squeeze_force = try xri.stringToPath(instance, string.ptr);

                string = try std.fmt.bufPrintZ(path_buffer[0..], "{}/input/grip/pose", .{subaction});
                const pose = try xri.stringToPath(instance, string.ptr);

                string = try std.fmt.bufPrintZ(path_buffer[0..], "{}/input/b/click", .{subaction});
                const b_click = try xri.stringToPath(instance, string.ptr);

                string = try std.fmt.bufPrintZ(path_buffer[0..], "{}/output/haptic", .{subaction});
                const haptic = try xri.stringToPath(instance, string.ptr);

                paths[i] = .{
                    .squeeze_force = squeeze_force,
                    .pose = pose,
                    .b_click = b_click,
                    .haptic = haptic,
                };
            }

            const index_profile = try xri.stringToPath(instance, "/interaction_profiles/valve/index_controller");
            const bindings = [_]xr.ActionSuggestedBinding{
                .{ .action = self.grab_object, .binding = paths[0].squeeze_force },
                .{ .action = self.grab_object, .binding = paths[1].squeeze_force },
                .{ .action = self.grip_pose, .binding = paths[0].pose },
                .{ .action = self.grip_pose, .binding = paths[1].pose },
                .{ .action = self.quit_session, .binding = paths[0].b_click },
                .{ .action = self.quit_session, .binding = paths[1].b_click },
                .{ .action = self.vibrate, .binding = paths[0].haptic },
                .{ .action = self.vibrate, .binding = paths[1].haptic },
            };
            try xri.suggestInteractionProfileBindings(instance, .{
                .interaction_profile = index_profile,
                .count_suggested_bindings = bindings.len,
                .suggested_bindings = &bindings,
            });
        }

        for (subaction_paths) |path, i| {
            const handedness = handedness_strs[i];
            const space = try xri.createActionSpace(session, .{
                .action = self.grip_pose,
                .pose_in_action_space = .{
                    // "origin": [-0.005154, 0.013042, 0.107171],
                    // "rotate_xyz" : [93.782, 0.0, 0.0]
                    // .position = .{ .x = 0.005154, .y = -0.013042, .z = -0.107171 },
                    // .orientation = .{ .x = 0.7300549, .y = 0, .z = 0, .w = 0.6833885 },
                },
                .subaction_path = path,
            });
            self.hands[i] = HandState.init(path, handedness, space);
        }

        try xri.attachSessionActionSets(session, .{
            .count_action_sets = 1,
            .action_sets = @ptrCast([*]xr.ActionSet, &self.set),
        });

        return self;
    }

    pub fn sync(self: *Actions) !void {
        const active_sets = [_]xr.ActiveActionSet{
            .{ .action_set = self.set, .subaction_path = 0 },
        };
        _ = try self.xri.syncActions(self.session, .{
            .count_active_action_sets = active_sets.len,
            .active_action_sets = &active_sets,
        });

        for (self.hands) |*hand| {
            try hand.sync(self);
        }
    }
};

const Session = struct {
    xrb: xr.BaseDispatch,
    xri: xr.InstanceDispatch,
    xrcmi: ?xr.CMInstanceDispatch,

    instance: xr.Instance,
    session: xr.Session,
    system: xr.SystemId,

    allocator: *Allocator,

    actions: Actions,
    state: xr.SessionState,
    views: ?ViewInfo,
    graphics: gfx.Graphics,
    layers: []*const xr.CompositionLayerBaseHeader,
    scene_space: xr.Space,

    hand_models: [2]?*scene.Model,

    const ViewInfo = struct {
        configurations: []xr.ViewConfigurationView,
        projections: []xr.CompositionLayerProjectionView,
    };

    pub fn init(allocator: *Allocator) !Session {
        var name: [128]u8 = undefined;
        mem.copy(u8, name[0..], "openxr-zig-test" ++ [_]u8{0});

        const xrb = try xr.BaseDispatch.load(xr.getProcAddr);

        const have_controller_model_ext = blk: {
            var extensions = [_]xr.ExtensionProperties{ xr.ExtensionProperties.empty() } ** 512;
            const extension_count = try xrb.enumerateInstanceExtensionProperties(null, extensions.len, &extensions);
            for (extensions[0..extension_count]) |ext| {
                const ext_name = @ptrCast([*:0]const u8, ext.extension_name[0..]);
                if (mem.eql(u8, mem.spanZ(ext_name), "XR_MSFT_controller_model")) {
                    std.debug.print("found XR_MSFT_controller_model extension version {}\n", .{ ext.extension_version },);
                    break :blk true;
                }
            }

            break :blk false;
        };

        const zero = [_:0]u8{};
        const extensions = [_][*:0]const u8{
            "XR_KHR_vulkan_enable2",
            "XR_MSFT_controller_model",
        };

        const instance = try xrb.createInstance(.{
            .create_flags = .{},
            .application_info = .{
                .application_name = name,
                .application_version = 0,
                .engine_name = name,
                .engine_version = 0,
                .api_version = xr.makeVersion(1, 0, 0),
            },
            .enabled_api_layer_count = 0,
            .enabled_api_layer_names = @ptrCast([*]const [*:0]const u8, &zero),
            .enabled_extension_count = if (have_controller_model_ext) 2 else 1,
            .enabled_extension_names = &extensions,
        });

        const xri = try xr.InstanceDispatch.load(instance, xr.getProcAddr);
        errdefer xri.destroyInstance(instance) catch unreachable;

        var xrcmi: ?xr.CMInstanceDispatch = null;
        if (have_controller_model_ext) {
            xrcmi = try xr.CMInstanceDispatch.load(instance, xr.getProcAddr);
        }

        const system = try xri.getSystem(instance, .{ .form_factor = .head_mounted_display });

        var system_properties = xr.SystemProperties.empty();
        try xri.getSystemProperties(instance, system, &system_properties);

        std.debug.print(
            \\system {}:
            \\  vendor Id: {}
            \\  systemName: {}
            \\  gfx
            \\    max swapchain image resolution: {}x{}
            \\    max layer count: {}
            \\  tracking
            \\    orientation tracking: {}
            \\    positional tracking: {}
            \\
        , .{
            system,
            system_properties.vendor_id,
            system_properties.system_name,
            system_properties.graphics_properties.max_swapchain_image_width,
            system_properties.graphics_properties.max_swapchain_image_height,
            system_properties.graphics_properties.max_layer_count,
            system_properties.tracking_properties.orientation_tracking,
            system_properties.tracking_properties.position_tracking,
        });

        var graphics = try gfx.Graphics.init(allocator, xri, instance, system);
        try graphics.setupDebugging();

        const binding = graphics.getBinding();
        const session = try xri.createSession(instance, .{
            .next = &binding,
            .create_flags = .{},
            .system_id = system,
        });
        errdefer xri.destroySession(session) catch unreachable;

        const actions = try Actions.init(xri, instance, session);

        const scene_space = try xri.createReferenceSpace(session, .{
            .reference_space_type = .stage,
            .pose_in_reference_space = .{
                .orientation = .{ .w = 1 },
            },
        });

        return Session{
            .xrb = xrb,
            .xri = xri,
            .xrcmi = xrcmi,

            .instance = instance,
            .session = session,
            .system = system,
            .allocator = allocator,

            .actions = actions,
            .state = .idle,
            .views = null,
            .graphics = graphics,
            .scene_space = scene_space,
            .layers = &[_]*const xr.CompositionLayerBaseHeader{},

            .hand_models = [_]?*scene.Model{ null, null },
        };
    }

    pub fn deinit(self: *const Session) void {
        if (self.views) |views| {
            self.allocator.free(views.configurations);
            self.allocator.free(views.projections);
        }

        for (self.layers) |layer| {
            self.allocator.destroy(layer);
        }

        self.graphics.deinit();

        _ = self.xri.destroySession(self.session) catch null;
        _ = self.xri.destroyInstance(self.instance) catch null;
    }

    fn pollEvents(self: *Session) !void {
        while (true) {
            var event = xr.EventDataBuffer.empty();
            const poll_result = try self.xri.pollEvent(self.instance, &event);

            if (poll_result == .event_unavailable) {
                break;
            }

            switch (event.type) {
                .event_data_session_state_changed => {
                    const state_event = @ptrCast(*const xr.EventDataSessionStateChanged, &event);
                    std.debug.print("session state {} → {}\n", .{ self.state, state_event.state });
                    self.state = state_event.state;
                    switch (self.state) {
                        .ready => {
                            _ = try self.xri.beginSession(
                                self.session,
                                .{ .primary_view_configuration_type = .primary_stereo },
                            );
                        },
                        .stopping => {
                            _ = try self.xri.endSession(self.session);
                        },
                        .idle, .synchronized, .visible, .focused, .exiting => {},
                        .loss_pending => return error.LossPending,
                        else => return error.InvalidState,
                    }
                },
                .event_data_interaction_profile_changed,
                .event_data_reference_space_change_pending,
                => {
                    std.debug.print("skipping event: {}\n", .{event.type});
                },

                .event_data_instance_loss_pending => break,
                .event_data_events_lost => return error.EventsLost,

                else => {
                    std.debug.print("unknown event: {}\n", .{event.type});
                },
            }
        }
    }

    pub fn loop(self: *Session) !void {
        const view_count = try self.xri.enumerateViewConfigurationViews(
            self.instance,
            self.system,
            .primary_stereo,
            0,
            null,
        );
        const views = ViewInfo{
            .configurations = try self.allocator.alloc(xr.ViewConfigurationView, view_count),
            .projections = try self.allocator.alloc(xr.CompositionLayerProjectionView, view_count),
        };

        mem.set(
            xr.ViewConfigurationView,
            views.configurations,
            xr.ViewConfigurationView.empty(),
        );
        mem.set(
            xr.CompositionLayerProjectionView,
            views.projections,
            xr.CompositionLayerProjectionView.empty(),
        );

        _ = try self.xri.enumerateViewConfigurationViews(
            self.instance,
            self.system,
            .primary_stereo,
            view_count,
            views.configurations.ptr,
        );

        for (views.configurations) |config, i| {
            views.projections[i] = xr.CompositionLayerProjectionView.empty();
        }

        try self.graphics.createSwapchain(self.session, views.configurations);
        var gfx_scene = try scene.Scene.init(&self.graphics);

        var helmet = scene.Model.initFile(self.allocator, &gfx_scene, "assets/DamagedHelmet.gltf");
        helmet.setTransform(glm.translation(glm.Vec3.init([_]f32{ 0, 1.5, -2 })));
        try gfx_scene.models.append(helmet);

        helmet = scene.Model.initFile(self.allocator, &gfx_scene, "assets/controller_right.glb");
        helmet.setTransform(
            glm.translation(glm.Vec3.init([_]f32{ -0.07, 1.6, -0.2 }))
            .mul(glm.rotation(0.8, glm.Vec3.init([_]f32{ 1, 0, 0 })))
        );
        try helmet.loadNodeTransform("r_button_b");
        try helmet.loadNodeTransform("r_trackpad_touch");
        try gfx_scene.models.append(helmet);

        const composition_layer = try self.allocator.create(xr.CompositionLayerProjection);
        composition_layer.* = .{
            .layer_flags = .{ .blend_texture_source_alpha_bit = true },
            .space = self.scene_space,
            .view_count = view_count,
            .views = views.projections.ptr,
        };

        self.views = views;
        self.layers = try self.allocator.alloc(*xr.CompositionLayerBaseHeader, 1);
        self.layers[0] = @ptrCast(*const xr.CompositionLayerBaseHeader, composition_layer);

        while (true) {
            try self.pollEvents();
            if (self.state == .stopping) break;
            try self.actions.sync();

            const did_render = try self.renderFrame(&gfx_scene);
            if (!did_render) {
                // no frame submitted, so compositor won't throttle by waiting for GPU
                // sleep to prevent busy-loop
                std.os.nanosleep(0, 250_000_000);
            }
        }
    }

    pub fn renderFrame(self: *Session, gfx_scene: *scene.Scene) !bool {
        if (self.views) |views| {
            self.graphics.lockQueue();
            errdefer self.graphics.unlockQueue();
            var frame_state = xr.FrameState.empty();
            _ = try self.xri.waitFrame(self.session, null, &frame_state);
            _ = try self.xri.beginFrame(self.session, null);
            self.graphics.unlockQueue();

            const predicted_seconds = @intToFloat(f32, frame_state.predicted_display_time) / 1_000_000_000;

            var transform = glm.translation(glm.Vec3.init([_]f32{ 0, 1.5, -2 }));
            transform = transform.mul(glm.rotation(predicted_seconds, glm.Vec3.init([_]f32{ 0, 1, 0 })));
            gfx_scene.models.items[0].setTransform(transform);

            {
                const model = &gfx_scene.models.items[1];
                var txm = &model.node_transforms.items[0];
                // motion.type == translate:l
                //   apply translation axis * lerp(in, value_mapping)
                txm.position = glm.Vec3.init([_]f32{0, -0.927, 0.0375}).mulScalar(
                    (std.math.sin(predicted_seconds * 3) / 2 + 0.5) * 0.002
                );

                var tp_mot_rot = glm.Mat4.IDENTITY;
                // motion.type == trackpad: only apply negative component_local transform
                //   apply translation axis * lerp(in, value_mapping)
                tp_mot_rot.mulAssign(glm.rotation(1.18682389, glm.Vec3.init([_]f32{ 1, 0, 0 })));
                tp_mot_rot.mulAssign(math.quat2mat(xr.Quaternionf{ .x=-0.17349398, .y=-0.79472193, .z=0.32537197, .w=-0.48213067 }));
                tp_mot_rot.mulAssign(glm.rotation(-1.18682389, glm.Vec3.init([_]f32{ 1, 0, 0 })));

                txm = &model.node_transforms.items[1];
                txm.position =
                    tp_mot_rot.apply(
                        glm.Vec4.init([_]f32{
                            std.math.clamp(std.math.sin(predicted_seconds * 3) * 2, -1, 1) * 0.009,
                            std.math.clamp(std.math.cos(predicted_seconds * 3) * 2, -1, 1) * 0.015,
                            0,
                            1
                        })
                    ).shrink(3);
                // txm.rotation = glm.Vec4.init([_]f32{ -0.17349398, -0.79472193,  0.32537197, -0.48213067 });

                model.flushTransforms();
                std.debug.print("transform pos: {}\n", .{ model.node_transforms.items[0].position });
            }

            var layer_count: u32 = 0;

            if (frame_state.should_render > 0) {
                if (renderdoc.api) |api| api.StartFrameCapture(null, null);

                var located_views = [_]xr.View{xr.View.empty()} ** 2;
                var view_state = xr.ViewState.empty();

                _ = try self.xri.locateViews(self.session, .{
                    .view_configuration_type = .primary_stereo,
                    .display_time = frame_state.predicted_display_time,
                    .space = self.scene_space,
                }, &view_state, 2, &located_views);

                if (!view_state.view_state_flags.contains(.{ .position_valid_bit = true, .orientation_valid_bit = true })) {
                    return error.InvalidPositionOrOrientation;
                }

                for (self.actions.hands) |*hand, i| {
                    try hand.locate(self.xri, self.scene_space, frame_state.predicted_display_time);
                    if (hand.active and self.hand_models[i] == null) {
                        var model = if (self.xrcmi) |xri| model: {
                            var model_state = xr.ControllerModelKeyStateMSFT.empty();
                            try xri.getControllerModelKeyMSFT(self.session, hand.subaction_path, &model_state);

                            if (model_state.model_key == 0)
                                continue;

                            const size = try xri.loadControllerModelMSFT(self.session, model_state.model_key, 0, null);
                            const buffer = try self.allocator.alloc(u8, size);
                            defer self.allocator.free(buffer);

                            _ = try xri.loadControllerModelMSFT(self.session, model_state.model_key, size, buffer.ptr);
                            break :model scene.Model.initMemory(self.allocator, gfx_scene, buffer);
                        } else model: {
                            var path_buffer = [_]u8{0} ** 128;
                            const path = try std.fmt.bufPrintZ(path_buffer[0..], "assets/controller_{}.glb", .{hand.handedness});
                            break :model scene.Model.initFile(self.allocator, gfx_scene, path);
                        };

                        const ptr = try gfx_scene.models.addOne();
                        ptr.* = model;
                        self.hand_models[i] = ptr;
                    }

                    if (hand.pose) |pose| {
                        if (self.hand_models[i]) |model| {
                            model.setTransform(pose);
                        }
                    }
                }

                const swapchain = self.graphics.swapchain;
                try swapchain.acquireImage();

                for (located_views) |located_view, i| {
                    views.projections[i] = .{
                        .pose = located_view.pose,
                        .fov = located_view.fov,
                        .sub_image = .{
                            .swapchain = swapchain.handle,
                            .image_rect = swapchain.image_rects[i],
                            .image_array_index = 0,
                        },
                    };
                }

                try gfx_scene.render(located_views[0..]);

                swapchain.presentImage();

                _ = try self.xri.endFrame(self.session, .{
                    .display_time = frame_state.predicted_display_time,
                    .environment_blend_mode = .@"opaque",
                    .layer_count = @intCast(u32, self.layers.len),
                    .layers = self.layers.ptr,
                });

                if (renderdoc.api) |api| _ = api.EndFrameCapture(null, null);

                return true;
            }

            self.graphics.lockQueue();
            errdefer self.graphics.unlockQueue();
            _ = try self.xri.endFrame(self.session, .{
                .display_time = frame_state.predicted_display_time,
                .environment_blend_mode = .@"opaque",
                .layer_count = 0,
                .layers = null,
            });
            self.graphics.unlockQueue();

            return false;
        }

        unreachable;
    }
};