#include "tensor/GpuBackend.h"
#include "tensor/Graph.h"
#include "tensor/KernelGen.h"
#include "tensor/KernelGenWgsl.h"
#include "tensor/Optimizer.h"

#include "gpgpu/ComputeShader.h"
#include "gpgpu/Gpgpu.h"
#include "gpgpu/GpuBuffer.h"
#include "gpgpu/Sequence.h"

#include "common/Exception.h"
#include "common/config.h"

#include <algorithm>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <limits>
#include <map>
#include <memory>
#include <vector>

namespace eve::tensor {
namespace {

constexpr int kLocalSize = 256;
constexpr int kAutotuneIters = 5;

#ifdef EVENGINE_WEBGPU
const char *kReduceSource = R"(
@group(0) @binding(0) var<storage, read_write> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> partial: array<f32>;
struct Params { data: array<vec4<f32>, 8>, }
@group(0) @binding(8) var<uniform> pc: Params;
var<workgroup> scratch: array<f32, 256>;
@compute @workgroup_size(256)
fn main(@builtin(local_invocation_id) lid: vec3<u32>,
        @builtin(global_invocation_id) gid: vec3<u32>,
        @builtin(workgroup_id) wid: vec3<u32>) {
    let size = u32(pc.data[0].x + 0.5);
    let op = u32(pc.data[0].y + 0.5);
    scratch[lid.x] = select(select(-3.402823e38f, 3.402823e38f, op == 1u), 0.0f, op == 0u);
    if (gid.x < size) { scratch[lid.x] = input[gid.x]; }
    workgroupBarrier();
    for (var step = 128u; step > 0u; step >>= 1u) {
        if (lid.x < step) {
            if (op == 0u) { scratch[lid.x] += scratch[lid.x + step]; }
            else if (op == 1u) { scratch[lid.x] = min(scratch[lid.x], scratch[lid.x + step]); }
            else { scratch[lid.x] = max(scratch[lid.x], scratch[lid.x + step]); }
        }
        workgroupBarrier();
    }
    if (lid.x == 0u) { partial[wid.x] = scratch[0]; }
}
)";
#else
const char *kReduceSource = R"(#version 450
layout(local_size_x = 256) in;
layout(set = 0, binding = 0) readonly buffer In { float a[]; };
layout(set = 0, binding = 1) writeonly buffer Out { float partial[]; };
layout(push_constant) uniform PC { float data[32]; } pc;
shared float sdata[256];
void main() {
    uint tid = gl_LocalInvocationID.x;
    uint gid = gl_GlobalInvocationID.x;
    uint size = uint(pc.data[0] + 0.5);
    int op = int(pc.data[1] + 0.5);
    float ident = (op == 0) ? 0.0 : (op == 1 ? 3.402823e38 : -3.402823e38);
    sdata[tid] = (gid < size) ? a[gid] : ident;
    barrier();
    for (uint s = 128u; s > 0u; s >>= 1u) {
        if (tid < s) {
            if (op == 0) sdata[tid] += sdata[tid + s];
            else if (op == 1) sdata[tid] = min(sdata[tid], sdata[tid + s]);
            else sdata[tid] = max(sdata[tid], sdata[tid + s]);
        }
        barrier();
    }
    if (tid == 0u) partial[gl_WorkGroupID.x] = sdata[0];
}
)";

#endif

struct ReduceKernels {
    gpgpu::Gpgpu *gpgpu = nullptr;
    gpgpu::ComputeShader *reduce = nullptr;
};

/** Lazily created, process-lifetime singleton. Returns nullptr if unavailable. */
ReduceKernels *getReduceKernels() {
    static ReduceKernels *kernels = nullptr;
    static bool failed = false;
    if (kernels) return kernels;
    if (failed) return nullptr;
    try {
        auto *gp = gpgpu::Gpgpu::create();
        if (!gp || !gp->isAvailable()) {
            failed = true;
            return nullptr;
        }
        auto *k = new ReduceKernels();
        k->gpgpu = gp;
        k->reduce = gp->newShader(kReduceSource);
        kernels = k;
        return kernels;
    } catch (...) {
        failed = true;
        return nullptr;
    }
}

