#include "HalideRuntime.h"
#include "device_buffer_utils.h"
#include "device_interface.h"
#include "printer.h"
#include "scoped_mutex_lock.h"

extern "C" {

extern void *malloc(size_t);
extern void free(void *);

}

namespace Halide { namespace Runtime { namespace Internal {

struct device_handle_wrapper {
    uint64_t device_handle;
    const halide_device_interface_t *interface;
};

// TODO: Coarser grained locking, also consider all things that need
// to be atomic with respect to each other. At present only
// halide_copy_to_host and halide_copy_to_device are atomic with
// respect to each other. halide_device_malloc and halide_device_free
// are also candidates, but to do so they likely need to be able to do
// a copy internaly as well.
WEAK halide_mutex device_copy_mutex;

WEAK int copy_to_host_already_locked(void *user_context, struct halide_buffer_t *buf) {
    if (!buf->device_dirty()) {
        return 0;  // my, that was easy
    }

    debug(user_context) << "copy_to_host_already_locked " << buf << " dev_dirty is true\n";
    const halide_device_interface_t *interface = buf->device_interface;
    if (buf->host_dirty()) {
        debug(user_context) << "copy_to_host_already_locked " << buf << " dev_dirty and host_dirty are true\n";
        return halide_error_code_copy_to_host_failed;
    }
    if (interface == NULL) {
        debug(user_context) << "copy_to_host_already_locked " << buf << " interface is NULL\n";
        return halide_error_code_no_device_interface;
    }
    int result = interface->impl->copy_to_host(user_context, buf);
    if (result != 0) {
        debug(user_context) << "copy_to_host_already_locked " << buf << " device copy_to_host returned an error\n";
        return halide_error_code_copy_to_host_failed;
    }
    buf->set_device_dirty(false);
    halide_msan_annotate_buffer_is_initialized(user_context, buf);

    return result;
}

}}} // namespace Halide::Runtime::Internal

