# 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.

{.experimental: "dotOperators".}

import strutils, tables, algorithm, sequtils, strformat, bitops
import memfiles
# import sugar
import hashes

# TODO!! ADD BOUND CHECKS TO ALL FNODES AND DNA1 PARSING!
# TODO: implement big endian, test 32 bit
type
    BlendFileVal* = object
        mem: pointer
        file_path*: string
        dna: TypeAttrsRef
        blocks: FNodesByAddress
        named_blocks*: NamedBlocks
        struct_to_type: StructToTypeMap
        type_lengths: TypeLengths
        # only needed for debugging and $
        type_names: seq[cstring]
        # if we close the file and we open it again later reusing blocks
        # we need to fix all the pointers relative to the old one
        # NOTE: also check that all blocks still have their own addr?
        old_mem_ptr: pointer
        # end pointer for bounds checking
        endp: pointer
    
    BlendFile* = ref BlendFileVal

    FBlock64 = object
        code: array[4, char]
        size: uint32
        address: uint64
        sdna_index: int32
        count: uint32
        data: char # anything to get its address

    # FBlock32 = object
    #     code: array[4, char]
    #     size: uint32
    #     address: uint32
    #     sdna_index: int32
    #     count: uint32
    #     data: char # anything to get its address
    
    # FBlockRef = ref (FBlock32 or FBlock64)

    FNodesByAddress = ref Table[uint64, FNode]

    NamedBlocks = Table[array[2, char], seq[FNode]]

    TypeLengths = ptr UncheckedArray[uint16]

    FNode* = object
        dna: TypeAttrsRef
        blocks: FNodesByAddress
        # we have to keep a reference to this otherwise it will be GC'd
        # TODO: remove the two pointers above to use less memory?
        blend_file: BlendFile

        p: pointer
        count: uint32
        tid: uint16

    AttrType = enum
        Value
        Pointer
        PointerArray

    AttrOffset = object
        tid, offset, count: uint16
        atype: AttrType

    AttrOffsets = Table[string, AttrOffset]

    TypeAttrsRef = ref seq[AttrOffsets]

    StructToTypeMap = seq[uint16]

    # Note: we're assuming blend files have these data types in this exact order
    BlendBasicType* = enum
        B_char, B_uchar, B_short, B_ushort, B_int32, B_long, B_ulong,
        B_float, B_double, B_int64, B_uint64, B_void, B_int8

proc hash*(n: FNode): Hash =
    # NOTE: assuming read only data (otherwise we should hash the contents)
    hash (n.p, n.count)

# proc `=destroy`(self: var BlendFileVal) =
#     try:
#         self.mem_file.close()
#     except OSError:
#         # this should never happen but compiler requires it for hooks
#         discard

const BASIC_TYPE_LENGTHS = [1, 1, 2, 2, 4, 4, 4, 4, 8, 8, 8, 0, 1]

template `[]`*(blocks: NamedBlocks, s: string): untyped = blocks[[s[0],s[1]]]
# iterator items*(s: NamedBlocks): seq[FNode] {.borrow.}
# iterator pairs*(s: NamedBlocks): (array[2, char], seq[FNode]) {.borrow.}

template fourCC(s:string): untyped = [s[0], s[1], s[2], s[3]]
template fourCC(p:pointer): untyped = cast[ptr array[4,char]](p)[]

# pointer arithmetic
template `$`*(p: pointer): string = cast[uint](p).tohex
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))
template `-`(p, q: pointer): int = cast[int](cast[int](p) -% cast[int](q))
template `+=`(p: pointer, i: Natural) = p = cast[pointer](cast[int](p) +% cast[int](i))

# this is for getting a string that may not be null terminated
# otherwise you can just do $cast[cstring](x.addr)
proc getString(data: pointer, max_size: Natural): string =
    let u8 = cast[ptr UncheckedArray[char]](data)
    for i in 0 ..< max_size:
        if u8[i] == '\0':
            break
        result &= u8[i]