int groupsFor(int count) { return (count + kLocalSize - 1) / kLocalSize; }

Result<KernelSpec> backendKernel(const Graph &graph, const FusedGroup &group,
                                 KernelVariant variant = KernelVariant::Default) {
#ifdef EVENGINE_WEBGPU
    return generateWgslKernel(graph, group, variant);
#else
    KernelSpec spec;
    const bool generated = variant == KernelVariant::TiledMatMul ? generateMatMulVariant(graph, group, true, spec)
                                                                 : generateKernel(graph, group, spec);
    if (!generated)
        return Result<KernelSpec>::failure(
            Diagnostic::error(DiagnosticCode::Unsupported, "Tensor GLSL kernel generation failed", "tensor.kernel"));
    return Result<KernelSpec>::success(std::move(spec));
#endif
}

/** Bind a single-pass kernel to its real input/output arena buffers. */
void bindKernel(gpgpu::ComputeShader *shader, const KernelSpec &spec,
                const std::vector<gpgpu::GpuBuffer *> &inputs, gpgpu::GpuBuffer *output) {
    for (int i = 0; i < spec.inputCount; ++i)
        shader->bindBuffer(i,
                           spec.inputRepresentatives.empty() || spec.inputRepresentatives[static_cast<size_t>(i)] == i
                               ? inputs[static_cast<size_t>(i)]
                               : nullptr);
    shader->bindBuffer(spec.inputCount, output);
}

double timeDispatch(gpgpu::Gpgpu *gp, gpgpu::ComputeShader *shader, int gx, int gy) {
    const auto t0 = std::chrono::steady_clock::now();
    for (int i = 0; i < kAutotuneIters; ++i) gp->dispatch(shader, gx, gy, 1);
    const auto t1 = std::chrono::steady_clock::now();
    return std::chrono::duration<double, std::milli>(t1 - t0).count() / kAutotuneIters;
}

}  // namespace

struct GpuProgram::Impl {
    struct GroupRuntime {
        KernelSpec spec;
        std::unique_ptr<gpgpu::ComputeShader> pass1;
        std::unique_ptr<gpgpu::ComputeShader> pass2;
        std::vector<gpgpu::GpuBuffer *> inputs;   // one per group input node
        std::vector<gpgpu::GpuBuffer *> stats;    // statsCount working buffers
        gpgpu::GpuBuffer *qScales = nullptr;      // per-group scales (int8/int4)
        gpgpu::GpuBuffer *output = nullptr;
    };

    gpgpu::Gpgpu *gpgpu = nullptr;
    std::vector<GroupRuntime> groups;
    std::vector<gpgpu::GpuBuffer *> slotBuffer;   // arena slot -> buffer
    std::vector<gpgpu::GpuBuffer *> placeholderBuffers;
    std::vector<int> placeholderSizes;
    std::vector<std::unique_ptr<gpgpu::GpuBuffer>> ownedBuffers;
    std::map<int, gpgpu::GpuBuffer *> qScalesByNode;  // quantized const node -> scales
    gpgpu::GpuBuffer *outputBuffer = nullptr;
    int outputSize = 0;
    std::unique_ptr<gpgpu::Sequence> sequence;
    std::unique_ptr<gpgpu::GpuBuffer> outputStaging;

    ~Impl() {
        // Retire command/descriptor references before their owning buffers.
        sequence.reset();
        groups.clear();
        outputStaging.reset();
    }

    gpgpu::GpuBuffer *alloc(int byteSize) {
        auto  buffer   = std::unique_ptr<gpgpu::GpuBuffer>(gpgpu->newBuffer(byteSize, "storage"));
        auto* borrowed = buffer.get();
        ownedBuffers.push_back(std::move(buffer));
        return borrowed;
    }

