#include <utility>
#include <vector>
#include <algorithm>
#include <numeric>
#include <thread>
#include <functional>
#include <mutex>

#include "aoc.hpp"
#include "aoc/utils.hpp"
#include "fmt/format.h"

using namespace aoc::types;
using namespace std::string_view_literals;

namespace ms {
class magic {
public:
    magic(std::string const& str) {
        m_stones = str | std::views::filter([](auto const& c) { return c != '\n'; })
                       | std::views::split(" "sv)
                       | std::views::transform([](auto const& c_str) {
                            return i64(std::stoll(std::string(std::begin(c_str), std::end(c_str))));
                       }) | std::ranges::to<std::vector>();
    }

    auto blink() -> void {
        std::vector<i64> stones{};
        for (usize i = 0; i < m_stones.size(); ++i) {
            auto const stone = m_stones[i];
            auto const str = std::to_string(stone);

            if (stone == 0) stones.emplace_back(1);
            else if (str.length() % 2 == 0) {
                auto const a = std::stoll(str.substr(0, str.length() / 2));
                auto const b = std::stoll(str.substr(str.length() / 2));
                stones.emplace_back(a);
                stones.emplace_back(b);
            } else {
                stones.emplace_back(stone * 2024);
            }
        }

        m_stones = stones;
    }

    auto blink(usize order, std::span<i64> const& stones) -> void {
        std::vector<i64> compute{};

        for (usize i = 0; i < stones.size(); ++i) {
            auto const stone = stones[i];
            auto const str = std::to_string(stone);

            if (stone == 0) compute.emplace_back(1);
            else if (str.length() % 2 == 0) {
                fmt::print("{}\n", str);
                auto const a = std::stoll(str.substr(0, str.length() / 2));
                auto const b = std::stoll(str.substr(str.length() / 2));
                compute.emplace_back(a);
                compute.emplace_back(b);
            } else {
                compute.emplace_back(stone * 2024);
            }
        }

        m_mutex.lock();
        m_tmp.push_back({order, compute});
        m_mutex.unlock();
    }

    auto size() const -> usize { return m_stones.size(); }

    auto join(i32 current) -> void {
        std::sort(std::begin(m_tmp), std::end(m_tmp), [](auto const& a, auto const& b) {
            return a.first < b.first;
        });

        m_stones = {};
        for (usize i = 0; i < m_tmp.size(); ++i) {
            for (usize j = 0; j < m_tmp[i].second.size(); ++j) {
                m_stones.push_back(m_tmp[i].second[j]);
            }
        }
        m_tmp = {};
        (void)current;

        // fmt::print("join... {}, {}\n", current, m_stones.size());
    }

    auto data() const -> std::vector<i64> const& { return m_stones; }
    auto str() const -> std::string {
        std::string str{};
        for (usize i = 0; i < m_stones.size(); ++i) {
            str += std::to_string(m_stones[i]);
            if (i < m_stones.size() - 1) str += " ";
        }
        return str;
    }

private:
    std::vector<i64> m_stones{};
    std::vector<std::pair<usize, std::vector<i64>>> m_tmp{};
    std::mutex m_mutex{};
};
}

auto aoc24::day11([[maybe_unused]]std::span<char const*> const& args) -> std::expected<void, aoc::error> {
    auto res = aoc::read_text("./dat/24/re/11.txt");
    // auto res = aoc::read_text("./dat/24/ex/11b.txt");
    if (!res) return std::unexpected(res.error());
    auto const txt = *res;

    ms::magic magic{txt};
    fmt::print("{}\n", magic.str());

    auto const max_threads = std::thread::hardware_concurrency();

    for (i32 i = 0; i < 75; ++i) {
        auto stones = magic.data();
        auto chunk_size = stones.size() / max_threads;
        if (chunk_size == 0) chunk_size = stones.size();

        std::vector<std::thread> threads{};
        for (usize k = 0; k < max_threads; ++k) {
            auto start_index = k * chunk_size;
            auto end_index = std::min((k + 1) * chunk_size, stones.size());
            if (start_index >= stones.size()) break;

            std::span<i64> sp{stones};
            std::span<i64> split = sp.subspan(start_index, end_index);
            threads.emplace_back([&]{ magic.blink(k, split); });
        }
        for (usize k = 0; k < threads.size(); ++k) {
            threads[k].join();
        }
        magic.join(i);
    }

    fmt::print("stones: {}\n", magic.size());

    return {};
}