#include "UnpackBuffers.h"
#include "IROperator.h"
#include "IRVisitor.h"

namespace Halide {
namespace Internal {

using std::map;
using std::pair;
using std::set;
using std::string;
using std::vector;

namespace {

struct BufferInfo {
    Expr handle;
    int dimensions;
};

class FindBufferSymbols : public IRVisitor {
    using IRVisitor::visit;

    void visit_param(const string &ref_name, const Parameter &param) {
        if (param.defined() && param.is_buffer()) {
            const string &name = param.name();
            buffers[name] =
                BufferInfo{Variable::make(type_of<halide_buffer_t *>(), name + ".buffer", param),
                           param.dimensions()};
        }
    }

    void visit_buffer(const string &ref_name, const Buffer<> &buffer) {
        if (buffer.defined()) {
            const string &name = buffer.name();
            buffers[name] =
                BufferInfo{Variable::make(type_of<halide_buffer_t *>(), name + ".buffer", buffer),
                           buffer.dimensions()};
        }
    }

    void visit(const Variable *op) override {
        visit_param(op->name, op->param);
        visit_buffer(op->name, op->image);
        symbols.insert(op->name);
    }

    void visit(const Load *op) override {
        visit_param(op->name, op->param);
        visit_buffer(op->name, op->image);
        symbols.insert(op->name);
        IRVisitor::visit(op);
    }

    void visit(const Store *op) override {
        visit_param(op->name, op->param);
        symbols.insert(op->name);
        IRVisitor::visit(op);
    }

public:
    set<string> symbols;
    map<string, BufferInfo> buffers;
};

}  // namespace

Stmt unpack_buffers(Stmt s) {
    FindBufferSymbols finder;
    s.accept(&finder);

    vector<pair<string, Expr>> lets;

    for (auto &p : finder.buffers) {
        const string &name = p.first;
        const BufferInfo &info = p.second;
        vector<Expr> args = {info.handle};

        string host_var = name;
        Expr host_val = Call::make(type_of<void *>(), Call::buffer_get_host, args, Call::Extern);
        lets.emplace_back(host_var, host_val);

        string dev_var = name + ".device";
        Expr dev_val = Call::make(type_of<uint64_t>(), Call::buffer_get_device, args, Call::Extern);
        lets.emplace_back(dev_var, dev_val);

        string dev_interface_var = name + ".device_interface";
        Expr dev_interface_val = Call::make(type_of<const halide_device_interface_t *>(),
                                            Call::buffer_get_device_interface, args, Call::Extern);
        lets.emplace_back(dev_interface_var, dev_interface_val);

        string type_code_var = name + ".type";
        Expr type_code_val = Call::make(UInt(32), Call::buffer_get_type, args, Call::Extern);
        lets.emplace_back(type_code_var, type_code_val);

        string host_dirty_var = name + ".host_dirty";
        Expr host_dirty_val = Call::make(Bool(), Call::buffer_get_host_dirty, args, Call::Extern);
        lets.emplace_back(host_dirty_var, host_dirty_val);

        string dev_dirty_var = name + ".device_dirty";
        Expr dev_dirty_val = Call::make(Bool(), Call::buffer_get_device_dirty, args, Call::Extern);
        lets.emplace_back(dev_dirty_var, dev_dirty_val);

        string dimensions_var = name + ".dimensions";
        Expr dimensions_val = Call::make(Int(32), Call::buffer_get_dimensions, args, Call::Extern);
        lets.emplace_back(dimensions_var, dimensions_val);

        for (int i = 0; i < info.dimensions; i++) {
            vector<Expr> args = {info.handle, i};

            string min_var = name + ".min." + std::to_string(i);
            Expr min_val = Call::make(Int(32), Call::buffer_get_min, args, Call::Extern);
            lets.emplace_back(min_var, min_val);

            string extent_var = name + ".extent." + std::to_string(i);
            Expr extent_val = Call::make(Int(32), Call::buffer_get_extent, args, Call::Extern);
            lets.emplace_back(extent_var, extent_val);

            string stride_var = name + ".stride." + std::to_string(i);
            Expr stride_val = Call::make(Int(32), Call::buffer_get_stride, args, Call::Extern);
            lets.emplace_back(stride_var, stride_val);
        }
    }

    while (!lets.empty()) {
        pair<string, Expr> l = lets.back();
        lets.pop_back();
        if (finder.symbols.count(l.first)) {
            s = LetStmt::make(l.first, l.second, s);
        }
    }

    // Create buffer is not null assertions
    for (auto &p : finder.buffers) {
        Expr buf = p.second.handle;
        Expr cond = reinterpret<uint64_t>(buf) != 0;
        Expr error = Call::make(Int(32), "halide_error_buffer_argument_is_null",
                                {p.first}, Call::Extern);
        Stmt check = AssertStmt::make(cond, error);
        s = Block::make(check, s);
    }

    return s;
}

}  // namespace Internal
}  // namespace Halide