#pragma once

#include <string>
#include <vector>
#include <optional>
#include <cstdint>
#include <array>
#include <memory>
#include <utility>

#include <nd-wfc/wfc.h>

// Forward declarations
struct NonogramHints;
struct NonogramSolution;

// Compact storage for nonogram hints
// Stores each hint as a uint8_t (max hint value 255)
class NonogramHintsStorage {
public:
    // Each hint group is stored as: [count, hint1, hint2, ..., hintN]
    std::vector<uint8_t> data;
    std::vector<uint16_t> offsets; // Start position of each row/column in data

    void clear() {
        data.clear();
        offsets.clear();
    }

    // Add a new row/column's hints
    void addHints(const std::vector<uint8_t>& hints) {
        offsets.push_back(static_cast<uint16_t>(data.size()));
        data.push_back(static_cast<uint8_t>(hints.size())); // Store count first
        for (uint8_t hint : hints) {
            data.push_back(hint);
        }
    }

    // Get hints for a specific row/column
    std::vector<uint8_t> getHints(size_t index) const {
        if (index >= offsets.size()) return {};

        size_t start = offsets[index];
        if (start >= data.size()) return {};

        uint8_t count = data[start];
        std::vector<uint8_t> result(count);
        for (size_t i = 0; i < count; ++i) {
            result[i] = data[start + 1 + i];
        }
        return result;
    }

    size_t size() const { return offsets.size(); }
};

// Compact bit storage for nonogram solution
// Each cell is 1 bit (filled=1, empty=0)
class NonogramSolutionStorage {
public:
    std::vector<uint8_t> data; // 8 bits per byte
    uint16_t width = 0;
    uint16_t height = 0;

    void resize(uint16_t w, uint16_t h) {
        width = w;
        height = h;
        size_t total_bits = static_cast<size_t>(w) * h;
        size_t total_bytes = (total_bits + 7) / 8; // Ceiling division
        data.resize(total_bytes, 0);
    }

    void clear() {
        data.clear();
        width = 0;
        height = 0;
    }

    // Get cell value at (row, col) - true if filled, false if empty
    bool get(size_t row, size_t col) const {
        if (row >= height || col >= width) return false;
        size_t bit_index = row * width + col;
        size_t byte_index = bit_index / 8;
        size_t bit_offset = bit_index % 8;
        return (data[byte_index] & (1 << (7 - bit_offset))) != 0;
    }

    // Set cell value at (row, col)
    void set(size_t row, size_t col, bool filled) {
        if (row >= height || col >= width) return;
        size_t bit_index = row * width + col;
        size_t byte_index = bit_index / 8;
        size_t bit_offset = bit_index % 8;
        if (filled) {
            data[byte_index] |= (1 << (7 - bit_offset));
        } else {
            data[byte_index] &= ~(1 << (7 - bit_offset));
        }
    }

    // Load from goal string (sequence of 0s and 1s)
    bool loadFromGoalString(const std::string& goal, uint16_t w, uint16_t h) {
        size_t expected_size = static_cast<size_t>(w) * h;
        if (goal.length() != expected_size) return false;

        resize(w, h);
        for (size_t i = 0; i < expected_size; ++i) {
            if (goal[i] == '1') {
                size_t row = i / w;
                size_t col = i % w;
                set(row, col, true);
            }
        }
        return true;
    }
};

// Memory-efficient nonogram class
class Nonogram {
public:
    Nonogram();
    explicit Nonogram(uint16_t w, uint16_t h);

    // Dimensions
    uint16_t getWidth() const { return width_; }
    uint16_t getHeight() const { return height_; }

    // Hints access
    std::vector<uint8_t> getRowHints(size_t row) const {
        return row_hints_.getHints(row);
    }

    std::vector<uint8_t> getColumnHints(size_t col) const {
        return col_hints_.getHints(col);
    }

    size_t getRowCount() const { return row_hints_.size(); }
    size_t getColumnCount() const { return col_hints_.size(); }

    // Solution access (if available)
    bool hasSolution() const { return solution_ != nullptr; }
    bool getSolutionCell(size_t row, size_t col) const {
        return solution_ ? solution_->get(row, col) : false;
    }

    // Load from various formats
    bool loadFromFile(const std::string& filename);
    bool loadFromString(const std::string& content);

    void clear();

private:
    uint16_t width_ = 0;
    uint16_t height_ = 0;
    NonogramHintsStorage row_hints_;
    NonogramHintsStorage col_hints_;
    std::unique_ptr<NonogramSolutionStorage> solution_;

    // Parser helpers
    bool parseLine(const std::string& line, int& section);
    bool parseHintsLine(const std::string& line, bool is_rows);
    bool parseGoalLine(const std::string& line);
    bool parseDimensions(const std::string& line);
    std::vector<uint8_t> parseHintSequence(const std::string& hint_str);
};

// Fast nonogram loader
class NonogramLoader {
public:
    static std::optional<Nonogram> fromFile(const std::string& filename);
    static std::optional<Nonogram> fromString(const std::string& content);
    static std::vector<Nonogram> fromDirectory(const std::string& dirname);

private:
    static std::vector<std::string> splitContent(const std::string& content);
};

// Utility functions
std::string trim(const std::string& str);
bool startsWith(const std::string& str, const std::string& prefix);
std::vector<std::string> split(const std::string& str, char delimiter);
std::vector<uint8_t> parseNumberSequence(const std::string& str, char delimiter);

using NonogramWFC = WFC::Builder<Nonogram, bool>
    ::Define<false, true>
    ::Build;