pm64: support backgrounds with two palettes

2 files changed, 132 insertions, 59 deletions

diff --git a/src/factories/pm64/BackgroundFactory.cpp b/src/factories/pm64/BackgroundFactory.cpp
index c2eda63..5a825d8 100644
--- a/src/factories/pm64/BackgroundFactory.cpp
+++ b/src/factories/pm64/BackgroundFactory.cpp
@@ -3,102 +3,158 @@
 #include "utils/Decompressor.h"
 #include "spdlog/spdlog.h"
 
-// PM64 background file structure (N64 ROM format):
-// BackgroundHeader (0x10 bytes):
-//   0x00: rasterAddr (u32) - N64 VRAM address (e.g., 0x80200210) - NOT a file offset!
-//   0x04: paletteAddr (u32) - N64 VRAM address (e.g., 0x80200010) - NOT a file offset!
-//   0x08: startX (u16)
-//   0x0A: startY (u16)
-//   0x0C: width (u16)
-//   0x0E: height (u16)
+// Each header is 16 bytes:
+//   0x00: rasterAddr  (u32)
+//   0x04: paletteAddr (u32)
+//   0x08: startX      (u16)
+//   0x0A: startY      (u16)
+//   0x0C: width       (u16)
+//   0x0E: height      (u16)
 //
