(with_import "./collections.kp"
(let (

    ; Vectors and Values
    ; Bytes encode themselves
    encode_LEB128_helper (rec-lambda recurse (allow_neg x)
        (cond (and allow_neg (< x 0x80)) [x]
                             (< x 0x40)  [x]
              true       (cons (| (& x 0x7F) 0x80) (recurse true (>> x 7))))
    )
    encode_u_LEB128    (lambda (x) (encode_LEB128_helper true x))
    encode_s8_LEB128   (lambda (x) (encode_LEB128_helper (< x 0) (& x 0xFF)))
    encode_s32_LEB128  (lambda (x) (encode_LEB128_helper (< x 0) (& x 0xFFFFFFFF)))
    encode_s33_LEB128  (lambda (x) (encode_LEB128_helper (< x 0) (& x 0x1FFFFFFFF)))
    encode_s64_LEB128  (lambda (x) (encode_LEB128_helper (< x 0) (& x 0xFFFFFFFFFFFFFFFF)))
    encode_vector (lambda (enc v)
        (concat (encode_u_LEB128 (len v)) (flat_map enc v) )
    )
    encode_floating_point (lambda (x) (error "unimplemented"))
    encode_name (lambda (name)
        (encode_vector (lambda (x) [x]) name)
    )
    encode_bytes encode_name

    ; Types
    ; TODO
    encode_limits (lambda (x)
        (cond (= 1 (len x)) (concat [0x00] (encode_u_LEB128 (idx x 0)))
              (= 2 (len x)) (concat [0x01] (encode_u_LEB128 (idx x 0)) (encode_u_LEB128 (idx x 1)))
              true          (error "trying to encode bad limits"))
    )
    encode_number_type (lambda (x)
        (cond   (= x 'i32) [0x7F]
                (= x 'i64) [0x7E]
                (= x 'f32) [0x7D]
                (= x 'f64) [0x7C]
                true       (error (str "bad number type " x)))
    )
    encode_valtype (lambda (x)
        ; we don't handle reference types yet
        (encode_number_type x)
    )
    encode_result_type (lambda (x)
        (encode_vector encode_valtype x)
    )
    encode_function_type (lambda (x)
        (concat [0x60] (encode_result_type (idx x 0))
                       (encode_result_type (idx x 1)))
    )

    ; Modules
    encode_type_section (lambda (x)
        (let (
            encoded (encode_vector encode_function_type x)
        ) (concat [0x01] (encode_u_LEB128 (len encoded)) encoded ))
    )
    encode_import (lambda (import)
        (let (
            (mod_name name type idx) import
        ) (concat (encode_name mod_name)
                  (encode_name name)
                  (cond (= type 'func)   (concat [0x00] (encode_u_LEB128 idx))
                        (= type 'table)  (concat [0x01] (error "can't encode table type"))
                        (= type 'memory) (concat [0x02] (error "can't encode memory type"))
                        (= type 'global) (concat [0x03] (error "can't encode global type"))
                        true                            (error (str "bad import type" type))))
        )
    )
    encode_import_section (lambda (x)
        (let (
            encoded (encode_vector encode_import x)
        ) (concat [0x02] (encode_u_LEB128 (len encoded)) encoded ))
    )

    encode_ref_type (lambda (t) (cond (= t 'funcref)   [0x70]
                                      (= t 'externref) [0x6F]
                                      true             (error (str "Bad ref type " t))))

    encode_table_type (lambda (t) (concat (encode_ref_type (idx t 0)) (encode_limits (idx t 1))))

    encode_table_section (lambda (x)
        (let (
            encoded (encode_vector encode_table_type x)
        ) (concat [0x04] (encode_u_LEB128 (len encoded)) encoded ))
    )
    encode_memory_section (lambda (x)
        (let (
            encoded (encode_vector encode_limits x)
        ) (concat [0x05] (encode_u_LEB128 (len encoded)) encoded ))
    )
    encode_export (lambda (export)
        (let (
            (name type idx) export
        ) (concat (encode_name name)
                  (cond (= type 'func)   [0x00]
                        (= type 'table)  [0x01]
                        (= type 'memory) [0x02]
                        (= type 'global) [0x03]
                        true             (error "bad export type"))
                  (encode_u_LEB128 idx)
        ))
    )
    encode_export_section (lambda (x)
        (let (
            encoded (encode_vector encode_export x)
        ) (concat [0x07] (encode_u_LEB128 (len encoded)) encoded ))
    )

    encode_start_section (lambda (x)
        (cond (= 0 (len x)) []
              (= 1 (len x)) (let (encoded (encode_u_LEB128 (idx x 0))) (concat [0x08] (encode_u_LEB128 (len encoded)) encoded ))
              true          (error (str "bad lenbgth for start section " (len x) " was " x)))
    )

    encode_function_section (lambda (x)
        (let (                                     ; nil functions are placeholders for improted functions
            _ (println "encoding function section " x)
            filtered (filter (lambda (i) (!= nil i)) x)
            _ (println "post filtered " filtered)
            encoded (encode_vector encode_u_LEB128 filtered)
        ) (concat [0x03] (encode_u_LEB128 (len encoded)) encoded ))
    )
    encode_blocktype (lambda (type) (cond (symbol? type)  (encode_valtype type)
                                         (= [] type)      [0x40] ; empty type
                                         true             (encode_s33_LEB128 typ)
                                   ))
    encode_ins (rec-lambda recurse (ins)
        (let (
            op (idx ins 0)
        ) (cond (= op 'unreachable)            [0x00]
                (= op 'nop)                    [0x01]
                (= op 'block)          (concat [0x02] (encode_blocktype (idx ins 1)) (flat_map recurse (idx ins 2)) [0x0B])
                (= op 'loop)           (concat [0x03] (encode_blocktype (idx ins 1)) (flat_map recurse (idx ins 2)) [0x0B])
                (= op 'if)             (concat [0x04] (encode_blocktype (idx ins 1)) (flat_map recurse (idx ins 2)) (if (!= 3 (len ins)) (concat [0x05] (flat_map recurse (idx ins 3)))
                                                                                                                                         [])                                        [0x0B])
                (= op 'br)             (concat [0x0C] (encode_u_LEB128 (idx ins 1)))
                (= op 'br_if)          (concat [0x0D] (encode_u_LEB128 (idx ins 1)))
                ;...
                (= op 'return)                 [0x0F]
                (= op 'call)           (concat [0x10] (encode_u_LEB128 (idx ins 1)))
                ; call_indirect
                ; skipping a bunch
                ; Parametric Instructions
                (= op 'drop)                   [0x1A]
                ; skip
                ; Variable Instructions
                (= op 'local.get)      (concat [0x20] (encode_u_LEB128 (idx ins 1)))
                (= op 'local.set)      (concat [0x21] (encode_u_LEB128 (idx ins 1)))
                (= op 'local.tee)      (concat [0x22] (encode_u_LEB128 (idx ins 1)))
                (= op 'global.get)     (concat [0x23] (encode_u_LEB128 (idx ins 1)))
                (= op 'global.set)     (concat [0x24] (encode_u_LEB128 (idx ins 1)))
                ; table
                ; memory
                (= op 'i32.load)       (concat [0x28] (encode_u_LEB128 (idx ins 1)) (encode_u_LEB128 (idx ins 2)))
                (= op 'i64.load)       (concat [0x29] (encode_u_LEB128 (idx ins 1)) (encode_u_LEB128 (idx ins 2)))
                (= op 'i32.store)      (concat [0x36] (encode_u_LEB128 (idx ins 1)) (encode_u_LEB128 (idx ins 2)))
                (= op 'i64.store)      (concat [0x37] (encode_u_LEB128 (idx ins 1)) (encode_u_LEB128 (idx ins 2)))
                ; Numeric Instructions
                (= op 'i32.const)      (concat [0x41] (encode_s32_LEB128 (idx ins 1)))
                (= op 'i64.const)      (concat [0x42] (encode_s64_LEB128 (idx ins 1)))
                ; skip
                (= op 'i32.add)                [0x6A]
        ))
    )
    encode_expr (lambda (expr) (concat (flat_map encode_ins expr) [0x0B]))
    encode_code (lambda (x)
        (let (
            (locals body) x
            enc_locals (encode_vector (lambda (loc)
                                    (concat (encode_u_LEB128 (idx loc 0)) (encode_valtype (idx loc 1)))) locals)
            enc_expr (encode_expr body)
            code_bytes (concat enc_locals enc_expr)
        ) (concat (encode_u_LEB128 (len code_bytes)) code_bytes))
    )
    encode_code_section (lambda (x)
        (let (
            encoded (encode_vector encode_code x)
        ) (concat [0x0A] (encode_u_LEB128 (len encoded)) encoded ))
    )

    encode_global_type (lambda (t) (concat (encode_valtype (idx t 0)) (cond (= (idx t 1) 'const) [0x00]
                                                                            (= (idx t 1) 'mut)   [0x01]
                                                                            true                 (error (str "bad mutablity " (idx t 1))))))
    encode_global_section (lambda (global_section)
        (let (
            encoded (encode_vector (lambda (x) (concat (encode_global_type (idx x 0)) (encode_expr (idx x 1)))) global_section)
        ) (concat [0x06] (encode_u_LEB128 (len encoded)) encoded ))
    )

    ; only supporting one type of element section for now, active funcrefs with offset
    encode_element (lambda (x) (concat [0x00] (encode_expr (idx x 0)) (encode_vector encode_u_LEB128 (idx x 1))))
    encode_element_section (lambda (x)
        (let (
            encoded (encode_vector encode_element x)
        ) (concat [0x09] (encode_u_LEB128 (len encoded)) encoded ))
    )

    encode_data (lambda (data) (cond (= 2 (len data)) (concat [0x00] (encode_expr (idx data 0)) (encode_bytes (idx data 1)))
                                     (= 1 (len data)) (concat [0x01] (encode_bytes (idx data 0)))
                                     (= 3 (len data)) (concat [0x02] (encode_u_LEB128 (idx data 0)) (encode_expr (idx data 1)) (encode_bytes (idx data 2)))
                                     true             (error (str "bad data" data))))
    encode_data_section (lambda (x)
        (let (
            encoded (encode_vector encode_data x)
        ) (concat [0x0B] (encode_u_LEB128 (len encoded)) encoded ))
    )


    wasm_to_binary (lambda (wasm_code)
        (let (
            (type_section import_section function_section table_section memory_section global_section export_section start_section element_section code_section data_section) wasm_code
            _ (println "type_section" type_section "import_section" import_section "function_section" function_section "memory_section" memory_section "global_section" global_section "export_section" export_section "start_section" start_section "element_section" element_section "code_section" code_section "data_section" data_section)
            magic    [ 0x00 0x61 0x73 0x6D ]
            version  [ 0x01 0x00 0x00 0x00 ]
            type     (encode_type_section type_section)
            import   (encode_import_section import_section)
            function (encode_function_section function_section)
            table    (encode_table_section table_section)
            memory   (encode_memory_section memory_section)
            global   (encode_global_section global_section)
            export   (encode_export_section export_section)
            start    (encode_start_section start_section)
            elem     (encode_element_section element_section)
            code     (encode_code_section code_section)
            data     (encode_data_section data_section)
            ;data_count (let (body (encode_u_LEB128 (len data_section))) (concat [0x0C] (encode_u_LEB128 (len body)) body))
            data_count []
        ) (concat magic version type import function table memory global export data_count start elem code data))
    )

    module (lambda (& args) (let (
        helper (rec-lambda recurse (entries i name_dict type import function table memory global export start elem code data)
            (if (= i (len entries)) [ type import function table memory global export start elem code data]
                (let (
                    (n_d t im f ta m g e s elm c d) ((idx entries i) name_dict type import function table memory global export start elem code data)
                 ) (recurse entries (+ i 1) n_d t im f ta m g e s elm c d))))
    ) (helper args 0 empty_dict [] [] [] [] [] [] [] [] [] [] [])))

    table (vau de (idx_name & limits_type) (lambda (name_dict type import function table memory global export start elem code data)
        [ (put name_dict idx_name (len table)) type import function (concat table [[ (idx limits_type -1) (map (lambda (x) (eval x de)) (slice limits_type 0 -2)) ]]) memory global export start elem code data ]))

    memory (vau de (idx_name & limits) (lambda (name_dict type import function table memory global export start elem code data)
        [ (put name_dict idx_name (len memory)) type import function table (concat memory [(map (lambda (x) (eval x de)) limits)]) global export start elem code data ]))

    func (vau de (name & inside) (lambda (name_dict type import function table memory global export start elem code data)
        (let (
            (params result locals body) ((rec-lambda recurse (i pe re)
                                            (cond (and (= nil pe) (< i (len inside)) (array? (idx inside i)) (< 0 (len (idx inside i))) (= 'param (idx (idx inside i) 0)))
                                                             (recurse (+ i 1) pe re)
                                                  (and (= nil pe) (= nil re) (< i (len inside)) (array? (idx inside i)) (< 0 (len (idx inside i))) (= 'result (idx (idx inside i) 0)))
                                                            ; only one result possible
                                                             (recurse (+ i 1) i (+ i 1))
                                                  (and (= nil re) (< i (len inside)) (array? (idx inside i)) (< 0 (len (idx inside i))) (= 'result (idx (idx inside i) 0)))
                                                            ; only one result possible
                                                             (recurse (+ i 1) pe (+ i 1))
                                                  (and            (< i (len inside)) (array? (idx inside i)) (< 0 (len (idx inside i))) (= 'local (idx (idx inside i) 0)))
                                                             (recurse (+ i 1) pe re)
                                                  true      [ (slice inside 0 (or pe 0)) (slice inside (or pe 0) (or re pe 0)) (slice inside (or re pe 0) i) (slice inside i -1) ]
                                            )
                                        ) 0 nil nil)
            result (if (!= 0 (len result)) (idx result 0)
                                           result)
            _ (println "params " params " result " result " locals " locals " body " body)
            outer_name_dict (put name_dict name (len function))
            (num_params inner_name_dict) (foldl (lambda (a x) [(+ (idx a 0) 1) (put (idx a 1) (idx x 1) (idx a 0))]) [ 0 outer_name_dict ] params)
            (num_locals inner_name_dict) (foldl (lambda (a x) [(+ (idx a 0) 1) (put (idx a 1) (idx x 1) (idx a 0))]) [ num_params inner_name_dict ] locals)
            _ (println "inner name dict" inner_name_dict)
            compressed_locals ((rec-lambda recurse (cur_list cur_typ cur_num i)
                                (cond (and (= i (len locals)) (= 0 cur_num)) cur_list
                                           (= i (len locals))                (concat cur_list [ [cur_num cur_typ] ])
                                      (= cur_typ (idx (idx locals i) 2))     (recurse cur_list                              cur_typ                (+ 1 cur_num) (+ 1 i))
                                      (= nil cur_typ)                        (recurse cur_list                              (idx (idx locals i) 2) 1             (+ 1 i))
                                      true                                   (recurse (concat cur_list [[cur_num cur_typ]]) (idx (idx locals i) 2) 1             (+ 1 i)))
                               ) [] nil 0 0)
            inner_env (add-dict-to-env de (put inner_name_dict 'depth 0))
            our_type [ (map (lambda (x) (idx x 2)) params) (slice result 1 -1) ]
            _ (println "about to get our_code")
            our_code (flat_map (lambda (x) (let (ins (eval x inner_env))
                                             (cond (array? ins) ins
                                                    true       (ins) ; un-evaled function, bare WAT
                                             )))
                                body)
            _ (println "resulting code " our_code)
        ) [
            outer_name_dict
            ; type
            (concat type [ our_type ])
            ; import
            import
            ; function
            (concat function [ (len function) ])
            ; table
            table
            ; memory
            memory
            ; global
            global
            ; export
            export
            ; start
            start
            ; element
            elem
            ; code
            (concat code [ [ compressed_locals our_code ] ])
            ; data
            data
        ])
    ))
    drop      (lambda ()                                                        [['drop]])
    i32.const (lambda (const)                                                   [['i32.const const]])
    i64.const (lambda (const)                                                   [['i64.const const]])
    local.get (lambda (const)                                                   [['local.get const]])
    i32.add   (lambda (& flatten)  (concat (flat_map (lambda (x) x) flatten)    [['i32.add]]))
    i32.load  (lambda (& flatten)  (concat (flat_map (lambda (x) x) flatten)    [['i32.load 2 0]]))
    i64.load  (lambda (& flatten)  (concat (flat_map (lambda (x) x) flatten)    [['i64.load 3 0]]))
    i32.store (lambda (& flatten)  (concat (flat_map (lambda (x) x) flatten)    [['i32.store 2 0]]))
    i64.store (lambda (& flatten)  (concat (flat_map (lambda (x) x) flatten)    [['i64.store 3 0]]))
    flat_eval_ins (lambda (instructions de) (flat_map (lambda (x) (let (ins (eval x de)) (cond (array? ins) ins
                                                                                              true        (ins)))) instructions))
    block_like_body (lambda (name de inner) (let (
                new_depth (+ 1 (eval 'depth de))
                inner_env (add-dict-to-env de [[ name [new_depth] ] [ 'depth new_depth ]])
               ) (flat_eval_ins inner inner_env)))
    block     (vau de (name & inner)                                            [['block [] (block_like_body name de inner)]])
    loop      (vau de (name & inner)                                            [['loop  [] (block_like_body name de inner)]])
    _if       (vau de (name & inner) (let (
                                            (end_idx else_section) (if (= 'else (idx (idx inner -1) 0))      [ -2            (slice (idx inner -1) 1 -1) ]
                                                                                                             [ -1             nil ])
                                            (end_idx then_section) (if (= 'then (idx (idx inner end_idx) 0)) [ (- end_idx 1) (slice (idx inner end_idx) 1 -1) ]
                                                                                                             [ (- end_idx 1) [ (idx inner end_idx) ] ])
                                            flattened (flat_eval_ins (slice inner 0 end_idx) de)
                                            _ (println "flattened " flattened " then_section " then_section " else_section " else_section)
                                            then_block (block_like_body name de then_section)
                                            else_block (if (!= nil else_section) [(block_like_body name de else_section)]
                                                                                 [])
                                           ) (concat flattened                  [(concat ['if [] then_block] else_block)])))

    br        (vau de (b) (let (block (eval b de)) (if (int? block)             [['br block]]
                                                                                [['br (eval [- 'depth (idx block 0)] de)]])))
    br_if     (vau de (b & flatten)  (let (block (eval b de)
                                           block_val (if (int? block) block
                                                                      (eval [- 'depth (idx block 0)] de))
                                           rest (flat_eval_ins flatten de)
                                      )                           (concat rest  [['br_if block_val]])))
    call      (lambda (f & flatten)  (concat (flat_map (lambda (x) x) flatten)  [['call f]]))
    import (vau de (mod_name name t_idx_typ) (lambda (name_dict type import function table memory global export start elem code data) (let (
            _ (if (!= 'func (idx t_idx_typ 0)) (error "only supporting importing functions rn"))
            (import_type idx_name param_type result_type) t_idx_typ
            actual_type_idx (len type)
            actual_type [ (slice param_type 1 -1) (slice result_type 1 -1) ]
        )
        [ (put name_dict idx_name (len function)) (concat type [actual_type]) (concat import [ [mod_name name import_type actual_type_idx] ]) (concat function [nil]) table memory global export start elem code data ])
    ))

    global (vau de (idx_name global_type expr) (lambda (name_dict type import function table memory global export start elem code data)
        [ (put name_dict idx_name (len global))
          type import function table memory
          (concat global [[(if (array? global_type) (reverse global_type) [global_type 'const]) (eval expr de) ]])
          export start elem code data ]
    ))

    export (vau de (name t_v) (lambda (name_dict type import function table memory global export start elem code data)
        [ name_dict type import function table memory global (concat export [ [ name (idx t_v 0) (get-value name_dict (idx t_v 1)) ] ]) start elem code data ]
    ))

    start (vau de (name) (lambda (name_dict type import function table memory global export start elem code data)
        [ name_dict type import function table memory global export (concat start [(get-value name_dict name)]) elem code data ]
    ))

    elem (vau de (offset & entries) (lambda (name_dict type import function table memory global export start elem code data)
        [ name_dict type import function table memory global export start (concat elem [[(eval offset de) (map (lambda (x) (get-value name_dict x)) entries)]]) code data ]
    ))

    data (lambda (& it) (lambda (name_dict type import function table memory global export start elem code data)
                                [name_dict type import function table memory global export start elem code (concat data [it])]))
)
    (provide wasm_to_binary
             module import table memory start elem func global export data
             drop i32.const i64.const local.get i32.add
             i32.load i64.load
             i32.store i64.store
             block loop _if br br_if call)
))