#include "Halide.h"
#include <map>
#include <stdio.h>
#include <string>

namespace {

using std::map;
using std::string;
using std::vector;
using namespace Halide;
using namespace Halide::Internal;

class FindErrorHandler : public IRVisitor {
public:
    bool result;
    FindErrorHandler()
        : result(false) {
    }
    using IRVisitor::visit;
    void visit(const Call *op) override {
        if (op->name == "halide_error_unaligned_host_ptr" &&
            op->call_type == Call::Extern) {
            result = true;
            return;
        }
        IRVisitor::visit(op);
    }
};

class ParseCondition : public IRVisitor {
public:
    Expr condition;

    using IRVisitor::visit;
    void visit(const Mod *op) override {
        condition = op;
    }

    void visit(const Call *op) override {
        if (op->is_intrinsic(Call::bitwise_and)) {
            condition = op;
        } else {
            IRVisitor::visit(op);
        }
    }
};

class CountHostAlignmentAsserts : public IRVisitor {
public:
    int count;
    std::map<string, int> alignments_needed;
    CountHostAlignmentAsserts(std::map<string, int> m)
        : count(0),
          alignments_needed(m) {
    }

    using IRVisitor::visit;

    void visit(const AssertStmt *op) override {
        Expr m = op->message;
        FindErrorHandler f;
        m.accept(&f);
        if (f.result) {
            Expr c = op->condition;
            ParseCondition p;
            c.accept(&p);
            if (p.condition.defined()) {
                Expr left, right;
                if (const Mod *mod = p.condition.as<Mod>()) {
                    left = mod->a;
                    right = mod->b;
                } else if (const Call *call = Call::as_intrinsic(p.condition, {Call::bitwise_and})) {
                    left = call->args[0];
                    right = call->args[1];
                }
                const Reinterpret *reinterpret = left.as<Reinterpret>();
                if (!reinterpret) return;
                Expr name = reinterpret->value;
                const Variable *V = name.as<Variable>();
                string name_host_ptr = V->name;
                int expected_alignment = alignments_needed[name_host_ptr];
                if (is_const(right, expected_alignment) || is_const(right, expected_alignment - 1)) {
                    count++;
                    alignments_needed.erase(name_host_ptr);
                }
            }
        }
    }
};

void set_alignment_host_ptr(ImageParam &i, int align, std::map<string, int> &m) {
    i.set_host_alignment(align);
    m.insert(std::pair<string, int>(i.name(), align));
}

int count_host_alignment_asserts(Func f, std::map<string, int> m) {
    Target t = get_jit_target_from_environment();
    t.set_feature(Target::NoBoundsQuery);
    f.compute_root();
    Stmt s = Internal::lower_main_stmt({f.function()}, f.name(), t);
    CountHostAlignmentAsserts c(m);
    s.accept(&c);
    return c.count;
}

int test() {
    Var x, y, c;
    std::map<string, int> m;
    ImageParam i1(Int(8), 1);
    ImageParam i2(Int(8), 1);
    ImageParam i3(Int(8), 1);

    set_alignment_host_ptr(i1, 128, m);
    set_alignment_host_ptr(i2, 32, m);

    Func f("f");
    f(x) = i1(x) + i2(x) + i3(x);
    f.output_buffer().set_host_alignment(128);
    m.insert(std::pair<string, int>("f", 128));
    int cnt = count_host_alignment_asserts(f, m);
    if (cnt != 3) {
        printf("Error: expected 3 host alignment assertions in code, but got %d\n", cnt);
        return 1;
    }

    printf("Success!\n");
    return 0;
}

}  // namespace

int main(int argc, char **argv) {
    return test();
}