#include "slime_solver_core.hpp"

#include <cuda_runtime.h>

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace sr = slime_research;

namespace {

#define CUDA_CHECK(call) do { \
    const cudaError_t error_ = (call); \
    if (error_ != cudaSuccess) { \
        std::cerr << #call << ": " << cudaGetErrorString(error_) << '\n'; \
        std::exit(1); \
    } \
} while (false)

struct Args {
    int32_t x = 0;
    int32_t z = 0;
    int32_t size = 5;
    std::string pattern;
    int threads = 128;
    uint64_t maxResults = 1000;
};

struct CellOffset {
    int32_t dx;
    int32_t dz;
    uint8_t requirement;
};

[[noreturn]] void usage(int code = 1) {
    std::cerr << "Usage: slime_pattern_cuda --x X --z Z "
                 "[-n N|--pattern ROW/ROW] [--threads N] [--max-results N]\n";
    std::exit(code);
}

const char* value(int argc, char** argv, int& i) {
    if (++i >= argc) usage();
    return argv[i];
}

Args parseArgs(int argc, char** argv) {
    Args args;
    for (int i = 1; i < argc; ++i) {
        const std::string arg = argv[i];
        if (arg == "--x") args.x = std::stoi(value(argc, argv, i));
        else if (arg == "--z") args.z = std::stoi(value(argc, argv, i));
        else if (arg == "-n" || arg == "--size") args.size = std::stoi(value(argc, argv, i));
        else if (arg == "--pattern") args.pattern = value(argc, argv, i);
        else if (arg == "--threads") args.threads = std::stoi(value(argc, argv, i));
        else if (arg == "--max-results") args.maxResults = std::stoull(value(argc, argv, i));
        else if (arg == "--help") usage(0);
        else usage();
    }
    if (args.size <= 0 || int64_t(args.size) * args.size > sr::kMaxPatternCells ||
        args.threads <= 0 || args.threads > 1024) usage();
    return args;
}

std::vector<sr::CellTerm> makeTerms(const Args& args) {
    std::vector<CellOffset> offsets;
    if (args.pattern.empty()) {
        for (int32_t dz = 0; dz < args.size; ++dz) {
            for (int32_t dx = 0; dx < args.size; ++dx) {
                offsets.push_back({dx, dz, sr::kSlime});
            }
        }
    } else {
        int32_t dx = 0;
        int32_t dz = 0;
        int32_t width = -1;
        for (size_t i = 0; i <= args.pattern.size(); ++i) {
            const char c = i == args.pattern.size() ? '/' : args.pattern[i];
            if (c == '/') {
                if (dx == 0 || (width >= 0 && dx != width)) {
                    throw std::runtime_error("pattern rows must have equal non-zero width");
                }
                width = dx;
                dx = 0;
                ++dz;
                continue;
            }
            if (c == '#') offsets.push_back({dx, dz, sr::kSlime});
            else if (c == '.') offsets.push_back({dx, dz, sr::kNonSlime});
            else if (c != '?') throw std::runtime_error("pattern accepts only #, ., ?, and /");
            ++dx;
        }
    }
    if (offsets.empty() || offsets.size() > sr::kMaxPatternCells) {
        throw std::runtime_error("pattern must contain 1..64 constrained cells");
    }

    std::vector<sr::CellTerm> terms;
    for (const CellOffset& cell : offsets) {
        terms.push_back({
            sr::coordinateTerm(args.x + cell.dx, args.z + cell.dz),
            cell.requirement,
        });
    }
    return terms;
}

std::vector<uint32_t> roots18(const std::vector<sr::CellTerm>& terms) {
    std::vector<uint32_t> roots;
    for (uint32_t seed = 0; seed < (1U << 18); ++seed) {
        if (sr::partialPatternPossible(seed, 18, terms.data(), int(terms.size()))) {
            roots.push_back(seed);
        }
    }
    return roots;
}

struct DeviceEmitter {
    unsigned long long* count;
    uint64_t* results;
    uint64_t capacity;