    /** Bind real arena buffers once at build time; bindings persist across runs. */
    void bindGroup(const GroupRuntime &g) const {
        if (g.spec.twoPass) {
            auto *p1 = g.pass1.get();
            for (int i = 0; i < g.spec.inputsReadPass1; ++i)
                p1->bindBuffer(
                    i, g.spec.inputRepresentatives.empty() || g.spec.inputRepresentatives[static_cast<size_t>(i)] == i
                           ? g.inputs[static_cast<size_t>(i)]
                           : nullptr);
            for (int s = 0; s < g.spec.statsCount; ++s)
                p1->bindBuffer(g.spec.inputsReadPass1 + s, g.stats[static_cast<size_t>(s)]);
        }
        auto *p2 = g.pass2.get();
        for (int i = 0; i < g.spec.inputCount; ++i)
            p2->bindBuffer(
                i, g.spec.inputRepresentatives.empty() || g.spec.inputRepresentatives[static_cast<size_t>(i)] == i
                       ? g.inputs[static_cast<size_t>(i)]
                       : nullptr);
        if (g.spec.scalesBinding >= 0 && g.qScales)
            p2->bindBuffer(g.spec.scalesBinding, g.qScales);
        const int outBinding = g.spec.outputBinding >= 0 ? g.spec.outputBinding
                                                         : g.spec.inputCount;
        if (g.spec.twoPass) {
            for (int s = 0; s < g.spec.statsCount; ++s)
                p2->bindBuffer(g.spec.inputCount + s, g.stats[static_cast<size_t>(s)]);
            p2->bindBuffer(g.spec.inputCount + g.spec.statsCount, g.output);
        } else {
            p2->bindBuffer(outBinding, g.output);
        }
    }

    /** Record all dispatches of one fused group into the shared sequence. */
    void recordGroup(const GroupRuntime &g, gpgpu::Sequence *seq) const {
        if (g.spec.twoPass)
            seq->recordDispatch(g.pass1.get(), g.spec.groupsX1, g.spec.groupsY1,
                                g.spec.groupsZ1);
        seq->recordDispatch(g.pass2.get(), g.spec.groupsX2, g.spec.groupsY2,
                            g.spec.groupsZ2);
    }
};

GpuProgram::GpuProgram() : impl_(new Impl()) {}
GpuProgram::~GpuProgram() { delete impl_; }

