declare let DEBUG: boolean

import { h } from './lib.js'
import { dbg_assert, dbg_log } from './log.js'

const STATE_VERSION = 6
const STATE_MAGIC = 0x86768676 | 0
const STATE_INDEX_MAGIC = 0
const STATE_INDEX_VERSION = 1
const STATE_INDEX_TOTAL_LEN = 2
const STATE_INDEX_INFO_LEN = 3
const STATE_INFO_BLOCK_START = 16

const ZSTD_MAGIC = 0xfd2fb528

export class StateLoadError extends Error {
    constructor(msg: string) {
        super(msg)
    }
}

// Minimal interface for the CPU fields state needs
interface StateCpu {
    get_state(): any[]

    set_state(state: any[]): void
    wasm_memory: WebAssembly.Memory
    zstd_create_ctx(len: number): number
    zstd_get_src_ptr(ctx: number): number
    zstd_read(ctx: number, len: number): number
    zstd_read_free(ptr: number, len: number): void
    zstd_free_ctx(ctx: number): void
}

const CONSTRUCTOR_TABLE: Record<string, any> = {
    Map: Map,
    Uint8Array: Uint8Array,
    Int8Array: Int8Array,
    Uint16Array: Uint16Array,
    Int16Array: Int16Array,
    Uint32Array: Uint32Array,
    Int32Array: Int32Array,
    Float32Array: Float32Array,
    Float64Array: Float64Array,
}

function save_object(obj: any, saved_buffers: Uint8Array[]): any {
    if (typeof obj !== 'object' || obj === null) {
        dbg_assert(typeof obj !== 'function')
        return obj
    }

    if (Array.isArray(obj)) {
        return obj.map((x) => save_object(x, saved_buffers))
    }

    if (obj instanceof Map) {
        return {
            __state_type__: 'Map',

            args: Array.from(obj.entries()).map(([k, v]: [any, any]) => [
                save_object(k, saved_buffers),
                save_object(v, saved_buffers),
            ]),
        }
    }

    if (obj.constructor === Object) {
        console.log(obj)
        dbg_assert(obj.constructor !== Object, 'Expected non-object')
    }

    if (obj.BYTES_PER_ELEMENT) {
        // Uint8Array, etc.
        const buffer = new Uint8Array(
            obj.buffer,
            obj.byteOffset,
            obj.length * obj.BYTES_PER_ELEMENT,
        )

        const constructor = obj.constructor.name.replace('bound ', '')

        dbg_assert(CONSTRUCTOR_TABLE[constructor])

        return {
            __state_type__: constructor,
            buffer_id: saved_buffers.push(buffer) - 1,
        }
    }

    if (DEBUG && !obj.get_state) {
        console.log('Object without get_state: ', obj)
    }

    const state = obj.get_state()

    const result: any[] = []

    for (let i = 0; i < state.length; i++) {
        const value = state[i]

        dbg_assert(typeof value !== 'function')

        result[i] = save_object(value, saved_buffers)
    }

    return result
}

function restore_buffers(obj: any, buffers: ArrayBuffer[]): any {
    if (typeof obj !== 'object' || obj === null) {
        dbg_assert(typeof obj !== 'function')
        return obj
    }

    if (Array.isArray(obj)) {
        for (let i = 0; i < obj.length; i++) {
            obj[i] = restore_buffers(obj[i], buffers)
        }

        return obj
    }

    const type = obj['__state_type__']
    dbg_assert(type !== undefined)

    const constructor = CONSTRUCTOR_TABLE[type]
    dbg_assert(constructor, 'Unkown type: ' + type)

    if (obj['args'] !== undefined) {
        return new constructor(obj['args'])
    }

    const buffer = buffers[obj['buffer_id']]
    return new constructor(buffer)
}

export function save_state(cpu: StateCpu): ArrayBuffer {
    const saved_buffers: Uint8Array[] = []
    const state = save_object(cpu, saved_buffers)

    const buffer_infos: { offset: number; length: number }[] = []
    let total_buffer_size = 0

    for (let i = 0; i < saved_buffers.length; i++) {
        const len = saved_buffers[i].byteLength

        buffer_infos[i] = {
            offset: total_buffer_size,
            length: len,
        }

        total_buffer_size += len

        // align
        total_buffer_size = (total_buffer_size + 3) & ~3
    }

    const info_object = JSON.stringify({
        buffer_infos: buffer_infos,
        state: state,
    })
    const info_block = new TextEncoder().encode(info_object)

    let buffer_block_start = STATE_INFO_BLOCK_START + info_block.length
    buffer_block_start = (buffer_block_start + 3) & ~3
    const total_size = buffer_block_start + total_buffer_size

    const result = new ArrayBuffer(total_size)

    const header_block = new Int32Array(result, 0, STATE_INFO_BLOCK_START / 4)
    new Uint8Array(result, STATE_INFO_BLOCK_START, info_block.length).set(
        info_block,
    )
    const buffer_block = new Uint8Array(result, buffer_block_start)

    header_block[STATE_INDEX_MAGIC] = STATE_MAGIC
    header_block[STATE_INDEX_VERSION] = STATE_VERSION
    header_block[STATE_INDEX_TOTAL_LEN] = total_size
    header_block[STATE_INDEX_INFO_LEN] = info_block.length

    for (let i = 0; i < saved_buffers.length; i++) {
        const buf = saved_buffers[i]
        dbg_assert(buf.constructor === Uint8Array)
        buffer_block.set(buf, buffer_infos[i].offset)
    }

    dbg_log('State: json size ' + (info_block.byteLength >> 10) + 'k')
    dbg_log(
        'State: Total buffers size ' + (buffer_block.byteLength >> 10) + 'k',
    )

    return result
}

