#include <utility>
#include <vector>
#include <algorithm>
#include <numeric>

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

using namespace aoc::types;

namespace dsk {
enum class format : u8 {
    file,
    free,
};

class block {
public:
    block(u32 id, dsk::format format) : m_id(id), m_format(format) {}

    auto id() const -> u32 { return m_id; }
    auto format() const -> dsk::format { return m_format; }
    auto str() const -> std::string {
        using namespace std::string_literals;
        std::string str{"{"};
        auto const id = m_id == aoc::lim<u32>::max() ? "na" : std::to_string(m_id);
        str += " id: "s + id + ",";
        str += " type: "s + std::to_string(u32(m_format)) + ",";
        str += " }";
        return str;
    }

private:
    u32 m_id;
    dsk::format m_format;
};

class page_block {
public:
    page_block(u32 id, dsk::format format, u32 size)
        : m_format(format), m_size(size) {
        for (u32 i = 0; i < m_size; ++i) {
            m_blocks.emplace_back(id, format);
        }
    }

    auto format() const -> dsk::format { return m_format; }
    auto size() const -> u32 { return m_size; }
    auto blocks() const -> std::vector<block> const& { return m_blocks; }

    auto str() const -> std::string {
        using namespace std::string_literals;
        std::string str{"page_block {"};
        str += " size: "s + std::to_string(u32(m_size)) + ",";
        str += " data: [ "s;
        for (usize i = 0; i < m_blocks.size(); ++i) {
            str += fmt::format("{}: ", i);
            str += m_blocks[i].str();
            if (i < m_blocks.size() - 1) str += ", ";
            else if (i == m_blocks.size() - 1) str += " ";
        }
        str += "]";
        str += " }";
        return str;
    }

private:
    dsk::format m_format;
    u32 m_size;
    std::vector<block> m_blocks{};
};

class disk_map {
public:
    disk_map(std::string const& input_str) {
        m_pages = input_str | std::views::filter([](auto const& c) { return c != '\n'; })
                            | std::views::transform(
                                 [i = 0u, id = 0u](auto const& size_char) mutable {
                                    auto const f = format(i++ % 2);
                                    auto curr_id = id;
                                    if (f == format::file)
                                        ++id;
                                    else
                                        curr_id = aoc::lim<u32>::max();
                                    return page_block(curr_id, f, u32(size_char - '0'));
                                 }
                           ) | std::ranges::to<std::vector>();
    }

    auto pages() const -> std::vector<page_block> {
        return m_pages;
    }

    auto to_blocks(std::vector<page_block> const& pages) const -> std::vector<block> {
        std::vector<block> tmp{};
        for (usize i = 0; i < pages.size(); ++i) {
            auto const& blocks = pages[i].blocks();
            for (usize j = 0; j < blocks.size(); ++j) {
                tmp.push_back(blocks[j]);
            }
        }
        return tmp;
    }

    auto defrag_block(std::vector<block> const& blocks) -> std::vector<block> {
        if (blocks.size() == 0) return {};

        std::vector<block> cpy = blocks;
        usize free_index = 0;
        usize file_index = cpy.size() - 1;
        while (free_index < cpy.size()) {
            free_index = next_free(cpy, free_index);
            if (free_index >= cpy.size() - 1) break;
            file_index = next_file(cpy, file_index);
            if (file_index == 0) break;

            if (free_index >= file_index) break;
            std::swap(cpy[free_index], cpy[file_index]);
        }
        return cpy;
    }

    // auto defrag_page(std::vector<page_block> const& pages) -> std::vector<page_block> {
    //     if (pages.size() == 0) return {};
    //
    //     std::vector<page_block> cpy = pages;
    //     usize free_index = 0;
    //     usize file_index = cpy.size() - 1;
    //     while (free_index < cpy.size()) {
    //         free_index = next_page_free(cpy, free_index);
    //         if (free_index >= cpy.size() - 1) break;
    //         auto const& tmp = cpy[free_index];
    //
    //         file_index = next_page_file(cpy, file_index, tmp.size());
    //         if (file_index == 0) break;
    //
    //         if (free_index >= file_index) break;
    //
    //         fmt::print("{}, {}\n", file_index, free_index);
    //         std::swap(cpy[free_index], cpy[file_index]);
    //     }
    //
    //     return cpy;
    // }

    auto checksum(std::vector<block> const& blocks) const -> u64 {
        u64 sum = 0;
        for (usize i = 0; i < blocks.size(); ++i) {
            auto const& block = blocks[i];
            if (block.format() == format::free) break;
            sum += i * block.id();
        }
        return sum;
    }

    auto debug_str(std::vector<block> const& blocks) const -> std::string {
        using namespace std::string_literals;
        std::string str{};

        for (usize i = 0; i < blocks.size(); ++i) {
            auto const& b = blocks[i];
            if (b.format() == format::file) {
                str += fmt::format("{}", b.id());
                if (b.id() > 9) str += "'";
            } else {
                str += ".";
            }
        }

        return str;
    }

    auto str() const -> std::string {
        using namespace std::string_literals;
        std::string str{"disk_map [\n"};
        auto const blocks = to_blocks(m_pages);

        for (usize i = 0; i < blocks.size(); ++i) {
            str += "  "s + blocks[i].str();
            if (i < blocks.size() - 1)
                str += ",\n";
        }
        str += "\n]";
        return str;
    }

private:
    auto next_free(std::vector<block> const& blocks, usize start) const -> usize {
        while (start < blocks.size()) {
            if (blocks[start].format() == format::free)
                break;
            ++start;
        }
        return start;
    }

    auto next_file(std::vector<block> const& blocks, usize start) const -> usize {
        while (start >= 0) {
            if (blocks[start].format() == format::file)
                break;
            --start;
        }
        return start;
    }

    auto next_page_free(std::vector<page_block> const& pages, usize start) const -> usize {
        while (start < pages.size()) {
            if (pages[start].format() == format::free)
                break;
            ++start;
        }
        return start;
    }

    auto next_page_file(std::vector<page_block> const& pages, usize start, u32 size) const -> usize {
        while (start >= 0) {
            if (pages[start].format() == format::file && pages[start].size() <= size)
                break;
            --start;
        }
        return start;
    }

private:
    std::vector<page_block> m_pages{};
};
}

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

    dsk::disk_map dsk{txt};
    fmt::print("{}\n", dsk.debug_str(dsk.to_blocks(dsk.pages())));
    auto const bl = dsk.defrag_block(dsk.to_blocks(dsk.pages()));
    fmt::print("{}\n", dsk.debug_str(bl));
    fmt::print("{}\n", dsk.checksum(bl));

    return {};
}