# The contents of this file are subject to the Common Public Attribution License
# Version 1.0 (the “License”); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
# https://myou.dev/licenses/LICENSE-CPAL. The License is based on the Mozilla
# Public License Version 1.1 but Sections 14 and 15 have been added to cover use
# of software over a computer network and provide for limited attribution for
# the Original Developer. In addition, Exhibit A has been modified to be
# consistent with Exhibit B.
#
# Software distributed under the License is distributed on an “AS IS” basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
# the specific language governing rights and limitations under the License.
#
# The Original Code is Myou Engine.
#
# the Original Developer is the Initial Developer.
#
# The Initial Developer of the Original Code is the Myou Engine developers.
# All portions of the code written by the Myou Engine developers are Copyright
# (c) 2024. All Rights Reserved.
#
# Alternatively, the contents of this file may be used under the terms of the
# GNU Affero General Public License version 3 (the [AGPL-3] License), in which
# case the provisions of [AGPL-3] License are applicable instead of those above.
#
# If you wish to allow use of your version of this file only under the terms of
# the [AGPL-3] License and not to allow others to use your version of this file
# under the CPAL, indicate your decision by deleting the provisions above and
# replace them with the notice and other provisions required by the [AGPL-3]
# License. If you do not delete the provisions above, a recipient may use your
# version of this file under either the CPAL or the [AGPL-3] License.

import ../types
import std/strformat
import vmath except Quat

when defined(nimdoc):
    type TYPES* = ProbeInfluenceType | ProbeParallaxType | CubemapProbe | CubemapProbeUniform | SH9Uniform

# Forward declarations and ob type methods
proc newCubemapProbe*(engine: MyouEngine, name: string="camera", scene: Scene = nil,
        influence_type: ProbeInfluenceType = SphereInfluence,
        influence_distance: float32 = 2.5,
        falloff: float32 = 0.2,
        intensity: float32 = 1,
        clipping_start: float32 = 0.8,
        clipping_end: float32 = 40,
        parallax_type: ProbeParallaxType = NoParallax,
        parallax_distance: float32 = 0.0, # 0 means auto
    ): CubemapProbe
func getCubemapSideMatrix*(side: int, position = vec3()): Mat4
proc render_cubemap*(self: CubemapProbe, use_roughness_prefiltering = false, mipmap_shader: Material = nil)
proc render_background_cubemap*(scene: Scene, use_roughness_prefiltering = false, mipmap_shader: Material = nil, world_to_cube_matrix: Mat4 = mat4(), upload_UBO = true)

method is_cubemap_probe*(self: GameObject): bool {.base.} =
    return false
method get_cubemap_probe*(self: GameObject): CubemapProbe {.base.} =
    return nil
method is_cubemap_probe*(self: CubemapProbe): bool =
    return true
method get_cubemap_probe*(self: CubemapProbe): CubemapProbe =
    return self
# End forward declarations and ob type methods

import ./gameobject
import ../postprocessing/effect_shaders
import ../graphics/framebuffer
import ../graphics/render
import ../scene
import ../graphics/texture
import ../graphics/ubo
import arr_ref


proc newCubemapProbe*(engine: MyouEngine, name: string="camera", scene: Scene = nil,
        influence_type: ProbeInfluenceType = SphereInfluence,
        influence_distance: float32 = 2.5,
        falloff: float32 = 0.2,
        intensity: float32 = 1,
        clipping_start: float32 = 0.8,
        clipping_end: float32 = 40,
        parallax_type: ProbeParallaxType = NoParallax,
        parallax_distance: float32 = 0.0, # 0 means auto
    ): CubemapProbe =
    var self = new CubemapProbe
    discard procCall(self.GameObject.initGameObject(engine, name))
    self.ubo_index = -1
    self.influence_type = influence_type
    self.influence_distance = influence_distance
    self.falloff = falloff
    self.intensity = intensity
    self.clipping_start = clipping_start
    self.clipping_end = clipping_end
    self.parallax_type = parallax_type
    self.parallax_distance = if parallax_distance != 0:
        parallax_distance
    else:
        influence_distance
    return self


func getCubemapSideMatrix*(side: int, position = vec3()): Mat4 =
    const CUBEMAP_DIRECTIONS = [
        vec3(1, 0, 0),
        vec3(-1, 0, 0),
        vec3(0, -1, 0), # NOTE: Y is reversed! because cubemaps
        vec3(0, 1, 0),  # have left handed coordinates
        vec3(0, 0, 1),
        vec3(0, 0, -1),
    ]
    const CUBEMAP_UP_VECTORS = [
        vec3(0, -1, 0),
        vec3(0, -1, 0),
        vec3(0, 0, -1), # reversed here too
        vec3(0, 0, 1),  #
        vec3(0, -1, 0),
        vec3(0, -1, 0),
    ]
    let dir = CUBEMAP_DIRECTIONS[side]
    let up = CUBEMAP_UP_VECTORS[side]
    let side = cross(up,-dir)
    return mat4(
        vec4(side, position.x),
        vec4(up, position.y),
        vec4(-dir, position.z),
        vec4(0,0,0,1))

