/**
 * Compiled `BinaryAsset` streams (`Data\*.bin` + `.manifest`).
 *
 * Retail RA3 does not ship its gameplay/art assets as source files: the
 * BinaryAssetBuilder compiles every XML/`.w3x` into a `.bin` *instance stream*
 * plus a `.manifest` index (`data\static.bin`, `data\global.bin`, ...). This
 * partition parses that index, exposes each asset's instance slice, its
 * relocation/import sidecars and its custom-data (`cdata`) blob, and implements
 * the Bob-Jenkins-style hash the builder names assets with.
 *
 * Asset names are `"Type:Instance"` (e.g. `W3DMesh:ABAIRFIELD`); `Type` is
 * case-sensitive and `Instance` is lowercased before hashing.
 */
export module ra3.assets:binary;

import std;

import :big;
import :bytes;
import :error;
import :refpack;

export namespace ra3::assets {

    /** The compiled instance streams a retail install ships. */
    enum class binary_stream {
        static_data, ///< `data\static.bin` - the main art/gameplay stream
        global,      ///< `data\global.bin` - sounds, AI, UI
        locale,      ///< `data\locale.bin` - localized on-demand textures
        static_low,  ///< `data\static_l.bin` - low-detail model LODs
        static_medium, ///< `data\static_m.bin` - medium-detail model LODs
    };

    /** The `data\<name>.manifest` stem for a stream. */
    [[nodiscard]] inline auto to_string(binary_stream stream) -> std::string_view {
        switch (stream) {
        case binary_stream::static_data: return "static";
        case binary_stream::global: return "global";
        case binary_stream::locale: return "locale";
        case binary_stream::static_low: return "static_l";
        case binary_stream::static_medium: return "static_m";
        }
        return "static";
    }

    /** Resolve a stream name (e.g. `"global"`), or `std::nullopt`. */
    [[nodiscard]] inline auto binary_stream_from_name(std::string_view name) -> std::optional<binary_stream> {
        for (const auto candidate: {binary_stream::static_data, binary_stream::global, binary_stream::locale, binary_stream::static_low, binary_stream::static_medium}) {
            if (to_string(candidate) == name) return candidate;
        }
        return std::nullopt;
    }

    /** Streams in the order `binary_stream_from_big` prefers when none is given. */
    inline constexpr std::array<binary_stream, 5> binary_stream_order{binary_stream::static_data, binary_stream::global, binary_stream::locale, binary_stream::static_low,
                                                                      binary_stream::static_medium};

    /** Types compiled from a `.w3x` model source. */
    inline constexpr std::array<std::string_view, 5> w3d_types{"W3DMesh", "W3DHierarchy", "W3DAnimation", "W3DContainer", "W3DCollisionBox"};

