#include "core/compute_workspace.h"

#include <algorithm>
#include <exception>
#include <map>
#include <unordered_map>
#include <unordered_set>

#include "core/ggml_compat.h"
#include "core/ggml_extend_backend.h"
#include "core/util.h"

namespace sd {
    ComputeWorkspace::~ComputeWorkspace() {
        segment_end();
        release();
        ggml_backend_free(cpu_backend_);
    }

    void ComputeWorkspace::set_extra_backends(const std::vector<ggml_backend_t>& backends) {
        if (extra_backends_ != backends) {
            GGML_ASSERT(!active_);
            release();
            extra_backends_ = backends;
        }
    }

    ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
        std::vector<ggml_backend_t> backends{backend_};
        backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
        if (!sd_backend_is_cpu(backend_)) {
            if (cpu_backend_ == nullptr) {
                cpu_backend_ = sd_backend_cpu_init();
            }
            if (cpu_backend_ == nullptr) {
                return nullptr;
            }
            backends.push_back(cpu_backend_);
        }
        std::vector<ggml_backend_buffer_type_t> bufts;
        for (auto backend : backends) {
            auto buft = backend == cpu_backend_
                            ? ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend_))
                            : nullptr;
            bufts.push_back(buft != nullptr ? buft : ggml_backend_get_default_buffer_type(backend));
        }
        return ggml_backend_sched_new(backends.data(), bufts.data(), static_cast<int>(backends.size()),
                                      graph_size, false, false);
    }

    bool ComputeWorkspace::prepare(const Measurement& measurement) {
        GGML_ASSERT(!active_);
        if (measurement.buffers.empty()) {
            return false;
        }
        const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
                                       [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
        if (grows) {
            release();
        }
        return true;
    }

    bool ComputeWorkspace::release_excess(const Measurement& measurement) {
        return std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
                           [&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) &&
               release();
    }

    bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
        GGML_ASSERT(!active_);
        const size_t capacity = static_cast<size_t>(ggml_graph_n_nodes(graph) + ggml_graph_n_leafs(graph)) + 8;
        if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
            release();
            scheduler_          = make_scheduler(capacity);
            scheduler_capacity_ = capacity;
        }
        if (scheduler_ == nullptr) {
            return false;
        }
        ggml_backend_sched_reset(scheduler_);
        assign_nodes(scheduler_, graph);
        // Scheduler allocation rewrites sources. Split the execution graph only once.
        if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
            release();
            return false;
        }
        active_ = true;
        return true;
    }

    ComputeWorkspace::Measurement ComputeWorkspace::measure(
        ggml_cgraph* graph,
        const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
        const AssignNodes& assign_nodes,
        bool cache_measurement) {
        std::vector<const ggml_tensor*> tensors;
        std::unordered_set<const ggml_tensor*> seen;
        auto visit = [&](const ggml_tensor* tensor) {
            if (tensor != nullptr && seen.insert(tensor).second) {
                tensors.push_back(tensor);
            }
        };
        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
            visit(ggml_graph_node(graph, i));
        }
        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
            visit(ggml_graph_leaf(graph, i));
        }
        for (size_t i = 0; i < tensors.size(); ++i) {
            visit(tensors[i]->view_src);
            for (auto source : tensors[i]->src) {
                visit(source);
            }
        }
        std::vector<uint8_t> key;
        if (cache_measurement) {
            key.reserve(tensors.size() * sizeof(ggml_tensor));
            auto append = [&](const auto& value) {
                const auto* bytes = reinterpret_cast<const uint8_t*>(&value);
                key.insert(key.end(), bytes, bytes + sizeof(value));
            };
            append(ggml_graph_n_nodes(graph));
            append(ggml_graph_n_leafs(graph));
            for (auto tensor : tensors) {
                append(tensor);
                append(tensor->type);
                append(tensor->ne);
                append(tensor->nb);
                append(tensor->op);
                append(tensor->op_params);
                append(tensor->flags);
                append(tensor->src);
                append(tensor->view_src);
                append(tensor->view_offs);
                append(tensor->name);
                const auto external = external_backend(tensor);
                append(external);
                if (external == nullptr || tensor->view_src != nullptr) {
                    append(tensor->buffer);
                    append(tensor->data);
                    append(tensor->extra);
                    if (tensor->buffer != nullptr) {
                        const auto usage = ggml_backend_buffer_get_usage(tensor->buffer);
                        const auto type  = ggml_backend_buffer_get_type(tensor->buffer);
                        append(usage);
                        append(type);
                    }
                } else {
                    append(tensor->buffer);
                    append(tensor->data);
                }
                const int32_t uses = ggml_graph_use_count(graph, tensor);
                append(uses);
            }
            // Match the full allocation metadata, including topology and external placement.
            if (key == measurement_key_) {
                return {measured_buffers_};
            }
        }
        const size_t graph_size = tensors.size() + 8;
        auto context            = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
        if (context == nullptr) {
            return {};
        }
        std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
        std::map<std::pair<ggml_backend_t, const void*>, ggml_backend_buffer_t> external_buffers;
        for (auto tensor : tensors) {
            auto copy      = ggml_dup_tensor(context, tensor);
            *copy          = *tensor;
            copies[tensor] = copy;
        }
        for (const auto& entry : copies) {
            auto source    = entry.first;
            auto copy      = entry.second;
            copy->view_src = source->view_src == nullptr ? nullptr : copies.at(source->view_src);
            for (int i = 0; i < GGML_MAX_SRC; ++i) {
                copy->src[i] = source->src[i] == nullptr ? nullptr : copies.at(source->src[i]);
            }
            auto external = external_backend(source);
            if (external != nullptr && source->view_src == nullptr) {
                // Preserve distinct external allocations so overlap checks do not invent aliases.
                const void* storage = source->buffer ? static_cast<const void*>(source->buffer) : source;
                auto& buffer        = external_buffers[{external, storage}];
                if (buffer == nullptr) {
                    buffer = ggml_backend_alloc_buffer(external, 0);
                    GGML_ASSERT(buffer != nullptr);
                    ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
                }
                copy->buffer = buffer;
                // Capability checks inspect alignment even before weights are resident.
                copy->data  = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
                copy->extra = nullptr;
            }
        }
        for (const auto& entry : copies) {
            const auto* source = entry.first;
            if (source->view_src && source->buffer) {
                const auto* owner = source->view_src;
                while (owner->view_src) {
                    owner = owner->view_src;
                }
                if (external_backend(owner) != nullptr) {
                    entry.second->buffer = copies.at(owner)->buffer;
                }
            }
        }
        auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
        auto map_tensor = [](const ggml_tensor* tensor, void* user_data) -> ggml_tensor* {
            const auto& mapped = *static_cast<const std::unordered_map<const ggml_tensor*, ggml_tensor*>*>(user_data);
            const auto found   = mapped.find(tensor);
            return found == mapped.end() ? nullptr : found->second;
        };
        GGML_ASSERT(ggml_graph_cpy_mapped(graph, copy_graph, map_tensor, &copies));
        Measurement result;
        auto scheduler = make_scheduler(graph_size);
        if (scheduler != nullptr) {
            assign_nodes(scheduler, copy_graph);
            std::vector<size_t> sizes(extra_backends_.size() + 2);
            ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data());
            result.buffers.push_back({backend_, sizes[0]});
            for (size_t i = 0; i < extra_backends_.size(); ++i) {
                result.buffers.push_back({extra_backends_[i], sizes[i + 1]});
            }
            ggml_backend_sched_free(scheduler);
        }
        for (const auto& entry : external_buffers) {
            ggml_backend_buffer_free(entry.second);
        }
        ggml_free(context);
        if (cache_measurement && !result.buffers.empty()) {
            measurement_key_  = std::move(key);
            measured_buffers_ = result.buffers;
        }
        return result;
    }

    void ComputeWorkspace::synchronize() const {
        if (scheduler_ != nullptr) {
            ggml_backend_sched_synchronize(scheduler_);
        } else {
            ggml_backend_synchronize(backend_);
        }
    }

    bool ComputeWorkspace::segment_end() noexcept {
        if (!active_) {
            return true;
        }
        // Outer cleanup guards must not retry a failed backend submission.
        active_ = false;
        try {
            synchronize();
            return true;
        } catch (const std::exception& error) {
            LOG_ERROR("%s workspace synchronization failed during segment cleanup: %s",
                      ggml_backend_name(backend_), error.what());
        } catch (...) {
            LOG_ERROR("%s workspace synchronization failed during segment cleanup: unknown exception",
                      ggml_backend_name(backend_));
        }
        return false;
    }

    bool ComputeWorkspace::release() {
        if (active_) {
            return false;
        }
        measurement_key_.clear();
        measured_buffers_.clear();
        ggml_backend_sched_free(scheduler_);
        scheduler_          = nullptr;
        scheduler_capacity_ = 0;
        return true;
    }

    size_t ComputeWorkspace::bytes(ggml_backend_t backend) const {
        if (scheduler_ != nullptr) {
            return ggml_backend_sched_get_buffer_size(scheduler_, backend);
        }
        return 0;
    }
}