GpuProgram *GpuProgram::tryBuild(const Graph &graph, const OptimizedGraph &opt, int outputNode) {
    auto *gp = gpgpu::Gpgpu::create();
    if (!gp || !gp->isAvailable()) return nullptr;
    if (outputNode < 0 || outputNode >= graph.nodeCount()) return nullptr;

    auto *prog = new GpuProgram();
    prog->impl_->gpgpu = gp;
    auto &impl = *prog->impl_;

    try {
        // arena buffers: one per memory-plan slot
        impl.slotBuffer.resize(opt.slotSize.size(), nullptr);
        for (size_t s = 0; s < opt.slotSize.size(); ++s)
            impl.slotBuffer[s] = impl.alloc(opt.slotSize[s] * int(sizeof(float)));

        impl.sequence.reset(gp->newSequence());

        // placeholder / const uploads
        impl.placeholderBuffers.clear();
        impl.placeholderSizes.clear();
        for (int id = 0; id < graph.nodeCount(); ++id) {
            const auto &nd = graph.node(id);
            if (opt.nodeSlot[static_cast<size_t>(id)] < 0) continue;
            auto *buf = impl.slotBuffer[static_cast<size_t>(opt.nodeSlot[static_cast<size_t>(id)])];
            if (nd.type == OpType::Placeholder) {
                const int slot = nd.placeholderSlot;
                if (slot < 0) throw eve::Exception("GpuProgram: bad placeholder slot");
                if (int(impl.placeholderBuffers.size()) <= slot) {
                    impl.placeholderBuffers.resize(size_t(slot) + 1, nullptr);
                    impl.placeholderSizes.resize(size_t(slot) + 1, 0);
                }
                impl.placeholderBuffers[static_cast<size_t>(slot)] = buf;
                impl.placeholderSizes[static_cast<size_t>(slot)] = nd.size;
            } else if (nd.type == OpType::Const) {
                if (!nd.constBytes.empty()) {
                    // Packed weights may end mid-word; WebGPU transfers require four-byte sizes.
                    if (nd.constBytes.size() % 4 != 0) {
                        auto bytes = nd.constBytes;
                        bytes.resize((bytes.size() + 3) & ~size_t(3), 0);
                        buf->uploadBytes(bytes.data(), bytes.size());
                    } else {
                        buf->uploadBytes(nd.constBytes.data(), nd.constBytes.size());
                    }
                    if (!nd.constScales.empty()) {
                        auto *sb = impl.alloc(int(nd.constScales.size()) * int(sizeof(float)));
                        sb->uploadBytes(nd.constScales.data(),
                                        sizeof(float) * nd.constScales.size());
                        impl.qScalesByNode[id] = sb;
                    }
                } else {
                    buf->uploadBytes(nd.constData.data(), sizeof(float) * size_t(nd.size));
                }
            }
        }

        // per-group kernels (topological execution order)
        for (int gi : opt.groupOrder) {
            const auto &grp = opt.groups[static_cast<size_t>(gi)];
            if (grp.kind == GroupKind::Alias) continue;  // pure buffer alias

            Impl::GroupRuntime rt;
            for (int input : grp.inputs) {
                const int slot = opt.nodeSlot[static_cast<size_t>(input)];
                if (slot < 0) throw eve::Exception("GpuProgram: input without slot");
                rt.inputs.push_back(impl.slotBuffer[static_cast<size_t>(slot)]);
                const GraphNode &inN = graph.node(input);
                if (!inN.constBytes.empty() && !inN.constScales.empty()) {
                    auto it = impl.qScalesByNode.find(input);
                    if (it != impl.qScalesByNode.end()) rt.qScales = it->second;
                }
            }
            const int outSlot = opt.nodeSlot[static_cast<size_t>(grp.outputNode)];
            if (outSlot < 0) throw eve::Exception("GpuProgram: output without slot");
            rt.output = impl.slotBuffer[static_cast<size_t>(outSlot)];

            KernelSpec spec;
            if (grp.kind == GroupKind::MatMul) {
                const GraphNode &mm = graph.node(grp.nodes.front());
                auto             naiveResult = backendKernel(graph, grp);
                if (!naiveResult) throw eve::Exception("GpuProgram: matmul codegen failed");
                KernelSpec naive = std::move(naiveResult).value();
                if (mm.rank == 2) {
                    auto tiledResult = backendKernel(graph, grp, KernelVariant::TiledMatMul);
                    std::unique_ptr<gpgpu::ComputeShader> naiveShader, tiledShader;
                    if (tiledResult) {
                        KernelSpec tiled = std::move(tiledResult).value();
                        naiveShader.reset(gp->newShader(naive.pass2));
                        tiledShader.reset(gp->newShader(tiled.pass2));
                        try {
                            // Time with the real arena buffers bound: dispatching
                            // unbound kernels on the 4-byte dummy SSBO faults.
                            bindKernel(naiveShader.get(), naive, rt.inputs, rt.output);
                            const double tNaive =
                                timeDispatch(gp, naiveShader.get(), naive.groupsX2,
                                             naive.groupsY2);
                            bindKernel(tiledShader.get(), tiled, rt.inputs, rt.output);
                            const double tTiled =
                                timeDispatch(gp, tiledShader.get(), tiled.groupsX2,
                                             tiled.groupsY2);
                            if (tTiled < tNaive) {
                                spec = tiled;
                                rt.pass2 = std::move(tiledShader);
                            } else {
                                spec = naive;
                                rt.pass2 = std::move(naiveShader);
                            }
                        } catch (...) {
                            spec = naive;
                            rt.pass2 = std::move(naiveShader);
                        }
                    } else {
                        spec = naive;
                        rt.pass2.reset(gp->newShader(naive.pass2));
                    }
                } else {
                    spec = naive;
                    rt.pass2.reset(gp->newShader(naive.pass2));
                }
            } else {
                auto result = backendKernel(graph, grp);
                if (!result) throw eve::Exception("GpuProgram: kernel codegen failed");
                spec = std::move(result).value();
                rt.pass2.reset(gp->newShader(spec.pass2));
                if (spec.twoPass) rt.pass1.reset(gp->newShader(spec.pass1));
            }
            rt.spec = spec;

            for (int s = 0; s < spec.statsCount; ++s)
                rt.stats.push_back(impl.alloc(spec.statsSize * int(sizeof(float))));
            impl.bindGroup(rt);
            impl.groups.push_back(std::move(rt));
        }

        // final output
        const int outSlot = opt.nodeSlot[static_cast<size_t>(outputNode)];
        if (outSlot < 0) throw eve::Exception("GpuProgram: final output without slot");
        impl.outputBuffer = impl.slotBuffer[static_cast<size_t>(outSlot)];
        impl.outputSize = graph.node(outputNode).size;
        impl.outputStaging.reset(
            gp->newBuffer(impl.outputSize * int(sizeof(float)), "staging"));
        return prog;
    } catch (const std::exception &e) {
        fprintf(stderr, "[tensor] GpuProgram::tryBuild failed: %s\n", e.what());
        delete prog;
        return nullptr;
    } catch (...) {
        delete prog;
        return nullptr;
    }
}

