export module ra3.fs;

import std;

export import ra3.core;

/**
 * Reading of the retail game's on-disk assets.
 *
 * Red Alert 3 ships its data in `BIG4` archives under `<install>\Data`, with
 * individual payloads compressed by EA's RefPack codec. This module implements
 * the container and codec so OpenRA3 can read a user's own installation.
 *
 * No game data is ever written into the repository; callers point the loader at
 * their local install (`--game-dir` / `RA3_GAME_DIR`, default `C:\Red Alert 3`).
 */
export namespace ra3::fs {
    using ra3::core::uint32;
    using ra3::core::uint8;
    using ra3::core::usize;

    /** Thrown when a `BIG4` archive is malformed or an entry is missing. */
    class archive_error : public std::runtime_error {
    public:
        using std::runtime_error::runtime_error;
    };

    /** Thrown when a RefPack stream is malformed. */
    class refpack_error : public std::runtime_error {
    public:
        using std::runtime_error::runtime_error;
    };

    inline constexpr std::array<uint8, 4> big_magic{'B', 'I', 'G', '4'};
    inline constexpr uint8 refpack_mask = 0x3EU;
    inline constexpr uint8 refpack_magic2 = 0xFBU;

    /** One file inside a `BIG4` archive. */
    struct big_entry {
        std::string name;
        uint32 offset = 0;
        uint32 size = 0;
    };

    /** True when `data` starts with a RefPack header (`0b??010000`, `0xFB`). */
    [[nodiscard]] inline auto is_refpack(std::span<const uint8> data) -> bool {
        return data.size() >= 2 && (data[0] & refpack_mask) == 0x10U && data[1] == refpack_magic2;
    }

    /**
     * Decompress an EA RefPack stream.
     *
     * @param data Compressed stream, starting at the header byte.
     * @return The decompressed bytes.
     * @throws refpack_error if the stream is malformed or the length disagrees.
     */
    [[nodiscard]] inline auto refpack_decompress(std::span<const uint8> data) -> std::vector<uint8> {
        if (!is_refpack(data)) throw refpack_error("not a RefPack stream");

        usize pos = 0;
        const auto header = data[pos++];
        const bool large_files = (header & 0x80U) != 0;
        const bool compressed_size_present = (header & 0x01U) != 0;
        pos++; // 0xFB

        const usize size_bytes = large_files ? 4U : 3U;
        auto read_size = [&]() -> uint32 {
            uint32 value = 0;
            for (usize i = 0; i < size_bytes; ++i) {
                if (pos >= data.size()) throw refpack_error("truncated RefPack size field");
                value = (value << 8) | data[pos++];
            }
            return value;
        };

        if (compressed_size_present) (void)read_size();
        const auto out_len = read_size();

        std::vector<uint8> out;
        out.reserve(out_len);

        auto copy_literals = [&](usize count) {
            if (pos + count > data.size()) throw refpack_error("truncated RefPack literals");
            out.insert(out.end(), data.begin() + static_cast<std::ptrdiff_t>(pos), data.begin() + static_cast<std::ptrdiff_t>(pos + count));
            pos += count;
        };

        auto copy_reference = [&](usize length, usize distance) {
            if (distance == 0 || distance > out.size()) throw refpack_error("RefPack back-reference out of range");
            usize start = out.size() - distance;
            for (usize i = 0; i < length; ++i) out.push_back(out[start + i]);
        };

        while (pos < data.size()) {
            const auto cmd = data[pos++];
            if ((cmd & 0x80U) == 0) { // 2-byte command
                if (pos >= data.size()) throw refpack_error("truncated 2-byte command");
                const auto b2 = data[pos++];
                copy_literals(cmd & 0x03U);
                copy_reference(((cmd & 0x1CU) >> 2) + 3, ((cmd & 0x60U) << 3) + b2 + 1);
            } else if ((cmd & 0x40U) == 0) { // 3-byte command
                if (pos + 1 >= data.size()) throw refpack_error("truncated 3-byte command");
                const auto b2 = data[pos];
                const auto b3 = data[pos + 1];
                pos += 2;
                copy_literals((b2 & 0xC0U) >> 6);
                copy_reference((cmd & 0x3FU) + 4, ((b2 & 0x3FU) << 8) + b3 + 1);
            } else if ((cmd & 0x20U) == 0) { // 4-byte command
                if (pos + 2 >= data.size()) throw refpack_error("truncated 4-byte command");
                const auto b2 = data[pos];
                const auto b3 = data[pos + 1];
                const auto b4 = data[pos + 2];
                pos += 3;
                copy_literals(cmd & 0x03U);
                copy_reference(((cmd & 0x0CU) << 6) + b4 + 5, ((cmd & 0x10U) << 12) + (b2 << 8) + b3 + 1);
            } else if (cmd < 0xFCU) { // long literal run
                copy_literals(((cmd & 0x1FU) + 1) << 2);
            } else { // stop
                copy_literals(cmd & 0x03U);
                break;
            }
        }

        if (out.size() != out_len) throw refpack_error("RefPack length mismatch");
        return out;
    }