export function restore_state(cpu: StateCpu, state: ArrayBuffer): void {
    const state_bytes = new Uint8Array(state)

    function read_state_header(
        data: Uint8Array,
        check_length: boolean,
    ): number {
        const len = data.length

        if (len < STATE_INFO_BLOCK_START) {
            throw new StateLoadError('Invalid length: ' + len)
        }

        const header_block = new Int32Array(data.buffer, data.byteOffset, 4)

        if (header_block[STATE_INDEX_MAGIC] !== STATE_MAGIC) {
            throw new StateLoadError(
                'Invalid header: ' + h(header_block[STATE_INDEX_MAGIC] >>> 0),
            )
        }

        if (header_block[STATE_INDEX_VERSION] !== STATE_VERSION) {
            throw new StateLoadError(
                'Version mismatch: dump=' +
                    header_block[STATE_INDEX_VERSION] +
                    ' we=' +
                    STATE_VERSION,
            )
        }

        if (check_length && header_block[STATE_INDEX_TOTAL_LEN] !== len) {
            throw new StateLoadError(
                "Length doesn't match header: " +
                    'real=' +
                    len +
                    ' header=' +
                    header_block[STATE_INDEX_TOTAL_LEN],
            )
        }

        return header_block[STATE_INDEX_INFO_LEN]
    }

    function read_info_block(info_block_buffer: Uint8Array) {
        const info_block = new TextDecoder().decode(info_block_buffer)
        return JSON.parse(info_block)
    }

    if (new Uint32Array(state_bytes.buffer, 0, 1)[0] === ZSTD_MAGIC) {
        const ctx = cpu.zstd_create_ctx(state_bytes.length)

        new Uint8Array(
            cpu.wasm_memory.buffer,
            cpu.zstd_get_src_ptr(ctx) >>> 0,
            state_bytes.length,
        ).set(state_bytes)

        let ptr = cpu.zstd_read(ctx, 16)
        const header_block = new Uint8Array(
            cpu.wasm_memory.buffer,
            ptr >>> 0,
            16,
        )
        const info_block_len = read_state_header(header_block, false)
        cpu.zstd_read_free(ptr, 16)

        ptr = cpu.zstd_read(ctx, info_block_len)
        const info_block_buffer = new Uint8Array(
            cpu.wasm_memory.buffer,
            ptr >>> 0,
            info_block_len,
        )
        const info_block_obj = read_info_block(info_block_buffer)
        cpu.zstd_read_free(ptr, info_block_len)

        let state_object = info_block_obj['state']
        const buffer_infos = info_block_obj['buffer_infos']
        const buffers: ArrayBuffer[] = []

        let position = STATE_INFO_BLOCK_START + info_block_len

        for (const buffer_info of buffer_infos) {
            const front_padding = ((position + 3) & ~3) - position
            const CHUNK_SIZE = 1 * 1024 * 1024

            if (buffer_info.length > CHUNK_SIZE) {
                const chunk_ptr = cpu.zstd_read(ctx, front_padding) >>> 0
                cpu.zstd_read_free(chunk_ptr, front_padding)

                const buffer = new Uint8Array(buffer_info.length)
                buffers.push(buffer.buffer)

                let have = 0
                while (have < buffer_info.length) {
                    const remaining = buffer_info.length - have
                    dbg_assert(remaining >= 0)
                    const to_read = Math.min(remaining, CHUNK_SIZE)

                    const read_ptr = cpu.zstd_read(ctx, to_read)
                    buffer.set(
                        new Uint8Array(
                            cpu.wasm_memory.buffer,
                            read_ptr >>> 0,
                            to_read,
                        ),
                        have,
                    )
                    cpu.zstd_read_free(read_ptr, to_read)

                    have += to_read
                }
            } else {
                const chunk_ptr = cpu.zstd_read(
                    ctx,
                    front_padding + buffer_info.length,
                )
                const offset = (chunk_ptr >>> 0) + front_padding
                buffers.push(
                    cpu.wasm_memory.buffer.slice(
                        offset,
                        offset + buffer_info.length,
                    ),
                )
                cpu.zstd_read_free(
                    chunk_ptr,
                    front_padding + buffer_info.length,
                )
            }

            position += front_padding + buffer_info.length
        }

        state_object = restore_buffers(state_object, buffers)
        cpu.set_state(state_object)

        cpu.zstd_free_ctx(ctx)
    } else {
        const info_block_len = read_state_header(state_bytes, true)

        if (info_block_len < 0 || info_block_len + 12 >= state_bytes.length) {
            throw new StateLoadError(
                'Invalid info block length: ' + info_block_len,
            )
        }

        const info_block_buffer = state_bytes.subarray(
            STATE_INFO_BLOCK_START,
            STATE_INFO_BLOCK_START + info_block_len,
        )
        const info_block_obj = read_info_block(info_block_buffer)
        let state_object = info_block_obj['state']
        const buffer_infos = info_block_obj['buffer_infos']
        let buffer_block_start = STATE_INFO_BLOCK_START + info_block_len
        buffer_block_start = (buffer_block_start + 3) & ~3

        const buffers = buffer_infos.map((buffer_info: any) => {
            const offset = buffer_block_start + buffer_info.offset
            return state_bytes.buffer.slice(offset, offset + buffer_info.length)
        })

        state_object = restore_buffers(state_object, buffers)
        cpu.set_state(state_object)
    }
}