#include <cstdint>
#include <poplar/DeviceManager.hpp>
#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/Program.hpp>
#include <algorithm>
#include <fstream>
#include <iostream>
#include <iterator>
#include <map>
#include <stdexcept>
#include <string>
#include <vector>

#ifndef LISP_CODELETS
#define LISP_CODELETS "lisp_codelets.gp"
#endif
constexpr unsigned bufferSize = 16384;
struct Result { std::vector<char> output; std::vector<int> status; std::vector<unsigned> gc; };

int main(int argc, char **argv) try {
    unsigned tiles = 4;
    bool test = false, stress = false, batch = false;
    std::string source;
    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        auto value = [&]() -> std::string {
            if (++i == argc) throw std::runtime_error("Missing value for " + a);
            return argv[i];
        };
        if (a == "--tiles") { auto v = value(); tiles = v == "all" ? 0 : std::stoul(v); }
        else if (a == "--self-test") test = true;
        else if (a == "--stress-gc") stress = true;
        else if (a == "--eval") { source = value(); batch = true; }
        else if (a == "--file") {
            std::ifstream f(value());
            if (!f) throw std::runtime_error("Cannot open Lisp file");
            source.assign(std::istreambuf_iterator<char>(f), {}); batch = true;
        } else if (a == "--help") {
            std::cout << "ipulisp [--tiles N|all] [--stress-gc] [--eval EXPR|--file FILE|--self-test]\n"
                         "Without a batch option, enter complete Lisp forms per line; :quit exits.\n";
            return 0;
        } else throw std::runtime_error("Unknown option: " + a);
    }
    auto manager = poplar::DeviceManager::createDeviceManager();
    auto devices = manager.getDevices(poplar::TargetType::IPU, 1);
    auto it = std::find_if(devices.begin(), devices.end(), [](poplar::Device &d) { return d.attach(); });
    if (it == devices.end()) throw std::runtime_error("No available hardware IPU");
    auto device = std::move(*it);
    unsigned available = device.getTarget().getNumTiles();
    if (!tiles) tiles = available;
    if (tiles > available) throw std::runtime_error("Requested more tiles than device contains");
    std::cerr << "Attached to hardware IPU " << device.getId() << "; " << tiles << " interpreters\n";
    poplar::Graph graph(device.getTarget());
    graph.addCodelets(LISP_CODELETS);
    auto input = graph.addVariable(poplar::CHAR, {bufferSize}, "source");
    auto output = graph.addVariable(poplar::CHAR, {tiles, bufferSize}, "output");
    auto status = graph.addVariable(poplar::INT, {tiles}, "status");
    auto gc = graph.addVariable(poplar::UNSIGNED_INT, {tiles}, "collections");
    auto control = graph.addVariable(poplar::UNSIGNED_INT, {3}, "control");
    graph.setTileMapping(input, 0);
    graph.setTileMapping(control, 0);
    auto cs = graph.addComputeSet("Lisp");
    for (unsigned t = 0; t < tiles; ++t) {
        graph.setTileMapping(output[t], t);
        graph.setTileMapping(status[t], t);
        graph.setTileMapping(gc[t], t);
        auto v = graph.addVertex(cs, "LispVertex", {{"source", input}, {"output", output[t]},
            {"length", control[0]}, {"reset", control[1]}, {"stress", control[2]},
            {"status", status[t]}, {"collections", gc[t]}});
        graph.setInitialValue(v["tile"], t);
        graph.setTileMapping(v, t);
    }
    graph.createHostWrite("source", input);
    graph.createHostWrite("control", control);
    graph.createHostRead("output", output);
    graph.createHostRead("status", status);
    graph.createHostRead("gc", gc);
    poplar::Engine engine(graph, poplar::program::Execute(cs));
    engine.load(device);
    Result result{std::vector<char>(tiles * bufferSize), std::vector<int>(tiles), std::vector<unsigned>(tiles)};
    auto run = [&](const std::string &code, bool reset, bool forceGC) -> const Result & {
        if (code.size() > bufferSize) throw std::runtime_error("Source exceeds 16384 bytes");
        if (code.find('\0') != std::string::npos) throw std::runtime_error("Source contains a NUL byte");
        std::vector<char> data(bufferSize);
        std::copy(code.begin(), code.end(), data.begin());
        unsigned flags[] = {unsigned(code.size()), unsigned(reset), unsigned(forceGC)};
        engine.writeTensor("source", data.data(), data.data() + data.size());
        engine.writeTensor("control", flags, flags + 3);
        engine.run();
        engine.readTensor("output", result.output.data(), result.output.data() + result.output.size());
        engine.readTensor("status", result.status.data(), result.status.data() + tiles);
        engine.readTensor("gc", result.gc.data(), result.gc.data() + tiles);
        return result;
    };
    auto check = [&](const std::string &name, const std::string &code, const std::string &expected,
                     bool reset, bool forceGC, int expectedStatus = 0) {
        run(code, reset, forceGC);
        for (unsigned t = 0; t < tiles; ++t) {
            std::string got(result.output.data() + t * bufferSize);
            auto want = expected == "@tile" ? std::to_string(t) + "\n" : expected;
            if (result.status[t] != expectedStatus || (expectedStatus == 0 && got != want))
                throw std::runtime_error(name + " failed on tile " + std::to_string(t) + ": " + got);
        }
        std::cout << "PASS " << name << " (" << tiles << " tiles" << (forceGC ? ", forced GC" : "") << ")\n";
    };
    if (test) {
        for (bool force : {false, true}) {
            check("arithmetic", "(+ 1 2 3) (- 9 4) (< 2 3)", "6\n5\nt\n", true, force);
            check("tile identity", "__tileid", "@tile", false, force);
            check("persistent define", "(define counter ((lambda (n) (lambda () (setq n (+ n 1)))) 40))", "<function>\n", false, force);
            check("persistent closure", "(counter) (counter)", "41\n42\n", false, force);
            check("lists", "(list 1 (+ 1 1) 3) (cdr '(a b . c))", "(1 2 3)\n(b . c)\n", false, force);
            check("macro", "(defmacro unless (c x) (list 'if c () x)) (unless () 42)", "<macro>\n42\n", false, force);
            check("recursion", "(defun fib (n) (if (< n 2) n (+ (fib (- n 1)) (fib (- n 2))))) (fib 10)", "<function>\n55\n", false, force);
            check("collector", "(define i 0) (while (< i 3000) (list i i i) (setq i (+ i 1))) i", "0\n()\n3000\n", false, force);
            if (*std::min_element(result.gc.begin(), result.gc.end()) == 0)
                throw std::runtime_error("GC test did not collect on every tile");
            std::cout << "Collections per tile: " << result.gc.front() << "\n";
            check("undefined symbol", "missing-symbol", "", false, force, 1);
            check("error recovery", "(counter)", "43\n", false, force);
            check("parse error", "'", "", false, force, 1);
            check("parse recovery", "(+ 20 22)", "42\n", false, force);
            check("output overflow", "(define i 0) (while (< i 9000) (println i) (setq i (+ i 1)))", "", false, force, 2);
            check("overflow recovery", "(+ 20 22)", "42\n", false, force);
        }
        std::cout << "All hardware tests passed.\n";
        return 0;
    }
    bool first = true;
    do {
        if (!batch) {
            std::cout << "lisp> " << std::flush;
            if (!std::getline(std::cin, source) || source == ":quit") break;
        }
        run(source, first, stress); first = false;
        std::map<std::string, unsigned> counts;
        for (unsigned t = 0; t < tiles; ++t) ++counts[result.output.data() + t * bufferSize];
        for (const auto &entry : counts) std::cout << "[" << entry.second << "x] " << entry.first;
        bool failed = false;
        for (unsigned t = 0; t < tiles; ++t) if (result.status[t]) {
            std::cerr << "Tile " << t << " status " << result.status[t] << "\n"; failed = true;
        }
        if (batch) return failed ? 1 : 0;
    } while (!batch);
    return 0;
} catch (const std::exception &e) {
    std::cerr << "Error: " << e.what() << "\n";
    return 1;
}