    /**
     * The hash BinaryAssetBuilder names assets with.
     *
     * A Bob-Jenkins `lookup3`-style mix. `seed` is the running value; the public
     * overloads seed it with the length (`hash_string`).
     */
    [[nodiscard]] inline auto fast_hash(std::span<const std::byte> data, std::uint32_t seed = 0) -> std::uint32_t {
        const auto length = data.size();
        if (length == 0U) return 0x1337C0DEU;
        const auto at16 = [&](std::size_t pos) -> std::uint32_t {
            return static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(data[pos])) |
                   (static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(data[pos + 1U])) << 8U);
        };
        std::uint32_t h = seed;
        std::size_t pos = 0;
        const auto extra = length & 3U;
        for (std::size_t i = 0; i < (length >> 2U); ++i) {
            h += at16(pos);
            h ^= (at16(pos + 2U) ^ (h << 5U)) << 11U;
            h += h >> 11U;
            pos += 4U;
        }
        if (extra == 1U) {
            h += std::to_integer<std::uint8_t>(data[pos]);
            h ^= h << 10U;
            h += h >> 1U;
        } else if (extra == 2U) {
            h += at16(pos);
            h ^= h << 11U;
            h += h >> 17U;
        } else if (extra == 3U) {
            h += at16(pos);
            h ^= h << 16U;
            h ^= static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(data[pos + 2U])) << 18U;
            h += h >> 11U;
        }
        h ^= h << 3U;
        h += h >> 5U;
        h ^= h << 2U;
        h += h >> 15U;
        h ^= h << 10U;
        return h;
    }

    /**
     * Hash an asset type (case-sensitive, the `TypeId`) or an instance name
     * (`case_sensitive = false`, the `InstanceId`).
     */
    [[nodiscard]] inline auto hash_string(std::string_view text, bool case_sensitive = true) -> std::uint32_t {
        std::string buffer{text};
        if (!case_sensitive) {
            std::ranges::transform(buffer, buffer.begin(), [](unsigned char ch) { return static_cast<char>(std::tolower(ch)); });
        }
        const auto *bytes = reinterpret_cast<const std::byte *>(buffer.data());
        return fast_hash(std::span<const std::byte>{bytes, buffer.size()}, static_cast<std::uint32_t>(buffer.size()));
    }

    /** One asset in a compiled stream. */
    struct binary_asset {
        std::uint32_t index = 0;
        std::uint32_t type_id = 0;
        std::uint32_t instance_id = 0;
        std::uint32_t type_hash = 0;
        std::uint32_t instance_hash = 0;
        std::vector<std::pair<std::uint32_t, std::uint32_t>> references; ///< (typeId, instanceId) pairs
        std::string name;   ///< "Type:Instance"
        std::string source; ///< the `.w3x`/XML it was compiled from
        std::uint32_t instance_size = 0;
        std::uint32_t relocation_size = 0;
        std::uint32_t imports_size = 0;
        bool tokenized = false;
        std::uint32_t instance_offset = 4;   ///< into the decompressed `.bin` (after its 4-byte checksum)
        std::uint32_t relocation_offset = 4; ///< into the `.relo`
        std::uint32_t imports_offset = 4;    ///< into the `.imp`

        [[nodiscard]] auto type_name() const -> std::string_view {
            const auto colon = this->name.find(':');
            return std::string_view{this->name}.substr(0, colon);
        }

        [[nodiscard]] auto instance_name() const -> std::string_view {
            const auto colon = this->name.find(':');
            return colon == std::string::npos ? std::string_view{this->name} : std::string_view{this->name}.substr(colon + 1U);
        }

        /**
         * The `cdata` blob name for a custom-data asset (`AudioFile`,
         * `OnDemandTexture`, ...):
         * `data\<stream>\cdata\<typeId>.<typeHash>.<instanceId>.<instanceHash>.cdata`.
         */
        [[nodiscard]] auto cdata_name(std::string_view stream) const -> std::string {
            return std::format("data\\{}\\cdata\\{:08x}.{:08x}.{:08x}.{:08x}.cdata", stream, this->type_id, this->type_hash, this->instance_id,
                               this->instance_hash);
        }
    };

    /** A callback that resolves a `cdata` blob name to its decompressed bytes. */
    using cdata_source = std::function<std::optional<std::vector<std::byte>>(const std::string &name)>;

    /**
     * A parsed `.manifest` (+ optional `.bin`/`.relo`/`.imp`).
     *
     * Only the manifest is mandatory; `need_data = false` when constructing
     * skips the (up to hundreds of megabytes) instance stream, which is all
     * listing the assets requires.
     */
    class binary_container {
    public:
        /** Parse a manifest, its instance stream and its optional fixup streams. */
        [[nodiscard]] static auto from_bytes(std::vector<std::byte> manifest, std::vector<std::byte> data = {}, std::vector<std::byte> relocation = {},
                                             std::vector<std::byte> imports = {}, std::string stream = "static") -> binary_container {
            binary_container container;
            container.stream_ = std::move(stream);
            container.manifest_ = std::move(manifest);
            container.data_ = std::move(data);
            container.relocation_ = std::move(relocation);
            container.imports_ = std::move(imports);
            container.parse();
            return container;
        }

        [[nodiscard]] auto stream() const -> const std::string & { return this->stream_; }
        [[nodiscard]] auto version() const -> std::uint16_t { return this->version_; }
        [[nodiscard]] auto is_linked() const -> bool { return this->is_linked_; }
        [[nodiscard]] auto assets() const -> const std::vector<binary_asset> & { return this->assets_; }
        [[nodiscard]] auto size() const -> std::size_t { return this->assets_.size(); }

        /** Count of assets per `Type` (the part before `:`). */
        [[nodiscard]] auto type_counts() const -> std::map<std::string, std::size_t> {
            std::map<std::string, std::size_t> counts;
            for (const auto &asset: this->assets_) ++counts[std::string{asset.type_name()}];
            return counts;
        }

        /** Map `(typeId << 32 | instanceId)` to the owning asset. */
        [[nodiscard]] auto asset_index() const -> const std::unordered_map<std::uint64_t, std::size_t> & { return this->asset_index_; }

        /** Every asset, in manifest order, whose `Type` matches (case-insensitive). */
        [[nodiscard]] auto of_type(std::string_view type) const -> std::vector<const binary_asset *> {
            std::vector<const binary_asset *> matches;
            for (const auto &asset: this->assets_) {
                if (std::ranges::equal(asset.type_name(), type, [](char a, char b) { return std::tolower(static_cast<unsigned char>(a)) == std::tolower(static_cast<unsigned char>(b)); }))
                    matches.push_back(&asset);
            }
            return matches;
        }

        /** Resolve `"Type:Instance"` (case-sensitive, then insensitive) or `"#index"`. */
        [[nodiscard]] auto find(std::string_view selector) const -> const binary_asset & {
            if (selector.starts_with('#')) {
                const auto index = static_cast<std::size_t>(std::stoul(std::string{selector.substr(1)}));
                if (index >= this->assets_.size()) throw binary_error("asset index out of range: " + std::string{selector});
                return this->assets_[index];
            }
            for (const auto &asset: this->assets_) {
                if (asset.name == selector) return asset;
            }
            for (const auto &asset: this->assets_) {
                if (std::ranges::equal(asset.name, selector, [](char a, char b) { return std::tolower(static_cast<unsigned char>(a)) == std::tolower(static_cast<unsigned char>(b)); }))
                    return asset;
            }
            throw binary_error("no such asset: " + std::string{selector});
        }

        auto set_cdata_source(cdata_source source) -> void { this->cdata_ = std::move(source); }

        [[nodiscard]] auto read_instance(const binary_asset &asset) const -> std::span<const std::byte> {
            if (asset.instance_size == 0U) return {};
            if (static_cast<std::uint64_t>(asset.instance_offset) + asset.instance_size > this->data_.size()) throw binary_error("instance extends past stream: " + asset.name);
            return std::span<const std::byte>{this->data_}.subspan(asset.instance_offset, asset.instance_size);
        }

        [[nodiscard]] auto read_relocation(const binary_asset &asset) const -> std::span<const std::byte> {
            if (asset.relocation_size == 0U) return {};
            if (static_cast<std::uint64_t>(asset.relocation_offset) + asset.relocation_size > this->relocation_.size()) throw binary_error("relocation extends past stream: " + asset.name);
            return std::span<const std::byte>{this->relocation_}.subspan(asset.relocation_offset, asset.relocation_size);
        }

        [[nodiscard]] auto read_imports(const binary_asset &asset) const -> std::span<const std::byte> {
            if (asset.imports_size == 0U) return {};
            if (static_cast<std::uint64_t>(asset.imports_offset) + asset.imports_size > this->imports_.size()) throw binary_error("imports extend past stream: " + asset.name);
            return std::span<const std::byte>{this->imports_}.subspan(asset.imports_offset, asset.imports_size);
        }

        /** The asset's `cdata` blob, or `std::nullopt` when it has none. */
        [[nodiscard]] auto read_cdata(const binary_asset &asset) const -> std::optional<std::vector<std::byte>> {
            if (!this->cdata_) return std::nullopt;
            return this->cdata_(asset.cdata_name(this->stream_));
        }

        /** The asset's real payload: its `cdata` blob when present, else its instance. */
        [[nodiscard]] auto read_payload(const binary_asset &asset) const -> std::vector<std::byte> {
            if (auto cdata = this->read_cdata(asset)) return *cdata;
            const auto instance = this->read_instance(asset);
            return {instance.begin(), instance.end()};
        }

    private:
        binary_container() = default;

        auto parse() -> void {
            byte_reader reader{this->manifest_};
            const auto is_big_endian = reader.read_u8();
            this->is_linked_ = reader.read_bool();
            this->version_ = reader.read_u16();
            if (is_big_endian != 0U) throw binary_error("big-endian manifests are not supported");
            if (this->version_ != 5U && this->version_ != 6U) throw binary_error("unsupported manifest version " + std::to_string(this->version_));
            (void) reader.read_u32(); // stream checksum
            (void) reader.read_u32(); // all-types hash
            const auto count = reader.read_u32();
            (void) reader.read_u32(); // total instance data size
            (void) reader.read_u32(); // max instance chunk size
            (void) reader.read_u32(); // max relocation chunk size
            (void) reader.read_u32(); // max imports chunk size
            const auto reference_buffer_size = reader.read_u32();
            const auto reference_name_buffer_size = reader.read_u32();
            const auto name_buffer_size = reader.read_u32();
            const auto source_buffer_size = reader.read_u32();

            constexpr std::size_t header_size = 48U;
            constexpr std::size_t entry_size = 48U;
            if (this->manifest_.size() < header_size + static_cast<std::size_t>(count) * entry_size) throw binary_error("manifest entry table is truncated");

            const auto entries_off = header_size;
            const auto references_off = entries_off + static_cast<std::size_t>(count) * entry_size;
            const auto reference_names_off = references_off + reference_buffer_size;
            const auto names_off = reference_names_off + reference_name_buffer_size;
            const auto sources_off = names_off + name_buffer_size;
            if (sources_off + source_buffer_size > this->manifest_.size()) throw binary_error("manifest string buffers are truncated");

            const auto cstr = [&](std::size_t base, std::int32_t offset) -> std::string {
                if (offset < 0) return {};
                const auto start = base + static_cast<std::size_t>(offset);
                if (start >= this->manifest_.size()) return {};
                std::size_t end = start;
                while (end < this->manifest_.size() && std::to_integer<std::uint8_t>(this->manifest_[end]) != 0U) ++end;
                return std::string{reinterpret_cast<const char *>(this->manifest_.data() + start), end - start};
            };

            this->assets_.reserve(count);
            std::uint32_t instance_offset = 4;
            std::uint32_t relocation_offset = 4;
            std::uint32_t imports_offset = 4;
            for (std::uint32_t i = 0; i < count; ++i) {
                byte_reader entry{std::span<const std::byte>{this->manifest_}.subspan(entries_off + static_cast<std::size_t>(i) * entry_size, entry_size)};
                binary_asset asset;
                asset.index = i;
                asset.type_id = entry.read_u32();
                asset.instance_id = entry.read_u32();
                asset.type_hash = entry.read_u32();
                asset.instance_hash = entry.read_u32();
                const auto reference_offset = entry.read_i32();
                const auto reference_count = entry.read_i32();
                const auto name_offset = entry.read_i32();
                const auto source_offset = entry.read_i32();
                asset.instance_size = entry.read_u32();
                asset.relocation_size = entry.read_u32();
                asset.imports_size = entry.read_u32();
                asset.tokenized = entry.read_u32() != 0U;
                asset.name = cstr(names_off, name_offset);
                asset.source = cstr(sources_off, source_offset);
                asset.instance_offset = instance_offset;
                asset.relocation_offset = relocation_offset;
                asset.imports_offset = imports_offset;
                instance_offset += asset.instance_size;
                relocation_offset += asset.relocation_size;
                imports_offset += asset.imports_size;

                for (std::int32_t j = 0; j < reference_count; ++j) {
                    byte_reader ref{std::span<const std::byte>{this->manifest_}.subspan(references_off + static_cast<std::size_t>(reference_offset) +
                                                                                          static_cast<std::size_t>(j) * 8U,
                                                                                      8U)};
                    asset.references.emplace_back(ref.read_u32(), ref.read_u32());
                }
                this->asset_index_.emplace((static_cast<std::uint64_t>(asset.type_id) << 32U) | asset.instance_id, this->assets_.size());
                this->assets_.push_back(std::move(asset));
            }
        }

        std::string stream_ = "static";
        std::vector<std::byte> manifest_;
        std::vector<std::byte> data_;
        std::vector<std::byte> relocation_;
        std::vector<std::byte> imports_;
        std::uint16_t version_ = 0;
        bool is_linked_ = false;
        std::vector<binary_asset> assets_;
        std::unordered_map<std::uint64_t, std::size_t> asset_index_;
        cdata_source cdata_;
    };

    namespace detail {
        [[nodiscard]] inline auto pick_manifest(std::string_view stream) -> std::string { return std::format("data\\{}.manifest", stream); }
    }

    /**
     * Load one BinaryAsset stream out of a `BIG4` archive.
     *
     * `stream` selects the `.manifest` (default: the first of
     * `binary_stream_order` found). When `need_data` is false only the manifest
     * is read, which is enough to enumerate assets without decompressing the
     * multi-hundred-megabyte instance stream.
     *
     * The returned container keeps a `cdata` reader that borrows `archive`; keep
     * the archive alive for as long as the container is used.
     */
    [[nodiscard]] inline auto binary_stream_from_big(const big_archive &archive, std::optional<binary_stream> stream = std::nullopt, bool need_data = true) -> binary_container {
        std::optional<binary_stream> chosen = stream;
        if (!chosen) {
            for (const auto candidate: binary_stream_order) {
                if (archive.contains(detail::pick_manifest(to_string(candidate)))) {
                    chosen = candidate;
                    break;
                }
            }
        }
        if (!chosen) throw binary_error("archive has no known .manifest entry");

        const auto manifest_name = detail::pick_manifest(to_string(*chosen));
        const auto base = manifest_name.substr(0, manifest_name.size() - std::string_view{".manifest"}.size());
        const auto manifest = maybe_decompress(archive.read(manifest_name, false));
        std::vector<std::byte> data;
        if (need_data) data = maybe_decompress(archive.read(base + ".bin", false));
        const auto relocation = maybe_decompress(archive.read(base + ".relo", false));
        const auto imports = maybe_decompress(archive.read(base + ".imp", false));

        auto container = binary_container::from_bytes(std::move(manifest), std::move(data), std::move(relocation), std::move(imports), std::string{to_string(*chosen)});
        container.set_cdata_source([&archive](const std::string &name) -> std::optional<std::vector<std::byte>> {
            const auto *entry = archive.find_exact(name);
            if (entry == nullptr) return std::nullopt;
            return maybe_decompress(archive.read(*entry, false));
        });
        return container;
    }

} // namespace ra3::assets