    /** Decompress `data` when it is RefPack, otherwise copy it unchanged. */
    [[nodiscard]] inline auto maybe_decompress(std::span<const uint8> data) -> std::vector<uint8> {
        if (is_refpack(data)) return refpack_decompress(data);
        return {data.begin(), data.end()};
    }

    /**
     * Read the declared output size from a RefPack header without decompressing.
     *
     * @throws refpack_error if the stream is malformed.
     */
    [[nodiscard]] inline auto refpack_output_size(std::span<const uint8> data) -> uint32 {
        if (!is_refpack(data)) throw refpack_error("not a RefPack stream");
        usize pos = 0;
        const auto header = data[pos++];
        const bool large_files = (header & 0x80U) != 0;
        const bool compressed_size_present = (header & 0x01U) != 0;
        pos++; // 0xFB
        const usize size_bytes = large_files ? 4U : 3U;
        auto read_size = [&]() -> uint32 {
            uint32 value = 0;
            for (usize i = 0; i < size_bytes; ++i) {
                if (pos >= data.size()) throw refpack_error("truncated RefPack size field");
                value = (value << 8) | data[pos++];
            }
            return value;
        };
        if (compressed_size_present) (void)read_size();
        return read_size();
    }

    [[nodiscard]] constexpr auto read_be32(const uint8 *p) -> uint32 {
        return (static_cast<uint32>(p[0]) << 24U) | (static_cast<uint32>(p[1]) << 16U) | (static_cast<uint32>(p[2]) << 8U) | static_cast<uint32>(p[3]);
    }

    [[nodiscard]] constexpr auto read_le32(const uint8 *p) -> uint32 {
        return static_cast<uint32>(p[0]) | (static_cast<uint32>(p[1]) << 8U) | (static_cast<uint32>(p[2]) << 16U) | (static_cast<uint32>(p[3]) << 24U);
    }

    /**
     * A parsed `BIG4` archive. Only the index is held in memory; payloads are
     * read from disk on demand so multi-hundred-megabyte archives stay cheap.
     */
    class big_archive {
    public:
        /**
         * Parse the index of a `BIG4` archive.
         *
         * @param path Archive path.
         * @throws archive_error if the file is missing, not `BIG4`, or truncated.
         */
        [[nodiscard]] static auto open(const std::filesystem::path &path) -> big_archive {
            big_archive archive;
            archive.path_ = path;
            std::error_code ec;
            archive.file_size_ = static_cast<usize>(std::filesystem::file_size(path, ec));
            if (ec) throw archive_error("cannot stat archive: " + path.string());

            std::ifstream in(path, std::ios::binary);
            if (!in) throw archive_error("cannot open archive: " + path.string());

            std::array<uint8, 16> header{};
            in.read(reinterpret_cast<char *>(header.data()), static_cast<std::streamsize>(header.size()));
            if (!in || std::memcmp(header.data(), big_magic.data(), big_magic.size()) != 0) throw archive_error("not a BIG4 archive: " + path.string());

            const auto count = read_be32(header.data() + 8);

            // Read the variable-length index, growing the window until parsed.
            usize window = std::min(archive.file_size_, std::max<usize>(1U << 20U, static_cast<usize>(count) * 256U));
            std::vector<uint8> index;
            for (;;) {
                index.resize(window);
                in.clear();
                in.seekg(0);
                in.read(reinterpret_cast<char *>(index.data()), static_cast<std::streamsize>(window));
                const auto got = static_cast<usize>(in.gcount());
                index.resize(got);

                archive.entries_.clear();
                archive.entries_.reserve(count);
                usize pos = 16;
                bool complete = true;
                for (uint32 i = 0; i < count; ++i) {
                    if (pos + 8 > index.size()) {
                        complete = false;
                        break;
                    }
                    big_entry entry;
                    entry.offset = read_be32(index.data() + pos);
                    entry.size = read_be32(index.data() + pos + 4);
                    pos += 8;
                    const auto *begin = reinterpret_cast<const char *>(index.data() + pos);
                    const auto *end = reinterpret_cast<const char *>(std::memchr(begin, '\0', index.size() - pos));
                    if (end == nullptr) {
                        complete = false;
                        break;
                    }
                    entry.name.assign(begin, end);
                    pos += static_cast<usize>(end - begin) + 1U;
                    archive.entries_.push_back(std::move(entry));
                }
                if (complete) break;
                if (window >= archive.file_size_) throw archive_error("truncated BIG4 index: " + path.string());
                window = std::min(archive.file_size_, window * 2U);
            }

            archive.index_.reserve(archive.entries_.size());
            for (usize i = 0; i < archive.entries_.size(); ++i) archive.index_.emplace(archive.entries_[i].name, i);
            return archive;
        }