std::vector<float> GpuProgram::run(const std::vector<const float *> &feeds) const {
    gpgpu::Sequence *seq = impl_->sequence.get();
    seq->begin();
    for (size_t i = 0; i < feeds.size(); ++i) {
        if (i >= impl_->placeholderBuffers.size() || !impl_->placeholderBuffers[i]) continue;
        seq->recordUpload(impl_->placeholderBuffers[i], feeds[i],
                          sizeof(float) * size_t(impl_->placeholderSizes[i]));
    }
    for (const auto &g : impl_->groups) impl_->recordGroup(g, seq);
    const uint64_t outBytes = sizeof(float) * size_t(impl_->outputSize);
    seq->recordDownload(impl_->outputBuffer, impl_->outputStaging.get(), outBytes);
    seq->submit();

    std::vector<float> out(static_cast<size_t>(impl_->outputSize));
    impl_->outputStaging->downloadBytes(out.data(), outBytes);
    return out;
}

bool gpuReduce(const float *data, int size, int op, float &outResult) {
    if (!data || size <= 0) return false;
    ReduceKernels *kernels = getReduceKernels();
    if (!kernels) return false;
    try {
        std::unique_ptr<gpgpu::GpuBuffer> in(
            kernels->gpgpu->newBuffer(size * int(sizeof(float)), "storage"));
        in->uploadBytes(data, sizeof(float) * size_t(size));

        const int groups = groupsFor(size);
        std::unique_ptr<gpgpu::GpuBuffer> partial(
            kernels->gpgpu->newBuffer(groups * int(sizeof(float)), "storage"));

        kernels->reduce->bindBuffer(0, in.get());
        kernels->reduce->bindBuffer(1, partial.get());
        kernels->reduce->setFloat(0, float(size));
        kernels->reduce->setFloat(1, float(op));
        kernels->gpgpu->dispatch(kernels->reduce, groups);

        std::vector<float> parts(static_cast<size_t>(groups));
        partial->downloadBytes(parts.data(), sizeof(float) * size_t(groups));

        float acc = op == 0 ? 0.f
                            : (op == 1 ? std::numeric_limits<float>::max()
                                       : -std::numeric_limits<float>::max());
        for (float v : parts) {
            if (op == 0) acc += v;
            else if (op == 1) acc = std::min(acc, v);
            else acc = std::max(acc, v);
        }
        outResult = acc;
        return true;
    } catch (...) {
        return false;
    }
}

}  // namespace eve::tensor