local mod = get_mod("servo_friend")

local unit = Unit
local pairs = pairs
local world = World
local class = class
local CLASS = CLASS
local light = Light
local vector3 = Vector3
local tostring = tostring
local managers = Managers
local quaternion = Quaternion
local unit_light = unit.light
local unit_alive = unit.alive
local vector3_zero = vector3.zero
local world_link_unit = world.link_unit
local light_set_enabled = light.set_enabled
local world_unlink_unit = world.unlink_unit
local world_destroy_unit = world.destroy_unit
local quaternion_identity = quaternion.identity
local world_spawn_unit_ex = world.spawn_unit_ex
local light_set_intensity = light.set_intensity
local light_set_falloff_end = light.set_falloff_end
local light_set_ies_profile = light.set_ies_profile
local unit_set_local_position = unit.set_local_position
local light_set_casts_shadows = light.set_casts_shadows
local light_set_falloff_start = light.set_falloff_start
local light_set_spot_reflector = light.set_spot_reflector
local light_set_spot_angle_end = light.set_spot_angle_end
local light_set_spot_angle_start = light.set_spot_angle_start
local light_color_with_intensity = light.color_with_intensity
local unit_set_vector3_for_materials = unit.set_vector3_for_materials
local light_set_volumetric_intensity = light.set_volumetric_intensity
local light_set_correlated_color_temperature = light.set_correlated_color_temperature

local flashlight_unit_small = "content/weapons/player/attachments/flashlights/flashlight_02/flashlight_02"
local flashlight_unit_large = "content/weapons/player/attachments/flashlights/flashlight_01/flashlight_01"
local packages_to_load = {
    flashlight_unit_large,
    flashlight_unit_small,
}
local flashlight_profiles = {
    small = {
        unit = flashlight_unit_small,
        ies_profile = "content/environment/ies_profiles/narrow/flashlight_custom_02",
        color_temperature = 7300,
        spot_reflector = false,
        intensity= 10,
        spot_angle_start = 0,
        spot_angle_end = 0.6,
        falloff_start = 0,
        falloff_end = 70,
        volumetric_intensity = 0.6,
    },
    large = {
        unit = flashlight_unit_large,
        ies_profile = "content/environment/ies_profiles/narrow/flashlight_custom_03",
        color_temperature = 6200,
        spot_reflector = false,
        intensity= 16,
        spot_angle_start = 0,
        spot_angle_end = 1.5,
        falloff_start = 0,
        falloff_end = 40,
        volumetric_intensity = 0.3,
    },
}

local ServoFriendFlashlightExtension = class("ServoFriendFlashlightExtension", "ServoFriendBaseExtension")

ServoFriendFlashlightExtension.init = function(self, extension_init_context, unit, extension_init_data)
    -- Base class
    ServoFriendFlashlightExtension.super.init(self, extension_init_context, unit, extension_init_data)
    -- Data
    self.event_manager = managers.event
    -- Events
    self.event_manager:register(self, "servo_friend_settings_changed", "on_settings_changed")
    self.event_manager:register(self, "servo_friend_spawned", "on_servo_friend_spawned")
    self.event_manager:register(self, "servo_friend_destroyed", "on_servo_friend_destroyed")
    -- Settings
    self:on_settings_changed()
end

ServoFriendFlashlightExtension.destroy = function(self)
    -- Events
    self.event_manager:unregister(self, "servo_friend_settings_changed")
    self.event_manager:unregister(self, "servo_friend_spawned")
    self.event_manager:unregister(self, "servo_friend_destroyed")
    -- Base class
    ServoFriendFlashlightExtension.super.destroy(self)
end

ServoFriendFlashlightExtension.update = function(self, dt, t)
    -- Base class
    ServoFriendFlashlightExtension.super.update(self, dt, t)
end

ServoFriendFlashlightExtension.on_settings_changed = function(self)
    self.flashlight = mod:get("mod_option_flashlight")
    self.flashlight_shadows = mod:get("mod_option_flashlight_shadows")
    self.flashlight_no_hub = mod:get("mod_option_flashlight_no_hub")
    self.flashlight_type = mod:get("mod_option_flashlight_type")
end

ServoFriendFlashlightExtension.on_servo_friend_spawned = function(self)
    self:spawn_flashlight()
end

ServoFriendFlashlightExtension.on_servo_friend_destroyed = function(self)
    self:destroy_flashlight()
end

ServoFriendFlashlightExtension.spawn_flashlight = function(self)
    local pt = self:pt()
    if pt.world and not pt.flashlight_unit or not unit_alive(pt.flashlight_unit) and self.all_packages_loaded then
        local player_position = self:player_position()
        local flashlight_profile = flashlight_profiles[self.flashlight_type]
        if flashlight_profile then
            local flashlight_unit = flashlight_profile.unit
            pt.flashlight_unit = world_spawn_unit_ex(pt.world, flashlight_unit, nil, player_position, quaternion_identity())
            world_link_unit(pt.world, pt.flashlight_unit, 1, pt.servo_friend_unit, 1)
            unit_set_local_position(pt.flashlight_unit, 1, vector3(.05, 0, 0))

            pt.light = unit_light(pt.flashlight_unit, 1)
            self:set_light()
        end
    end
end

ServoFriendFlashlightExtension.destroy_flashlight = function(self)
    local pt = self:pt()
    if pt.servo_friend_unit and unit_alive(pt.servo_friend_unit) then
        world_unlink_unit(pt.world, pt.flashlight_unit)
        world_destroy_unit(pt.world, pt.flashlight_unit)
    end
end

ServoFriendFlashlightExtension.respawn_flashlight = function(self)
    local pt = self:pt()
    if pt.servo_friend_unit and unit_alive(pt.servo_friend_unit) then
        self:destroy_flashlight()
        self:spawn_flashlight()
    end
end

ServoFriendFlashlightExtension.wants_flashlight_on = function(self)
    local pt = self:pt()
    local hub = not pt.hub or not self.flashlight_no_hub
    local dark_mission = self.flashlight == "only_dark_missions" and pt.dark_mission
    return (self.flashlight == "always_on" or dark_mission) and hub
end

ServoFriendFlashlightExtension.set_light = function(self)
    local pt = self:pt()
    if pt.player_spawned and pt.light then
        light_set_enabled(pt.light, self:wants_flashlight_on())
        light_set_casts_shadows(pt.light, self.flashlight_shadows)
        local template = flashlight_profiles[self.flashlight_type]
        if template then
            light_set_ies_profile(pt.light, template.ies_profile)
            light_set_correlated_color_temperature(pt.light, template.color_temperature)
            light_set_spot_reflector(pt.light, template.spot_reflector)
            light_set_intensity(pt.light, template.intensity)
            light_set_spot_angle_start(pt.light, template.spot_angle_start)
            light_set_spot_angle_end(pt.light, template.spot_angle_end)
            light_set_falloff_start(pt.light, template.falloff_start)
            light_set_falloff_end(pt.light, template.falloff_end)
            light_set_volumetric_intensity(pt.light, template.volumetric_intensity)
            local color = light_color_with_intensity(pt.light) or vector3_zero()
            unit_set_vector3_for_materials(pt.flashlight_unit, "light_color", color)
        end
    end
end