    __device__ void operator()(uint64_t seed) {
        const unsigned long long index = atomicAdd(count, 1ULL);
        if (index < capacity) results[index] = seed;
    }
};

__global__ void solveKernel(
    const uint32_t* roots,
    uint64_t workItems,
    const sr::CellTerm* cells,
    int cellCount,
    unsigned long long* resultCount,
    uint64_t* results,
    uint64_t resultCapacity
) {
    constexpr uint64_t extensions = 1ULL << (35 - 18);
    const uint64_t index = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x;
    if (index >= workItems) return;
    const uint64_t rootIndex = index >> (35 - 18);
    const uint64_t extension = index & (extensions - 1ULL);
    const uint64_t lower35 = roots[rootIndex] | (extension << 18);
    DeviceEmitter emit{resultCount, results, resultCapacity};
    sr::solveUpper13(lower35, cells, cellCount, emit);
}

}  // namespace

int main(int argc, char** argv) {
    const Args args = parseArgs(argc, argv);
    const std::vector<sr::CellTerm> terms = makeTerms(args);
    const std::vector<uint32_t> roots = roots18(terms);
    const uint64_t workItems = uint64_t(roots.size()) << (35 - 18);
    if (workItems == 0) {
        std::cout << "lower18_roots=0 lower35_checked=0 exact_matches=0 milliseconds=0\n";
        return 0;
    }

    uint32_t* deviceRoots = nullptr;
    sr::CellTerm* deviceTerms = nullptr;
    uint64_t* deviceResults = nullptr;
    unsigned long long* deviceCount = nullptr;
    CUDA_CHECK(cudaMalloc(&deviceRoots, roots.size() * sizeof(uint32_t)));
    CUDA_CHECK(cudaMalloc(&deviceTerms, terms.size() * sizeof(sr::CellTerm)));
    CUDA_CHECK(cudaMalloc(&deviceResults, args.maxResults * sizeof(uint64_t)));
    CUDA_CHECK(cudaMalloc(&deviceCount, sizeof(unsigned long long)));
    CUDA_CHECK(cudaMemcpy(deviceRoots, roots.data(), roots.size() * sizeof(uint32_t), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(deviceTerms, terms.data(), terms.size() * sizeof(sr::CellTerm), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemset(deviceCount, 0, sizeof(unsigned long long)));

    cudaEvent_t start;
    cudaEvent_t stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));
    CUDA_CHECK(cudaEventRecord(start));
    const uint64_t blocks64 = (workItems + uint64_t(args.threads) - 1) / uint64_t(args.threads);
    if (blocks64 > 0x7FFFFFFFULL) throw std::runtime_error("grid exceeds prototype limit");
    solveKernel<<<static_cast<unsigned>(blocks64), args.threads>>>(
        deviceRoots, workItems, deviceTerms, int(terms.size()), deviceCount,
        deviceResults, args.maxResults
    );
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));

    unsigned long long count = 0;
    CUDA_CHECK(cudaMemcpy(&count, deviceCount, sizeof(count), cudaMemcpyDeviceToHost));
    const size_t stored = size_t(std::min<unsigned long long>(count, args.maxResults));
    std::vector<uint64_t> results(stored);
    CUDA_CHECK(cudaMemcpy(results.data(), deviceResults, stored * sizeof(uint64_t), cudaMemcpyDeviceToHost));
    std::sort(results.begin(), results.end());
    float milliseconds = 0.0F;
    CUDA_CHECK(cudaEventElapsedTime(&milliseconds, start, stop));

    std::cout << "lower18_roots=" << roots.size() << " lower35_checked=" << workItems
              << " exact_matches=" << count << " milliseconds=" << milliseconds << '\n';
    for (uint64_t seed : results) std::cout << seed << '\n';

    cudaEventDestroy(start);
    cudaEventDestroy(stop);
    cudaFree(deviceRoots);
    cudaFree(deviceTerms);
    cudaFree(deviceResults);
    cudaFree(deviceCount);
}