#include "slime_solver_core.hpp"

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <mutex>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

namespace sr = slime_research;

namespace {

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

struct Args {
    std::string command;
    int32_t x = 0;
    int32_t z = 0;
    int32_t size = 5;
    std::string pattern;
    uint64_t seed = 0;
    uint64_t lower35 = 0;
    int bits = 24;
    int threads = std::max(1U, std::thread::hardware_concurrency());
    uint64_t maxResults = 1000;
};

[[noreturn]] void usage(int code = 1) {
    std::cerr
        << "Usage:\n"
        << "  slime_pattern_cpp verify --seed N --x X --z Z [-n N|--pattern ROW/ROW]\n"
        << "  slime_pattern_cpp upper --lower35 N --x X --z Z [-n N|--pattern ...]\n"
        << "  slime_pattern_cpp roots --x X --z Z [--bits 18..35] [-n N|--pattern ...]\n"
        << "  slime_pattern_cpp search --x X --z Z [-p THREADS] [--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) {
    if (argc < 2) usage();
    Args args;
    args.command = argv[1];
    for (int i = 2; 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 == "--seed") args.seed = std::stoull(value(argc, argv, i));
        else if (arg == "--lower35") args.lower35 = std::stoull(value(argc, argv, i));
        else if (arg == "--bits") args.bits = std::stoi(value(argc, argv, i));
        else if (arg == "-p" || 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 || args.bits < 18 || args.bits > 35 || args.threads <= 0) usage();
    return args;
}

std::vector<CellOffset> parsePattern(const Args& args) {
    std::vector<CellOffset> offsets;
    if (args.pattern.empty()) {
        if (int64_t(args.size) * args.size > sr::kMaxPatternCells) {
            throw std::runtime_error("pattern exceeds the 64-cell prototype limit");
        }
        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});
            }
        }
        return offsets;
    }

    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");
    }
    return offsets;
}

std::vector<sr::CellTerm> makeTerms(const Args& args) {
    std::vector<sr::CellTerm> terms;
    for (const CellOffset& cell : parsePattern(args)) {
        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 VectorEmitter {
    std::vector<uint64_t>& values;
    void operator()(uint64_t seed) { values.push_back(seed); }
};

int commandVerify(const Args& args, const std::vector<sr::CellTerm>& terms) {
    const bool match = sr::verifyPattern(args.seed & sr::kMask48, terms.data(), int(terms.size()));
    std::cout << "seed48=" << (args.seed & sr::kMask48) << ' '
              << (match ? "match" : "no-match") << '\n';
    return match ? 0 : 1;
}

int commandUpper(const Args& args, const std::vector<sr::CellTerm>& terms) {
    std::vector<uint64_t> results;
    VectorEmitter emit{results};
    sr::solveUpper13(args.lower35, terms.data(), int(terms.size()), emit);
    std::sort(results.begin(), results.end());
    std::cout << "lower35=" << (args.lower35 & ((1ULL << 35) - 1ULL))
              << " exact_matches=" << results.size() << '\n';
    for (uint64_t seed : results) std::cout << seed << '\n';
    return 0;
}

int commandRoots(const Args& args, const std::vector<sr::CellTerm>& terms) {
    std::vector<uint32_t> roots = roots18(terms);
    std::cout << "bits=18 candidates=" << roots.size() << '\n';
    for (int bits = 19; bits <= args.bits; ++bits) {
        const uint32_t high = 1U << (bits - 1);
        const size_t oldSize = roots.size();
        roots.reserve(oldSize * 2);
        for (size_t i = 0; i < oldSize; ++i) roots.push_back(roots[i] | high);
        roots.erase(std::remove_if(roots.begin(), roots.end(), [&](uint32_t seed) {
            return !sr::partialPatternPossible(seed, bits, terms.data(), int(terms.size()));
        }), roots.end());
        std::cout << "bits=" << bits << " candidates=" << roots.size() << '\n';
    }
    return 0;
}

int commandSearch(const Args& args, const std::vector<sr::CellTerm>& terms) {
    const std::vector<uint32_t> roots = roots18(terms);
    constexpr uint64_t extensions = 1ULL << (35 - 18);
    const uint64_t workItems = uint64_t(roots.size()) * extensions;
    std::atomic<uint64_t> next{0};
    std::atomic<uint64_t> totalResults{0};
    std::mutex resultMutex;
    std::vector<uint64_t> results;
    const auto start = std::chrono::steady_clock::now();

    auto worker = [&] {
        std::vector<uint64_t> local;
        local.reserve(64);
        VectorEmitter emit{local};
        constexpr uint64_t batch = 256;
        for (;;) {
            const uint64_t begin = next.fetch_add(batch, std::memory_order_relaxed);
            if (begin >= workItems) break;
            const uint64_t end = std::min(begin + batch, workItems);
            for (uint64_t index = begin; index < end; ++index) {
                const uint64_t rootIndex = index >> (35 - 18);
                const uint64_t extension = index & (extensions - 1);
                const uint64_t lower35 = roots[rootIndex] | (extension << 18);
                sr::solveUpper13(lower35, terms.data(), int(terms.size()), emit);
            }
            if (!local.empty()) {
                totalResults.fetch_add(local.size(), std::memory_order_relaxed);
                std::lock_guard<std::mutex> lock(resultMutex);
                const size_t room = results.size() < args.maxResults
                    ? size_t(args.maxResults - results.size()) : 0;
                const size_t take = std::min(room, local.size());
                results.insert(results.end(), local.begin(), local.begin() + take);
                local.clear();
            }
        }
    };

    std::vector<std::thread> threads;
    for (int i = 0; i < args.threads; ++i) threads.emplace_back(worker);
    for (std::thread& thread : threads) thread.join();
    std::sort(results.begin(), results.end());
    const std::chrono::duration<double> elapsed = std::chrono::steady_clock::now() - start;
    std::cout << "lower18_roots=" << roots.size() << " lower35_checked=" << workItems
              << " exact_matches=" << totalResults.load() << " seconds=" << elapsed.count() << '\n';
    for (uint64_t seed : results) std::cout << seed << '\n';
    return 0;
}

}  // namespace

int main(int argc, char** argv) {
    try {
        const Args args = parseArgs(argc, argv);
        const std::vector<sr::CellTerm> terms = makeTerms(args);
        if (args.command == "verify") return commandVerify(args, terms);
        if (args.command == "upper") return commandUpper(args, terms);
        if (args.command == "roots") return commandRoots(args, terms);
        if (args.command == "search") return commandSearch(args, terms);
        usage();
    } catch (const std::exception& error) {
        std::cerr << "error: " << error.what() << '\n';
        return 2;
    }
}