-// The decompressed data has a FIXED layout:
-//   0x0000: Header (16 bytes)
-//   0x0010: Palette (256 colors * 2 bytes = 512 bytes, RGBA16)
-//   0x0210: Raster (width * height bytes, CI8 indexed)
-//
-// The N64 VRAM addresses in the header are meaningless on PC - we use fixed offsets.
-
-static void ByteSwapBackgroundData(uint8_t* data, size_t size) {
-    if (size < 0x10) {
-        SPDLOG_WARN("Background data too small: {}", size);
-        return;
-    }
-
-    uint32_t* header32 = reinterpret_cast<uint32_t*>(data);
-    uint16_t* header16 = reinterpret_cast<uint16_t*>(data);
-
-    // Swap 16-bit dimension fields first (we need these for validation)
-    header16[4] = BSWAP16(header16[4]); // startX at offset 0x08
-    header16[5] = BSWAP16(header16[5]); // startY at offset 0x0A
-    header16[6] = BSWAP16(header16[6]); // width at offset 0x0C
-    header16[7] = BSWAP16(header16[7]); // height at offset 0x0E
+// Most backgrounds have 1 header + 1 palette + 1 raster (CI8).
+// Some (e.g. sbk_bg) have multiple headers sharing the same raster
+// but pointing to different palettes.
 
-    // The N64 header contains absolute VRAM addresses (0x802xxxxx) that cannot be
-    // converted to file offsets. The background data has a fixed layout:
-    //   - Palette at offset 0x10 (right after 16-byte header)
-    //   - Raster at offset 0x210 (after header + 512-byte palette)
-    // We ignore the N64 addresses and write the correct fixed offsets.
-    constexpr uint32_t paletteOffset = 0x10; // Right after 16-byte header
-    constexpr uint32_t rasterOffset = 0x210; // After header (16) + palette (512)
+static constexpr uint32_t VRAM_BASE = 0x80200000;
 
-    header32[0] = rasterOffset;
-    header32[1] = paletteOffset;
+static uint32_t ReadBE32(const uint8_t* p) {
+    return (p[0] << 24) | (p[1] << 16) | (p[2] << 8) | p[3];
+}
 
-    // Background palettes are swapped internally by libultraship, no need to swap them here
-    // Also, Raster data is CI8 (byte indices) - no swap needed
+static uint16_t ReadBE16(const uint8_t* p) {
+    return (p[0] << 8) | p[1];
 }
 
 std::optional<std::shared_ptr<IParsedData>> PM64BackgroundFactory::parse(std::vector<uint8_t>& buffer,
                                                                          YAML::Node& node) {
     auto offset = GetSafeNode<uint32_t>(node, "offset");
 
-    // Check if compressed (YAY0)
+    std::vector<uint8_t> bgData;
     auto compressionType = Decompressor::GetCompressionType(buffer, offset);
 
     if (compressionType == CompressionType::YAY0) {
         auto decoded = Decompressor::Decode(buffer, offset, CompressionType::YAY0);
         if (!decoded || decoded->size == 0) {
-            SPDLOG_ERROR("Failed to decompress YAY0 background data at offset 0x{:X}", offset);
+            SPDLOG_ERROR("Failed to decompress background at 0x{:X}", offset);
             return std::nullopt;
         }
-
-        std::vector<uint8_t> bgData(decoded->data, decoded->data + decoded->size);
-        ByteSwapBackgroundData(bgData.data(), bgData.size());
-
-        return std::make_shared<RawBuffer>(bgData);
+        bgData.assign(decoded->data, decoded->data + decoded->size);
     } else {
-        // Uncompressed - read raw data with size from YAML
         auto size = GetSafeNode<size_t>(node, "size");
         auto [_, segment] = Decompressor::AutoDecode(node, buffer, size);
+        bgData.assign(segment.data, segment.data + segment.size);
+    }
 
-        std::vector<uint8_t> bgData(segment.data, segment.data + segment.size);
-        ByteSwapBackgroundData(bgData.data(), bgData.size());
+    if (bgData.size() < 0x10) {
+        SPDLOG_ERROR("Background data too small at 0x{:X}: {} bytes", offset, bgData.size());
+        return std::nullopt;
+    }
+
+    uint8_t* data = bgData.data();
+    size_t dataSize = bgData.size();
+
+    uint32_t rasterOffset = ReadBE32(data + 0x00) - VRAM_BASE;
+    uint32_t palette0Offset = ReadBE32(data + 0x04) - VRAM_BASE;
+    uint16_t width = ReadBE16(data + 0x0C);
+    uint16_t height = ReadBE16(data + 0x0E);
+    uint32_t rasterSize = width * height;
 
-        return std::make_shared<RawBuffer>(bgData);
+    if (rasterOffset + rasterSize > dataSize) {
+        SPDLOG_ERROR("Background raster out of bounds: 0x{:X}+0x{:X} > 0x{:X}", rasterOffset, rasterSize, dataSize);
+        return std::nullopt;
     }
+    if (palette0Offset + 512 > dataSize) {
+        SPDLOG_ERROR("Background palette 0 out of bounds");
+        return std::nullopt;
+    }
+
+    std::vector<uint8_t> raster(data + rasterOffset, data + rasterOffset + rasterSize);
+
+    std::vector<std::vector<uint8_t>> palettes;
+    palettes.emplace_back(data + palette0Offset, data + palette0Offset + 512);
+
+    // Scan for additional headers (e.g. sbk_bg has 2 palettes sharing the same raster)
+    for (uint32_t hdrOff = 0x10; hdrOff + 0x10 <= rasterOffset; hdrOff += 0x10) {
+        uint32_t nextRasterAddr = ReadBE32(data + hdrOff);
+        uint32_t nextPalAddr = ReadBE32(data + hdrOff + 4);
+
+        if ((nextRasterAddr & 0xFFF00000) != 0x80200000 || (nextPalAddr & 0xFFF00000) != 0x80200000) {
+            break;
+        }
+
+        uint32_t nextPalOff = nextPalAddr - VRAM_BASE;
+        if (nextRasterAddr - VRAM_BASE != rasterOffset || nextPalOff == palette0Offset) {
+            break;
+        }
+        if (nextPalOff + 512 > dataSize) {
+            break;
+        }
+
+        palettes.emplace_back(data + nextPalOff, data + nextPalOff + 512);
+    }
+
+    SPDLOG_INFO("Background: {}x{} palettes={}", width, height, palettes.size());
+
+    return std::make_shared<PM64BackgroundData>(width, height, ReadBE16(data + 0x08), ReadBE16(data + 0x0A),
+                                                std::move(raster), std::move(palettes));
 }
 
 ExportResult PM64BackgroundBinaryExporter::Export(std::ostream& write, std::shared_ptr<IParsedData> raw,
                                                   std::string& entryName, YAML::Node& node, std::string* replacement) {
-    auto writer = LUS::BinaryWriter();
-    auto data = std::static_pointer_cast<RawBuffer>(raw)->mBuffer;
+    auto bg = std::static_pointer_cast<PM64BackgroundData>(raw);
 
-    // Write as Blob type - game loads as raw binary
-    WriteHeader(writer, Torch::ResourceType::Blob, 0);
-    writer.Write(static_cast<uint32_t>(data.size()));
-    writer.Write(reinterpret_cast<char*>(data.data()), data.size());
+    // Raster as Texture (CI8)
+    auto writer = LUS::BinaryWriter();
+    WriteHeader(writer, Torch::ResourceType::Texture, 0);
+    writer.Write((uint32_t)TextureType::Palette8bpp);
+    writer.Write((uint32_t)bg->width);
+    writer.Write((uint32_t)bg->height);
+    writer.Write((uint32_t)bg->raster.size());
+    writer.Write((char*)bg->raster.data(), bg->raster.size());
     writer.Finish(write);
 
+    // Each palette as a companion Texture (RGBA16, 256×1)
+    for (uint32_t i = 0; i < bg->palettes.size(); i++) {
+        auto palWriter = LUS::BinaryWriter();
+        WriteHeader(palWriter, Torch::ResourceType::Texture, 0);
+        palWriter.Write((uint32_t)TextureType::RGBA16bpp);
+        palWriter.Write((uint32_t)256);
+        palWriter.Write((uint32_t)1);
+        palWriter.Write((uint32_t)bg->palettes[i].size());
+        palWriter.Write((char*)bg->palettes[i].data(), bg->palettes[i].size());
+
+        std::stringstream ss;
+        palWriter.Finish(ss);
+        std::string str = ss.str();
+        std::vector<char> palData(str.begin(), str.end());
+
+        std::string palPath = entryName + "_pal" + std::to_string(i);
+        auto lastSlash = palPath.rfind('/');
+        if (lastSlash != std::string::npos) {
+            palPath = palPath.substr(lastSlash + 1);
+        }
+
+        Companion::Instance->RegisterCompanionFile(palPath, palData);
+    }
+
     return std::nullopt;
 }
 
 ExportResult PM64BackgroundHeaderExporter::Export(std::ostream& write, std::shared_ptr<IParsedData> raw,
                                                   std::string& entryName, YAML::Node& node, std::string* replacement) {
+    auto bg = std::static_pointer_cast<PM64BackgroundData>(raw);
     const auto symbol = GetSafeNode(node, "symbol", entryName);
 
     if (Companion::Instance->IsOTRMode()) {
-        write << "static const ALIGN_ASSET(2) char " << symbol << "[] = \"__OTR__" << (*replacement) << "\";\n\n";
+        write << "static const ALIGN_ASSET(2) char " << symbol << "[] = \"__OTR__" << (*replacement) << "\";\n";
+
+        for (uint32_t i = 0; i < bg->palettes.size(); i++) {
+            std::string palEntryName = (*replacement) + "_pal" + std::to_string(i);
+            write << "static const ALIGN_ASSET(2) char " << symbol << "_pal" << i << "[] = \"__OTR__" << palEntryName
+                  << "\";\n";
+        }
+
+        write << "\n";
         return std::nullopt;
     }
 
diff --git a/src/factories/pm64/BackgroundFactory.h b/src/factories/pm64/BackgroundFactory.h
index 099bafb..cbe4d80 100644
--- a/src/factories/pm64/BackgroundFactory.h
+++ b/src/factories/pm64/BackgroundFactory.h
@@ -1,7 +1,24 @@
 #pragma once
 
 #include "factories/BaseFactory.h"
-#include "types/RawBuffer.h"
+#include "utils/TextureUtils.h"
+
+// Parsed background data: raster (CI8) + palette(s) + metadata
+struct PM64BackgroundData : public IParsedData {
+    uint16_t width;
+    uint16_t height;
+    uint16_t startX;
+    uint16_t startY;
+    uint32_t palCount;
+    std::vector<uint8_t> raster;                // CI8 pixel indices
+    std::vector<std::vector<uint8_t>> palettes; // Each palette is 512 bytes (256 × RGBA16)
+
+    PM64BackgroundData(uint16_t w, uint16_t h, uint16_t sx, uint16_t sy, std::vector<uint8_t> raster,
+                       std::vector<std::vector<uint8_t>> palettes)
+        : width(w), height(h), startX(sx), startY(sy), palCount(palettes.size()), raster(std::move(raster)),
+          palettes(std::move(palettes)) {
+    }
+};
 
 class PM64BackgroundBinaryExporter : public BaseExporter {
     ExportResult Export(std::ostream& write, std::shared_ptr<IParsedData> data, std::string& entryName, YAML::Node& node, std::string* replacement) override;