        [[nodiscard]] auto path() const -> const std::filesystem::path & { return path_; }
        [[nodiscard]] auto entries() const -> const std::vector<big_entry> & { return entries_; }
        [[nodiscard]] auto size() const -> usize { return entries_.size(); }
        [[nodiscard]] auto contains(std::string_view name) const -> bool { return index_.contains(std::string{name}); }

        /** Entry names whose path contains `needle`, in index order. */
        [[nodiscard]] auto find(std::string_view needle) const -> std::vector<const big_entry *> {
            std::vector<const big_entry *> matches;
            for (const auto &entry: entries_) {
                if (entry.name.find(needle) != std::string::npos) matches.push_back(&entry);
            }
            return matches;
        }

        /**
         * Read an entry's payload.
         *
         * @param name Entry name (backslash-separated, case-sensitive).
         * @param decompress RefPack-decompress the payload when true.
         * @throws archive_error if the entry is missing or unreadable.
         */
        [[nodiscard]] auto read(std::string_view name, bool decompress = true) const -> std::vector<uint8> {
            const auto it = index_.find(std::string{name});
            if (it == index_.end()) throw archive_error("no such entry: " + std::string{name});
            const auto &entry = entries_[it->second];

            std::ifstream in(path_, std::ios::binary);
            if (!in) throw archive_error("cannot open archive: " + path_.string());
            in.seekg(static_cast<std::streamoff>(entry.offset));
            std::vector<uint8> raw(entry.size);
            in.read(reinterpret_cast<char *>(raw.data()), static_cast<std::streamsize>(raw.size()));
            if (!in) throw archive_error("short read for entry: " + entry.name);
            return decompress ? maybe_decompress(raw) : raw;
        }

        /** Read the first `count` stored bytes of an entry (no decompression). */
        [[nodiscard]] auto read_prefix(std::string_view name, usize count) const -> std::vector<uint8> {
            const auto it = index_.find(std::string{name});
            if (it == index_.end()) throw archive_error("no such entry: " + std::string{name});
            const auto &entry = entries_[it->second];

            std::ifstream in(path_, std::ios::binary);
            if (!in) throw archive_error("cannot open archive: " + path_.string());
            in.seekg(static_cast<std::streamoff>(entry.offset));
            const auto want = std::min(count, static_cast<usize>(entry.size));
            std::vector<uint8> raw(want);
            in.read(reinterpret_cast<char *>(raw.data()), static_cast<std::streamsize>(raw.size()));
            raw.resize(static_cast<usize>(in.gcount()));
            return raw;
        }

    private:
        big_archive() = default;

        std::filesystem::path path_;
        usize file_size_ = 0;
        std::vector<big_entry> entries_;
        std::unordered_map<std::string, usize> index_;
    };

    /**
     * Locate a Red Alert 3 installation.
     *
     * Resolution order: the explicit argument, then `$RA3_GAME_DIR`, then the
     * default `C:\Red Alert 3`. A directory qualifies only if it has a `Data`
     * subdirectory.
     *
     * @return The install root, or `std::nullopt` when none is found.
     */
    [[nodiscard]] inline auto find_game_dir(const std::optional<std::filesystem::path> &explicit_dir = std::nullopt) -> std::optional<std::filesystem::path> {
        auto qualifies = [](const std::filesystem::path &candidate) {
            std::error_code ec;
            return !candidate.empty() && std::filesystem::is_directory(candidate / "Data", ec);
        };

        if (explicit_dir && qualifies(*explicit_dir)) return explicit_dir;
        if (const auto *env = std::getenv("RA3_GAME_DIR"); env != nullptr && *env != '\0') {
            const std::filesystem::path candidate{env};
            if (qualifies(candidate)) return candidate;
        }
        const std::filesystem::path default_dir{"C:/Red Alert 3"};
        if (qualifies(default_dir)) return default_dir;
        return std::nullopt;
    }
}