# 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

# Forward declarations
proc swap_lines(p: pointer, line_stride, line_count: int) {.gcsafe.}
# End forward declarations

import std/strformat
import vmath except Quat, quat
import arr_ref
import float16

import stb_image/read as stbi
const myouConvertHdrToFloat16 {.booldefine.} = true

when not defined(nimdoc):
    import tinyexr

proc stride*(format: TextureFormat): int =
    case format:
    of SRGB_u8: 3
    of SRGB_Alpha_u8: 4
    of R_u8: 1
    of RG_u8: 2
    of RGB_u8: 3
    of RGBA_u8: 4
    of R_u16, R_f16: 2
    of RG_u16, RG_f16: 4
    of RGB_u16, RGB_f16: 6
    of RGBA_u16, RGBA_f16: 8
    of R_f32: 4
    of RG_f32: 8
    of RGB_f32: 12
    of RGBA_f32: 16
    of Depth_u16: 2
    of Depth_u24: 3 # TODO: are you sure?
    # of Depth_u24_s8: 4
    of Depth_f32: 4

proc channel_count*(format: TextureFormat): int =
    case format:
        of R_u8, R_u16, R_f16, R_f32, Depth_u16, Depth_u24, Depth_f32: 1
        of RG_u8, RG_u16, RG_f16, RG_f32: 2
        of SRGB_u8, RGB_u8, RGB_u16, RGB_f16, RGB_f32: 3
        else: 4

func component_type*(format: TextureFormat): DataType =
    case format:
        of SRGB_u8, SRGB_Alpha_u8, R_u8, RG_u8, RGB_u8, RGBA_u8: UByte
        of R_f16, RG_f16, RGB_f16, RGBA_f16: HalfFloat
        of R_f32, RG_f32, RGB_f32, RGBA_f32, Depth_f32: Float
        of R_u16, RG_u16, RGB_u16, RGBA_u16, Depth_u16: UShort
        of Depth_u24: UInt

func resize*(format: TextureFormat, channel_count: int): TextureFormat =
    if format in [SRGB_u8, SRGB_Alpha_u8] and channel_count < 3:
        raise ValueError.newException "Can't resize sRGB to " & $channel_count
    if format in [Depth_u16, Depth_u24, Depth_f32]:
        raise ValueError.newException "Can't resize depth format"
    let f = format.int
    TextureFormat(f - format.channel_count + channel_count)

func format_depth*(tex: Texture): int {.inline.} =
    if tex.tex_type == TexCube: 1
    else: tex.depth

template toOpenArrayByte(p: pointer, a,b: untyped): untyped =
    cast[ptr UncheckedArray[byte]](p).toOpenArray(a,b)

# template toOpenArray[T](p: pointer, a,b: untyped): untyped =
#     cast[ptr UncheckedArray[T]](p).toOpenArray(a,b)

proc f32_to_f16(source: ptr UncheckedArray[float32], dest: ptr UncheckedArray[Float16], len: int) =
    for i in 0 ..< len:
        dest[i] = source[i].tofloat16(clamp=true)

proc getDimensionsFormat*(p: pointer, len: int, min_channels=0): (int, int, TextureFormat) =
    when not defined(nimdoc):
        if isEXR(p, len):
            let dims = getEXRDimensions(p, len)
            return (dims[0], dims[1], RGBA_f16)
        # TODO, IMPORTANT: 2 channels for stb_image means grey+alpha
        # We should handle those cases properly
        var width, height, channels = 0
        if not infoFromMemory(p.toOpenArrayByte(0, len-1), width, height, channels):
            raise ValueError.newException "Could not read image"
        channels = max(min_channels, channels)
        let hdr = isHDRFromMemory(p.toOpenArrayByte(0, len-1))
        let is16 = is16BitFromMemory(p.toOpenArrayByte(0, len-1))
        # Calculate format with channels, and whether it's hdr or 16 bit
        assert (RG_u8.int - R_u8.int) == 1 # (just in case someone changes the enum)
        const toHDR = when myouConvertHdrToFloat16: (R_f16.int-R_u8.int) else: (R_f32.int-R_u8.int)
        let format = (R_u8.int - 1 + channels +
            hdr.int * toHDR + is16.int * (R_u16.int-R_u8.int)).TextureFormat
        return (width, height, format)