proc build_dna_table(dna_table: ptr FBlock64): (TypeAttrsRef, StructToTypeMap, seq[cstring], TypeLengths)  =
    var type_attrs = new TypeAttrsRef
    var struct_to_type: StructToTypeMap
    let pointer_size = sizeof dna_table.address
    var u8 = cast[ptr UncheckedArray[uint8]](dna_table.data.addr)
    var count = cast[ptr int32](u8[8].addr)[]
    var offset = 12
    var names: seq[cstring]
    for i in 0 ..< count:
        let name = cast[cstring](u8[offset].addr)
        offset += name.len + 1
        names.add name
    
    if (offset and 3) != 0: offset += 4 - (offset and 3)
    assert u8[offset].addr.fourCC == "TYPE".fourCC
    offset += 4 # skip header
    count = cast[ptr int32](u8[offset].addr)[]
    offset += 4
    var types: seq[cstring]
    for i in 0 ..< count:
        let t = cast[cstring](u8[offset].addr)
        offset += t.len + 1
        types.add t
    
    # we'll assume basic types always appear in the same order, to simplify checks
    assert types[0 .. 11] == @["char".cstring, "uchar", "short", "ushort", "int", "long", "ulong", "float", "double", "int64_t", "uint64_t", "void"]
    # also int8_t in newer files
    
    if (offset and 3) != 0: offset += 4 - (offset and 3)
    assert u8[offset].addr.fourCC == "TLEN".fourCC
    offset += 4 # skip header
    var lengths = cast[ptr UncheckedArray[uint16]](u8[offset].addr)
    let basic_lengths = cast[ptr array[12, uint16]](u8[offset].addr)
    assert basic_lengths[] == [1'u16, 1, 2, 2, 4, 4, 4, 4, 8, 8, 8, 0]
    offset += count * 2
    
    if (offset and 3) != 0: offset += 4 - (offset and 3)
    assert u8[offset].addr.fourCC == "STRC".fourCC
    offset += 4 # skip header
    count = cast[ptr int32](u8[offset].addr)[]
    offset += 4
    
    struct_to_type.setLen count
    type_attrs[].setLen types.len

    for i in 0 ..< count:
        let (type_id, count) = cast[ptr (uint16, uint16)](u8[offset].addr)[]
        offset += 4
        var aoffset = 0'u16
        let fields = cast[ptr UncheckedArray[(uint16, uint16)]](u8[offset].addr)
        for j in 0 ..< count.int:
            let (tid, nid) = fields[j]
            var name = $names[nid]
            # if types[type_id] == "Mesh":
            #     dump (name, aoffset, types[tid])
            # stripping ( ) so callback functions are treated just like pointers
            name = name.strip(true, true, {'(',')'})
            let is_pointer = name[0] == '*'
            let is_pointer_array = is_pointer and name[1] == '*'
            var asize: uint16 = if is_pointer: pointer_size.uint16 else: lengths[tid]
            var acount = 1
            if name[^1] == ']':
                let parts = name.split(sep='[')
                name = parts[0]
                for i in 1..parts.high:
                    acount *= parts[i][0 ..< ^1].parseInt
            if is_pointer:
                name = name.strip(true, false, {'*'})
            asize *= acount.uint16
            struct_to_type[i] = type_id
            type_attrs[type_id][name] = AttrOffset(
                tid: tid,
                offset: aoffset,
                count: acount.uint16,
                atype: if is_pointer_array: PointerArray
                elif is_pointer: Pointer
                else: Value
            )
            aoffset += asize
            offset += 4
        # if true:
        # # if lengths[type_id] != aoffset:
        #     echo &"{types[type_id]}: {lengths[type_id]}, {aoffset}"
        #     var offsets = toSeq(type_attrs[type_id].pairs)
        #     for _,(k,v) in offsets.sortedByIt(it[1].offset):
        #         echo &"{types[type_id]}.{k}:\t\t{v.offset}"
        assert lengths[type_id] == aoffset
        # type_attrs[type_id]["(size)"] = AttrOffset(offset: lengths[type_id])
    return (type_attrs, struct_to_type, types, lengths)

proc openBlendFile*(path: string, data: pointer, len: int): BlendFile =
    result = new BlendFile
    # result.mem_file = memfiles.open(path, mode=fmRead)
    result.file_path = path
    # let mem = result.mem_file.mem
    # result.mem = mem
    # let file_length = result.mem_file.size
    let mem = data
    let file_length = len
    assert file_length > 32, "Invalid file size"
    result.endp = mem + file_length
    let u8 = cast[ptr UncheckedArray[uint8]](mem)

    let version = u8[0].addr.getString(12)
    # echo "version: ", version
    assert version.startsWith("BLENDER"), "Not a Blender file"
    let pointer_size: uint16 = if version[7] == '_': 4 else: 8
    let big_endian = version[7] == 'V'
    assert pointer_size == 8, "32 bit files not supported at the moment"
    assert big_endian == false, "Big endian not supported at the moment"
    # let version_number = version[9..11].parseInt

    var blocks: Table[uint64, ptr FBlock64]
    var dna_table: ptr FBlock64

    var offset = 12
    while true:
        let blk = cast[ptr FBlock64](u8[offset].addr)
        offset += 16 + pointer_size.int + blk.size.int
        assert offset < file_length, "invalid or incomplete file"
        if blk.code == "DNA1".fourCC:
            dna_table = blk
            assert blk.size > 16*8, "invalid or incomplete file"
            # we're stopping here not only because Blender also stops,
            # but because we can guarantee that array[16, int64] 
            # can't go beyond the end
            break
        blocks[blk.address] = blk
    
    let (type_attrs, struct_to_type, type_names, type_lengths) = build_dna_table(dna_table)
    # turn all blocks into FNodes so we can have each with references
    # and without using struct IDs
    var node_blocks = new FNodesByAddress
    for address,blk in blocks.pairs:
        let tid = struct_to_type[blk.sdna_index]
        let real_count = blk.size div type_lengths[tid]
        let node = FNode(
            p: blk.data.addr,
            tid: tid,
            count: if blk.count == 1: real_count else: blk.count,
            dna: type_attrs,
            blocks: node_blocks,
            blend_file: result)
        node_blocks[address] = node
        if blk.code[2] == '\0':
            let t = [blk.code[0], blk.code[1]]
            if t notin result.named_blocks:
                result.named_blocks[t] = @[]
            result.named_blocks[t].add(node)
    result.blocks = node_blocks
    result.type_names = type_names
    result.struct_to_type = struct_to_type
    result.type_lengths = type_lengths

proc `[]`[T](a: UncheckedArray[T], r: HSlice): seq[T] =
    for i in r: result.add a[i]

proc `[]`*[T](a: ptr array[16, T], r: HSlice): seq[T] =
    for i in r: result.add a[i]

proc i8*(n: FNode): ptr array[16, int8] =
    assert n.tid == B_int8.uint16 or n.tid == B_char.uint16 #or n.tid == uint16.high
    cast[ptr array[16, int8]](n.p)

proc u8*(n: FNode): ptr array[16, uint8] =
    assert n.tid == B_uchar.uint16 or n.tid == B_char.uint16 #or n.tid == uint16.high
    cast[ptr array[16, uint8]](n.p)

proc i16*(n: FNode): ptr array[16, int16] =
    assert n.tid == B_short.uint16 #or n.tid == uint16.high
    cast[ptr array[16, int16]](n.p)

proc u16*(n: FNode): ptr array[16, uint16] =
    assert n.tid == B_ushort.uint16 #or n.tid == uint16.high
    cast[ptr array[16, uint16]](n.p)

proc i32*(n: FNode): ptr array[16, int32] =
    assert n.tid == B_int32.uint16 #or n.tid == uint16.high
    cast[ptr array[16, int32]](n.p)

# is long and ulong unused?

proc i64*(n: FNode): ptr array[16, int64] =
    assert n.tid == B_int64.uint16 #or n.tid == uint16.high, "uuuh " & $n.tid
    cast[ptr array[16, int64]](n.p)

proc u64*(n: FNode): ptr array[16, uint64] =
    assert n.tid == B_uint64.uint16 #or n.tid == uint16.high
    cast[ptr array[16, uint64]](n.p)

proc f32*(n: FNode): ptr array[16, float32] {.inline.} =
    assert n.tid == B_float.uint16 #or n.tid == uint16.high
    cast[ptr array[16, float32]](n.p)

proc f64*(n: FNode): ptr array[16, float64] =
    assert n.tid == B_double.uint16 #or n.tid == uint16.high
    cast[ptr array[16, float64]](n.p)

proc cstr*(n: FNode): cstring =
    assert n.tid == B_char.uint16 #or n.tid == uint16.high
    cast[cstring](n.p)

template str*(n: FNode): string = $(n.cstr)

template valid*(n: FNode): bool = n.p != nil
template isNil*(n: FNode): bool = n.p == nil

proc strip2*(s: string): string {.inline.} =
    if s.len > 2:
        return s[2..^1]
    return s

# const zero_values: cstring = "\0\0\0\0\0\0\0\0"

template `?.`*(n: FNode, y: untyped): untyped =
    var x = n
    # x.tid = uint16.high
    var v: typeof(x.y)
    if x.valid: v = x.y
    v

proc get*[T](x: ptr array[16, T], y=0): T =
    if x != nil: return x[y]

template get_fblock_of_node(n: FNode): ptr FBlock64 =
    cast[ptr FBlock64](n.p - FBlock64.offsetOf(data))

proc contains*(n: FNode, name: string): bool =
    if n.p == nil:
        return
    return name in n.dna[n.tid]

proc `[]`*(n: FNode, name: string): FNode =
    if n.p == nil:
        return
    assert not n.count.testBit(31), "This is an array, expected index instead of \"" & name & "\""
    let a = n.dna[n.tid][name]
    if a.atype != Value:
        let ptri = cast[ptr uint64](n.p+a.offset)[]
        if ptri == 0:
            return
        when defined(disallow_missing_blocks):
            assert ptri in n.blocks, "Missing block pointing to: " & name
        if ptri notin n.blocks:
            # This may happen if e.g. the user has deleted an image used by a node
            # so we'll just return an invalid node
            echo "Missing block pointing to: " & name
            return
        let blk = n.blocks[ptri]
        var count: uint32 = a.count
        if a.atype == PointerArray:
            # the sdna_index of this block is 0, which would give an incorrect struct size
            # and the count is always 1, so we have to get the actual count from the block size
            count = blk.get_fblock_of_node.size div sizeof(uint64).uint32
            # mark it as pointer array with the last bit
            count.setBit 31
        # use blk.count?
        var node = FNode(p: blk.p, tid: a.tid, count: count,
            dna: n.dna, blocks: n.blocks, blend_file: n.blend_file)
        if a.tid == B_void.uint16:
            node.tid = n.blend_file.struct_to_type[get_fblock_of_node(node).sdna_index]
        return node
    else:
        return FNode(p: n.p+a.offset, tid: a.tid, count: a.count,
            dna: n.dna, blocks: n.blocks, blend_file: n.blend_file)

proc `[]`*(n: FNode, index: Natural): FNode =
    var count = n.count
    if count.testBit 31:
        count.clearBit 31
        assert index.uint32 < count, "Index out of bounds"
        let ptri = cast[ptr UncheckedArray[uint64]](n.p)[index]
        if ptri == 0:
            return
        assert ptri in n.blocks, "Missing block"
        let blk = n.blocks[ptri]
        return FNode(p: blk.p, tid: n.tid, count: 1,
            dna: n.dna, blocks: n.blocks, blend_file: n.blend_file)
    else:
        var n = n
        let size = if n.tid < BASIC_TYPE_LENGTHS.len:
            BASIC_TYPE_LENGTHS[n.tid]
        else:
            # n.dna[n.tid]["(size)"].offset.int
            n.blend_file.type_lengths[n.tid].int
        n.p = n.p + index * size
        n.count = 1
        return n

iterator items*(n: FNode, stride=0): FNode =
    var count = n.count
    if count.testBit 31:
        count.clearBit 31
        let arr = cast[ptr UncheckedArray[uint64]](n.p)
        for i in 0 ..< count:
            let ptri = arr[i]
            if ptri == 0:
                yield FNode()
                continue
            assert ptri in n.blocks, "Missing block"
            let blk = n.blocks[ptri]
            yield FNode(p: blk.p, tid: n.tid, count: 1,
                dna: n.dna, blocks: n.blocks)
    else:
        var n = n
        let stride = if stride == 0:
            n.blend_file.type_lengths[n.tid].int
        else:
            stride
        n.count = 1
        for i in 0 ..< count:
            yield n
            n.p += stride

type FNodeSlice[T] = object
    node: Fnode
    first, last, stride: T

proc len*(node: FNode): uint32 = return node.count and 0x7fffffff

proc `[]`*[U, V: Ordinal](node: FNode; s: HSlice[U, V]): FNodeSlice[V] =
    # if we set both types to be the same and we 
    # pass different types the compilers crashes 
    if node.isNil: return
    let stride = node.blend_file.type_lengths[node.tid].V
    return FNodeSlice[V](node: node, first: s.a.V, last: s.b, stride: stride)

template arr_len*(node: FNode; arr, len: untyped): FNodeSlice[int32] =
    let n = node
    if n.valid and n.len.i32[0] > 0 and n.arr.valid:
        n.arr[0 ..< n.len.i32[0]]
    else:
        var v = FNodeSlice[int32]()
        v.last = -1
        v

iterator items*[T](ns: FNodeSlice[T]): FNode =
    var n = ns.node
    n.p += ns.first.int * ns.stride.int
    n.count = 1
    for i in ns.first.int .. ns.last.int:
        yield n
        n.p += ns.stride

iterator pairs*[T](ns: FNodeSlice[T]): (int, FNode) =
    var n = ns.node
    n.p += ns.first.int * ns.stride
    n.count = 1
    for i in ns.first.int .. ns.last.int:
        yield (i,n)
        n.p += ns.stride

proc `[]`*(ns: FNodeSlice, name: string): FNodeSlice =
    var ns = ns
    ns.node = ns.node[name]
    return ns

template `.`*[T: FNode or FNodeSlice](n: T, field: untyped): T = n[astToStr(field)]
template type_name*(n: FNode): cstring = n.blend_file.type_names[n.tid]

# proc structIdFromDNA(n: FNode): uint16 =
#     # assuming fnode points to the data of a fblock, get its struct id
#     return cast[ptr uint16](n.p - 8)[]

# proc fileOffset(b: BlendFile, n: FNode): int =
#     return cast[int](n.p - b.mem)

proc to_str[T](arr: ptr array[16, T], len: Natural): string =
    if len == 0: return "[]"
    let arr = cast[ptr UncheckedArray[T]](arr)
    result = $arr[0]
    for i in 1 ..< len:
        result &= ", " & $arr[i]
    return "[" & result & "]"

proc `$`*(n: FNode, depth=0, max_depth=3, just_name=false, dump_pointers=false): string =
    let bf = n.blend_file
    if bf.isNil:
        return "(nil)"
    let tname = $bf.type_names[n.tid]
    if n.isNil:
        return tname & " (nil)"
    if just_name or depth >= max_depth:
        return tname & ": ..."
    var ret = @[tname & ":"]
    var depth = depth+1
    let indent = "    ".repeat depth
    if n.count != 1:
        for i in 0 ..< (n.count and 0x7fffffff):
            let kk = n[i]
            let kkstr = `$`(kk, depth, max_depth)
            ret.add indent & "[" & $i & "]: " & kkstr
        return ret.join "\n"
    var attrs = toSeq(n.dna[n.tid].pairs)
    try:
        let name = n.id.name.str
        ret.add indent & "id.name: " & name
    except: discard
    for _,(k,v) in attrs.sortedByIt(it[0]):
        # if k == "(size)" or k.startsWith "_pad":
        if k.startsWith "_pad":
            continue
        if v.atype == Pointer and k in ["next", "prev", "first", "last"]:
            # ret.add indent & k & ": " & `$`(n[k], depth, max_depth, true)
            continue
        let n2 = try:
            n[k]
        except:
            ret.add indent & k & ": error"
            continue
        if n2.isNil:
            ret.add indent & k & ": " & $bf.type_names[v.tid] & " (nil)"
            continue
        if dump_pointers and v.atype == Pointer:
            ret.add indent & k & ".p: 0x" & cast[uint](n2.p).to_hex
        if n2.count != 1 and v.tid.int > BlendBasicType.high.int:
            if v.tid.int == B_char.int:
                ret.add indent & k & ".cstr: \"" & n2.str & "\""
            else:
                ret.add indent & k & ": " & `$`(n2, depth, max_depth)
            continue
        let count = n2.count and 0x7fffffff
        let str = case v.tid.int:
        of B_char.int: k & ".cstr: " & repr(n2.cstr)
        # of B_char.int: k & ".i8: " & $n2.i8.to_str(count)
        of B_uchar.int: k & ".u8: " & $n2.u8.to_str(count)
        of B_short.int: k & ".i16: " & $n2.i16.to_str(count)
        of B_ushort.int: k & ".u16: " & $n2.u16.to_str(count)
        of B_int32.int: k & ".i32: " & $n2.i32.to_str(count)
        of B_float.int: k & ".f32: " & $n2.f32.to_str(count)
        of B_double.int: k & ".f64: " & $n2.f64.to_str(count)
        of B_int64.int: k & ".i64: " & $n2.i64.to_str(count)
        of B_uint64.int: k & ".u64: " & $n2.u64.to_str(count)
        # of B_void.int: k & ".p: 0x" & cast[uint](n2.p).to_hex
        of B_int8.int: k & ".i8: " & $n2.i8.to_str(count)
        else: k & ": " & `$`(n2, depth, max_depth, dump_pointers=dump_pointers)
        ret.add indent & str
    return ret.join "\n"

# #todo don't export these
# template get_ptr*(n: FNode): pointer = n.p
# template get_ptr_u*(n: FNode): int = cast[uint](n.p)
# template get_end_ptr*(n: FNode): pointer = n.blend_file.endp
# template get_end_ptr_u*(n: FNode): pointer = cast[uint](n.blend_file.endp)

# gets an unchecked array but checking bounds don't go over end of file
template get_array*(n: FNode, size: Natural, T: typedesc): ptr UncheckedArray[T] =
    var p = cast[ptr UncheckedArray[T]](n.p)
    if cast[uint](addr(p[size])) > cast[uint](n.blend_file.endp):
        raise newException(RangeDefect, "index out of bounds")
    p

proc set_type*(n: var FNode, s: string) =
    for i,name in n.blend_file.type_names:
        if name == s:
            n.tid = cast[uint16](i)
            return
    raise newException(ValueError, "unknown type")

proc basic_type*(n: FNode): BlendBasicType =
    if n.tid <= BlendBasicType.high.uint16:
        return n.tid.BlendBasicType
    return B_void

proc find_pointer_usage*(n: FNode) =
    var p = n.blend_file.mem
    var q = cast[int](n.p)
    while p < n.blend_file.endp:
        if cast[ptr int](p)[] == q:
            echo "offset ", p - n.blend_file.mem
        p += 1

template `==`*(n, n2: FNode): bool =
    n.p == n2.p


proc debug_dump*(f: BlendFile) =
    echo "start"
    for address,blk in f.blocks:
        echo "0x", address.to_hex, " = ", `$`(blk, 0, 1, dump_pointers=true)
    echo "end"