template getCubemapSideAxis*(side: int): int = side div 2

const MIN_ROUGHNESS_LOD = 1

var roughness_prefilter: Material = nil

proc make_roughness_prefilter(engine: MyouEngine) =
    roughness_prefilter = engine.newEffectShader(@["cube"],
        texture_types = @[TexCube],
        library=staticRead("../shaders/cube_prefilter.glsl"),
        code = &"""
            float roughness = inputs[0].lod/(inputs[0].px_lod-{MIN_ROUGHNESS_LOD.float32});
            outColor = vec4(PrefilterEnvMap(cube, roughness, normalize(coord3D)), 1.0);
        """,
    )

template get_roughness_prefilter(engine: MyouEngine): Material =
    if roughness_prefilter == nil:
        make_roughness_prefilter(engine)
    roughness_prefilter

proc render_cubemap*(self: CubemapProbe, use_roughness_prefiltering = false, mipmap_shader: Material = nil) =
    if self.ubo_index == -1:
        self.scene.ensure_cubemaps()
    let cube2world = self.world_matrix * scale(vec3(self.influence_distance))
    let world2cube = cube2world.inverse
    if not defined(release):
        assert self.ubo_index < self.scene.max_cubemaps, "max_cubemaps is too low, make sure you call calculate_max_lights_and_cubemaps()"
    # self.scene.cubemap_UBO.storage(CubemapProbeUniform)[self.ubo_index] = CubemapProbeUniform(
    self.scene.cubemap_UBO.byte_storage.to(CubemapProbeUniform)[self.ubo_index] = CubemapProbeUniform(
        world2cube: world2cube,
        resolution: self.cubemap.width.float32,
        # resolution_inverse_squared: (1/(self.cubemap.width^2)).float32,
        falloff_inv: 1/max(0.0001, self.falloff),
        intensity: self.intensity,
        influence_type: self.influence_type.float32,
        parallax_type: self.parallax_type.float32,
        par_dist_rel_inv: 1 / (self.parallax_distance / self.influence_distance),
        roughness_lod: (self.cubemap.texture.mipmapHigh - MIN_ROUGHNESS_LOD).float32,
    )
    self.engine.renderer.draw_cubemap(self.scene, 
        self.cubemap,
        cube2world,
        world2cube,
        self.clipping_start,
        self.clipping_end,
        false)
    let mipmap_shader = if use_roughness_prefiltering and mipmap_shader == nil:
        self.engine.get_roughness_prefilter()
    else:
        mipmap_shader
    self.cubemap.generate_mipmap(mipmap_shader)
    self.cubemap.disable()

proc render_background_cubemap*(scene: Scene, use_roughness_prefiltering = false, mipmap_shader: Material = nil, world_to_cube_matrix: Mat4 = mat4(), upload_UBO = true) =
    if scene.background_cubemap == nil:
        scene.ensure_cubemaps()
    if not defined(release):
        assert 0 < scene.max_cubemaps, "max_cubemaps is too low, make sure you call calculate_max_lights_and_cubemaps()"
    scene.cubemap_UBO.storage(CubemapProbeUniform)[0] = CubemapProbeUniform(
        world2cube: world_to_cube_matrix,
        resolution: scene.background_cubemap.width.float32,
        # resolution_inverse_squared: (1/(scene.background_cubemap.width^2)).float32,
        falloff_inv: 1,
        intensity: 1,
        parallax_type: -1,
        roughness_lod: (scene.background_cubemap.texture.mipmapHigh - MIN_ROUGHNESS_LOD).float32,
    )
    scene.engine.renderer.draw_cubemap(scene, 
        scene.background_cubemap, mat4(), mat4(), 1, 100, true)
    let mipmap_shader = if use_roughness_prefiltering and mipmap_shader == nil:
        scene.engine.get_roughness_prefilter()
    else:
        mipmap_shader
    scene.background_cubemap.generate_mipmap(mipmap_shader)
    if upload_UBO:
        scene.cubemap_UBO.update()

proc generate_cubemap_mipmaps*(scene: Scene, cubemap: Framebuffer, use_roughness_prefiltering = false, mipmap_shader: Material = nil) =
    let mipmap_shader = if use_roughness_prefiltering and mipmap_shader == nil:
        scene.engine.get_roughness_prefilter()
    else:
        mipmap_shader
    cubemap.generate_mipmap(mipmap_shader)