proc loadFileFromSlices*(tex: Texture, slices: seq[SliceMem[byte]],
                        callback: proc(tex: Texture, data: SliceMem[byte]),
                        flip = true, min_channels = 0) {.gcsafe.} =

    when not defined(nimdoc):
        assert tex.tex_type != TexCube, "Loading a cube texture from file is not supported yet"
        let format = if min_channels <= tex.format.channel_count:
            tex.format
        else:
            tex.format.resize(min_channels)
        # TODO: Don't do this!!
        # Change the format when the decoder has detected less channels!!
        let min_channels = format.channel_count

        let layer_stride = tex.width * tex.height * format.stride
        assert tex.depth == slices.len
        var out_buffer = newSliceMem[byte](layer_stride * tex.depth)
        var pos = 0
        for slice in slices:
            var image: imagePixelData[byte]
            var image_16: imagePixelData[uint16]
            var image_f: imagePixelData[float32]
            var buffer: ArrRef[byte]
            # a reference to this pointer is kept with one of the vars above
            var pixels_ptr: pointer
            var pixels_len: int
            var flip = flip
            if isEXR(slice.data, slice.byte_len):
                let (width, height, pixels) = decodeEXR(slice.data, slice.byte_len)
                assert width == tex.width and height == tex.height, "Image size mismatch"
                buffer = pixels.to byte
                pixels_ptr = buffer[0].addr
                pixels_len = buffer.len
            else:
                setFlipVerticallyOnLoad(flip)
                flip = false
                var w,h,c = 0
                if isHDRFromMemory(slice.toOpenArrayByte):
                    image_f = loadFFromMemory(slice.toOpenArrayByte, w,h,c, min_channels)
                    pixels_ptr = image_f.data
                    pixels_len = image_f.byteLen
                    when myouConvertHdrToFloat16:
                        f32_to_f16(
                            cast[ptr UncheckedArray[float32]](pixels_ptr),
                            cast[ptr UncheckedArray[Float16]](pixels_ptr),
                            image_f.len)
                        pixels_len = pixels_len div 2
                elif is16BitFromMemory(slice.toOpenArrayByte):
                    image_16 = load16FromMemory(slice.toOpenArrayByte, w,h,c, min_channels)
                    pixels_ptr = image_16.data
                    pixels_len = image_16.byteLen
                else:
                    image = loadFromMemory(slice.toOpenArrayByte, w,h,c, min_channels)
                    pixels_ptr = image.data
                    pixels_len = image.len
            if layer_stride != pixels_len:
                echo "Format: ", format
                raise Defect.newException &"Image '{tex.name}' has a length" &
                                    &" of {pixels_len}, expected {layer_stride}"
            # TODO: make a swap_lines that copies them elsewhere
            # to avoid copying two times
            if flip:
                swap_lines(pixels_ptr, tex.width * format.stride, tex.height)
            copyMem(out_buffer[pos].addr, pixels_ptr, layer_stride)
            pos += layer_stride
        callback(tex, out_buffer)


proc swap_lines(p: pointer, line_stride, line_count: int) {.gcsafe.} =
    template `+`(p: pointer, i: Natural): pointer = cast[pointer](cast[int](p) +% cast[int](i))
    template `-`(p: pointer, i: Natural): pointer = cast[pointer](cast[int](p) -% cast[int](i))
    let int_stride = line_stride div sizeof(int)
    let int_stride_bytes = int_stride * sizeof(int)
    var p1 = p
    var p2 = p + line_stride*(line_count-1)
    var a1, a2: ptr UncheckedArray[int]
    var b1, b2: ptr UncheckedArray[byte]
    for i in 0 ..< line_count div 2:
        a1 = cast[ptr UncheckedArray[int]](p1)
        a2 = cast[ptr UncheckedArray[int]](p2)
        b1 = cast[ptr UncheckedArray[byte]](p1)
        b2 = cast[ptr UncheckedArray[byte]](p2)
        for j in 0 ..< int_stride:
            swap(a1[j], a2[j])
        for j in int_stride_bytes ..< line_stride:
            swap(b1[j], b2[j])
        p1 = p1 + line_stride
        p2 = p2 - line_stride