extern "C" {

/** Release all data associated with the current GPU backend, in particular
 * all resources (memory, texture, context handles) allocated by Halide. Must
 * be called explicitly when using AOT compilation. */
WEAK void halide_device_release(void *user_context, const halide_device_interface_t *device_interface) {
    device_interface->impl->device_release(user_context);
}

/** Copy image data from device memory to host memory. This must be called
 * explicitly to copy back the results of a GPU-based filter. */
WEAK int halide_copy_to_host(void *user_context, struct halide_buffer_t *buf) {
    ScopedMutexLock lock(&device_copy_mutex);

    debug(NULL) << "halide_copy_to_host " << buf << "\n";

    return copy_to_host_already_locked(user_context, buf);
}

/** Copy image data from host memory to device memory. This should not be
 * called directly; Halide handles copying to the device automatically. */
WEAK int halide_copy_to_device(void *user_context,
                               struct halide_buffer_t *buf,
                               const halide_device_interface_t *device_interface) {
    int result = 0;

    ScopedMutexLock lock(&device_copy_mutex);

    debug(user_context)
        << "halide_copy_to_device " << buf
        << ", host: " << buf->host
        << ", dev: " << buf->device
        << ", host_dirty: " << buf->host_dirty()
        << ", dev_dirty: " << buf->device_dirty() << "\n";
    if (device_interface == NULL) {
        debug(user_context) << "halide_copy_to_device " << buf << " interface is NULL\n";
        if (buf->device_interface == NULL) {
            debug(user_context) << "halide_copy_to_device " << buf << " no interface error\n";
            return halide_error_code_no_device_interface;
        }
        device_interface = buf->device_interface;
    }

    if (buf->device && buf->device_interface != device_interface) {
        debug(user_context) << "halide_copy_to_device " << buf << " flipping buffer to new device\n";
        if (buf->device_interface != NULL && buf->device_dirty()) {
            halide_assert(user_context, !buf->host_dirty());
            result = copy_to_host_already_locked(user_context, buf);
            if (result != 0) {
                debug(user_context) << "halide_copy_to_device " << buf << " flipping buffer halide_copy_to_host failed\n";
                return result;
            }
        }
        result = halide_device_free(user_context, buf);
        if (result != 0) {
            debug(user_context) << "halide_copy_to_device " << buf << " flipping buffer halide_device_free failed\n";
            return result;
        }
        buf->set_host_dirty(true); // force copy back to new device below.
    }

    if (buf->device == 0) {
        result = halide_device_malloc(user_context, buf, device_interface);
        if (result != 0) {
            debug(user_context) << "halide_copy_to_device " << buf
                                << " halide_copy_to_device call to halide_device_malloc failed\n";
            return result;
        }
    }

    if (buf->host_dirty()) {
        debug(user_context) << "halide_copy_to_device " << buf << " host is dirty\n";
        if (buf->device_dirty()) {
            debug(user_context) << "halide_copy_to_device " << buf << " dev_dirty is true error\n";
            return halide_error_code_copy_to_device_failed;
        } else {
            result = device_interface->impl->copy_to_device(user_context, buf);
            if (result == 0) {
                buf->set_host_dirty(false);
            } else {
                debug(user_context) << "halide_copy_to_device "
                                    << buf << "device copy_to_device returned an error\n";
                return halide_error_code_copy_to_device_failed;
            }
        }
    }

    return 0;
}

/** Wait for current GPU operations to complete. Calling this explicitly
 * should rarely be necessary, except maybe for profiling. */
WEAK int halide_device_sync(void *user_context, struct halide_buffer_t *buf) {
    const halide_device_interface_t *device_interface = NULL;
    if (buf) {
        device_interface = buf->device_interface;
    }
    if (device_interface == NULL) {
        debug(user_context) << "halide_device_sync on buffer with no interface\n";
        return halide_error_code_no_device_interface;
    }
    int result = device_interface->impl->device_sync(user_context, buf);
    if (result) {
        return halide_error_code_device_sync_failed;
    } else {
        return 0;
    }
}

/** Allocate device memory to back a halide_buffer_t. */
WEAK int halide_device_malloc(void *user_context, struct halide_buffer_t *buf,
                              const halide_device_interface_t *device_interface) {
    const halide_device_interface_t *current_interface = buf->device_interface;
    debug(user_context) << "halide_device_malloc: " << buf
                        << " interface " << device_interface
                        << " host: " << buf->host
                        << ", dev: " << buf->device
                        << ", host_dirty: " << buf->host_dirty()
                        << ", dev_dirty:" << buf->device_dirty()
                        << " buf current interface: " << current_interface << "\n";

    // halide_device_malloc does not support switching interfaces.
    if (current_interface != NULL && current_interface != device_interface) {
        error(user_context) << "halide_device_malloc doesn't support switching interfaces\n";
        return halide_error_code_device_malloc_failed;
    }

    // Ensure code is not freed prematurely.
    // TODO: Exception safety...
    device_interface->impl->use_module();
    int result = device_interface->impl->device_malloc(user_context, buf);
    device_interface->impl->release_module();

    if (result) {
        return halide_error_code_device_malloc_failed;
    } else {
        return 0;
    }
}

/** Free any device memory associated with a halide_buffer_t. */
WEAK int halide_device_free(void *user_context, struct halide_buffer_t *buf) {
    debug(user_context) << "halide_device_free: " << buf
                        << " buf dev " << buf->device
                        << " interface " << buf->device_interface << "\n";
    if (buf != NULL) {
        const halide_device_interface_t *device_interface = buf->device_interface;
        if (device_interface != NULL) {
            // Ensure interface is not freed prematurely.
            // TODO: Exception safety...
            device_interface->impl->use_module();
            int result = device_interface->impl->device_free(user_context, buf);
            device_interface->impl->release_module();
            halide_assert(user_context, buf->device == 0);
            if (result) {
                return halide_error_code_device_free_failed;
            } else {
                return 0;
            }
        }
    }
    buf->set_device_dirty(false);
    return 0;
}

WEAK int halide_weak_device_free(void *user_context, struct halide_buffer_t *buf) {
    return halide_device_free(user_context, buf);
}

/** Free any device memory associated with a halide_buffer_t and ignore any
 * error. Used when freeing as a destructor on an error. */
WEAK void halide_device_free_as_destructor(void *user_context, void *obj) {
    struct halide_buffer_t *buf = (struct halide_buffer_t *)obj;
    halide_device_free(user_context, buf);
}

/** Allocate host and device memory to back a buffer_t. Ideally this
 * will be a zero copy setup, but the default implementation may
 * separately allocate the host memory using halide_malloc and the
 * device memory using halide_device_malloc. */
WEAK int halide_device_and_host_malloc(void *user_context, struct halide_buffer_t *buf,
                                       const halide_device_interface_t *device_interface) {
    const halide_device_interface_t *current_interface = buf->device_interface;
    debug(user_context) << "halide_device_and_host_malloc: " << buf
                        << " interface " << device_interface
                        << " host: " << buf->host
                        << ", device: " << buf->device
                        << ", host_dirty: " << buf->host_dirty()
                        << ", dev_dirty:" << buf->device_dirty()
                        << " buf current interface: " << current_interface << "\n";

    // halide_device_malloc does not support switching interfaces.
    if (current_interface != NULL && current_interface != device_interface) {
        halide_error(user_context, "halide_device_and_host_malloc doesn't support switching interfaces\n");
        return halide_error_code_device_malloc_failed;
    }

    // Ensure code is not freed prematurely.
    // TODO: Exception safety...
    device_interface->impl->use_module();
    int result = device_interface->impl->device_and_host_malloc(user_context, buf);
    device_interface->impl->release_module();

    if (result) {
        halide_error(user_context, "allocating host and device memory failed\n");
        return halide_error_code_device_malloc_failed;
    } else {
        return 0;
    }
}

/** Free host and device memory associated with a buffer_t. */
WEAK int halide_device_and_host_free(void *user_context, struct halide_buffer_t *buf) {
    debug(user_context) << "halide_device_and_host_free: " << buf
                        << " buf dev " << buf->device
                        << " interface " << buf->device_interface << "\n";
    if (buf != NULL) {
        const halide_device_interface_t *device_interface = buf->device_interface;
        if (device_interface != NULL) {
            // Ensure interface is not freed prematurely.
            // TODO: Exception safety...
            device_interface->impl->use_module();
            int result = device_interface->impl->device_and_host_free(user_context, buf);
            device_interface->impl->release_module();
            halide_assert(user_context, buf->device == 0);
            if (result) {
                return halide_error_code_device_free_failed;
            } else {
                return 0;
            }
        } else if (buf->host) {
            // device_free must have been called on this buffer (which
            // must be legal for the device interface that was
            // used). We'd better still free the host pointer.
            halide_free(user_context, buf->host);
            buf->host = NULL;
        }
    }
    buf->set_device_dirty(false);
    return 0;
}

WEAK int halide_default_device_and_host_malloc(void *user_context, struct halide_buffer_t *buf,
                                               const halide_device_interface_t *device_interface) {
    size_t size = buf->size_in_bytes();
    buf->host = (uint8_t *)halide_malloc(user_context, size);
    if (buf->host == NULL) {
        return -1;
    }
    int result = halide_device_malloc(user_context, buf, device_interface);
    if (result != 0) {
        halide_free(user_context, buf->host);
        buf->host = NULL;
    }
    return result;
}

WEAK int halide_default_device_and_host_free(void *user_context, struct halide_buffer_t *buf,
                                             const halide_device_interface_t *device_interface) {
    int result = halide_device_free(user_context, buf);
    if (buf->host) {
        halide_free(user_context, buf->host);
        buf->host = NULL;
    }
    buf->set_host_dirty(false);
    buf->set_device_dirty(false);
    return result;
}


WEAK int halide_device_wrap_native(void *user_context, struct halide_buffer_t *buf, uint64_t handle,
                                   const halide_device_interface_t *device_interface) {
    const halide_device_interface_t *current_interface = buf->device_interface;

    if (current_interface != NULL && current_interface != device_interface) {
        error(user_context) << "halide_device_wrap_native doesn't support switching interfaces\n";
        return halide_error_code_device_wrap_native_failed;
    }

    device_interface->impl->use_module();
    buf->device_interface = device_interface;
    int result = device_interface->impl->wrap_native(user_context, buf, handle);
    device_interface->impl->release_module();

    if (result) {
        return halide_error_code_device_malloc_failed;
    } else {
        return 0;
    }
}

WEAK int halide_device_detach_native(void *user_context, struct halide_buffer_t *buf) {
    const halide_device_interface_t *device_interface = buf->device_interface;
    int result = 0;
    if (device_interface != NULL) {
        device_interface->impl->use_module();
        result = device_interface->impl->detach_native(user_context, buf);
        device_interface->impl->release_module();
        halide_assert(user_context, buf->device == 0);
        if (result) {
            result = halide_error_code_device_detach_native_failed;
        }
    }
    return result;
}

WEAK int halide_default_device_wrap_native(void *user_context, struct halide_buffer_t *buf, uint64_t handle) {
    buf->device_interface->impl->use_module();
    buf->device = handle;
    return 0;
}

WEAK int halide_default_device_detach_native(void *user_context, struct halide_buffer_t *buf) {
    if (buf->device == 0) {
        return 0;
    }
    buf->device_interface->impl->release_module();
    buf->device = 0;
    buf->device_interface = NULL;
    return 0;
}

/** Free any host and device memory associated with a buffer_t and ignore any
 * error. Used when freeing as a destructor on an error. */
WEAK void halide_device_and_host_free_as_destructor(void *user_context, void *obj) {
    struct halide_buffer_t *buf = (struct halide_buffer_t *)obj;
    halide_device_and_host_free(user_context, buf);
}

/** TODO: Find a way to elide host free without this hack. */
WEAK void halide_device_host_nop_free(void *user_context, void *obj) {
}

} // extern "C" linkage