path: root/Scripts/DataDecoder.cs

using UdonSharp;
using UnityEngine;
using VRC.SDKBase;
using VRC.Udon;
using VRC.SDK3.Rendering;
using VRC.Udon.Common.Interfaces;

public class DataDecoder : UdonSharpBehaviour
{
  public RenderTexture sourceTexture;
  public int tileSize = 8;
  private const int MaxTiles = 8192;

  private bool isValid;
  private Color[] pixelDataFloat;
  private Color32[] pixelDataBytes;
  private bool hasData = false;
  private bool hasReportedTileLimit;

  private const int PixelFormatNone = 0;
  private const int PixelFormatColor32 = 1;
  private const int PixelFormatFloat = 2;

  private int currentPixelFormat = PixelFormatNone;

  void Start() {}

  private bool IsValid()
  {
    if (sourceTexture == null)  {
      Debug.LogWarning("DataDecoder: sourceTexture missing");
      return false;
    }
    if (tileSize <= 0)  {
      Debug.LogWarning("DataDecoder: invalid tile size");
      return false;
    }
    return true;
  }

  void Update()
  {
    // Only call IsValid() until it succeeds, then cache.
    if (!isValid)
    {
      isValid = IsValid();
      if (!isValid) {
        return;
      }
    }

    int pixelCount = sourceTexture.width * sourceTexture.height;
    if (pixelDataBytes == null || pixelCount != pixelDataBytes.Length)
    {
      Debug.Log($"DataDecoder: allocating Color32 buffer for {pixelCount} pixels");
      pixelDataBytes = new Color32[pixelCount];
    }

    if (pixelDataFloat == null || pixelCount != pixelDataFloat.Length)
    {
      pixelDataFloat = new Color[pixelCount];
    }

    // Request readback every frame (multiple requests can be in flight)
    VRCAsyncGPUReadback.Request(sourceTexture, 0, (IUdonEventReceiver)this);

    // Process data if available
    if (hasData)
    {
      DecodeData(tileSize);
      hasData = false;
    }
  }

  public override void OnAsyncGpuReadbackComplete(VRCAsyncGPUReadbackRequest request)
  {
    if (request.hasError)
    {
      Debug.LogError("DataDecoder: GPU readback error");
      return;
    }

    // Store the pixel data (overwrites previous data with latest)
    if (pixelDataBytes != null && request.TryGetData(pixelDataBytes))
    {
      currentPixelFormat = PixelFormatColor32;
      hasData = true;
      return;
    }

    if (pixelDataFloat != null && request.TryGetData(pixelDataFloat))
    {
      currentPixelFormat = PixelFormatFloat;
      hasData = true;
      return;
    }

    Debug.LogWarning("DataDecoder: Unable to read GPU data into any supported buffer format");
  }

  /// <summary>
  /// Prints the RGB values of the first 4 tiles for debugging purposes.
  /// </summary>
  private void PrintFirst4Tiles(int tileSize, int tilesPerColumn)
  {
    int tilesPerRow = (sourceTexture.width + tileSize - 1) / tileSize;
    int totalTiles = Mathf.Max(Mathf.Min(tilesPerColumn * tilesPerRow, MaxTiles), 0);

    if (totalTiles == 0)
    {
      Debug.Log("DataDecoder: tiles=<none>");
      return;
    }

    int tilesToLog = Mathf.Min(4, totalTiles);
    string log = "DataDecoder: tiles=";

    for (int tileIdx = 0; tileIdx < tilesToLog; tileIdx++)
    {
      GetTileRGB(tileIdx, tileSize, tilesPerColumn, out byte r, out byte g, out byte b);
      log += $"[{tileIdx}:{r},{g},{b}]";
      if (tileIdx < tilesToLog - 1)
      {
        log += " ";
      }
    }

    if (totalTiles > tilesToLog)
    {
      log += " ...";
    }

    Debug.Log(log);
  }

  /// <summary>
  /// Decodes data from a texture where the data is encoded in RGB channels of tiles.
  /// The first tile contains the length (number of bytes to decode).
  /// Subsequent tiles contain the actual data.
  /// </summary>
  /// <param name="tileSize">The size of each tile (e.g., 8 for 8x8 tiles)</param>
  /// <returns>An array of decoded bytes</returns>
  public byte[] DecodeData(int tileSize)
  {
    if (!HasPixelData())
    {
      Debug.LogWarning("DataDecoder: No pixel data available");
      return new byte[0];
    }

    int tilesPerColumn = (sourceTexture.height + tileSize - 1) / tileSize;
    int tilesPerRow = (sourceTexture.width + tileSize - 1) / tileSize;

    PrintFirst4Tiles(tileSize, tilesPerColumn);

    // Read the length from the first tile
    int length = ReadLength(tileSize, tilesPerColumn);
    //Debug.Log($"Parsed length: {length}");
    if (length <= 0) {
      Debug.LogWarning($"DataDecoder: Length is 0 or negative: {length}");
      return new byte[0];
    }

    // Calculate how many data tiles we need (excluding the length tile)
    int bytesPerTile = 3; // RGB from center pixel only
    int tilesNeeded = (length + bytesPerTile - 1) / bytesPerTile; // Ceiling division

    // Allocate result array
    byte[] data = new byte[length];

    // Calculate tiles per column
    // tilesPerColumn/Row calculated above

    if (tilesPerColumn <= 0 || tilesPerRow <= 0)
    {
      Debug.LogError("DataDecoder: Texture dimensions smaller than tile size");
      return new byte[0];
    }

    int totalTilesRaw = tilesPerColumn * tilesPerRow;
    int totalTiles = Mathf.Min(totalTilesRaw, MaxTiles);

    if (totalTilesRaw <= MaxTiles)
    {
      hasReportedTileLimit = false;
    }
    else if (!hasReportedTileLimit)
    {
      Debug.LogWarning($"DataDecoder: Limiting decoding to first {MaxTiles} tiles (of {totalTilesRaw})");
      hasReportedTileLimit = true;
    }

    // Read data from each tile (starting after the length tile)
    int byteIndex = 0;
    for (int tileIdx = 0; tileIdx < tilesNeeded && byteIndex < length; tileIdx++)
    {
      // Skip tile 0 (the length tile) by using actualTileIdx
      int actualTileIdx = tileIdx + 1;

      if (actualTileIdx >= totalTiles)
      {
        Debug.LogWarning($"DataDecoder: Ran out of tiles at index {actualTileIdx} while decoding {length} bytes");
        break;
      }

      // Get RGB bytes from this tile
      GetTileRGB(actualTileIdx, tileSize, tilesPerColumn, out byte r, out byte g, out byte b);

      // Write bytes to output array
      if (byteIndex < length) data[byteIndex++] = r;
      if (byteIndex < length) data[byteIndex++] = g;
      if (byteIndex < length) data[byteIndex++] = b;
    }

    return data;
  }

  /// <summary>
  /// Gets the RGB bytes from the center pixel of the nth tile.
  /// Tiles are indexed in column-major order (tile 0 is top-left).
  /// </summary>
  /// <param name="tileIndex">The index of the tile (0-based)</param>
  /// <param name="tileSize">The size of each tile</param>
  /// <param name="r">Output: Red channel as byte (0-255)</param>
  /// <param name="g">Output: Green channel as byte (0-255)</param>
  /// <param name="b">Output: Blue channel as byte (0-255)</param>
  /// <returns>The center pixel color, or Color.clear if invalid data detected</returns>
  private Color GetTileRGB(int tileIndex, int tileSize, int tilesPerColumn, out byte r, out byte g, out byte b)
  {
    // Calculate tile position (column-major order)
    int columnIdx = tileIndex / tilesPerColumn;
    int rowIdx = tileIndex % tilesPerColumn;

    int tileX = columnIdx * tileSize;
    int tileY = rowIdx * tileSize;

    int xStart = tileX;
    int xEnd = Mathf.Min(tileX + tileSize, sourceTexture.width);
    int yStart = tileY;
    int yEnd = Mathf.Min(tileY + tileSize, sourceTexture.height);

    if (xStart >= xEnd || yStart >= yEnd)
    {
      r = 0;
      g = 0;
      b = 0;
      return Color.clear;
    }

    // Sample roughly from the tile center, clamped to the valid area.
    int centerX = Mathf.Clamp(tileX + tileSize / 2, xStart, xEnd - 1);
    int centerY = Mathf.Clamp(tileY + tileSize / 2, yStart, yEnd - 1);
    int sampleY = Mathf.Clamp(sourceTexture.height - 1 - centerY, 0, sourceTexture.height - 1);
    int index = sampleY * sourceTexture.width + centerX;

    switch (currentPixelFormat)
    {
      case PixelFormatColor32:
      {
        Color32 c32 = pixelDataBytes[index];
        r = c32.r;
        g = c32.g;
        b = c32.b;
        return new Color(c32.r / 255f, c32.g / 255f, c32.b / 255f, c32.a / 255f);
      }
      case PixelFormatFloat:
      {
        Color pixel = pixelDataFloat[index];

        // Check for invalid data
        if (float.IsNaN(pixel.r) || float.IsInfinity(pixel.r) ||
            float.IsNaN(pixel.g) || float.IsInfinity(pixel.g) ||
            float.IsNaN(pixel.b) || float.IsInfinity(pixel.b))
        {
          Debug.LogWarning($"DataDecoder: Invalid pixel data detected in tile {tileIndex}: {pixel}");
          r = 0;
          g = 0;
          b = 0;
          return Color.clear;
        }

        // Convert float (0-1) to bytes (0-255), clamping to prevent overflow
        r = FloatToByte(pixel.r);
        g = FloatToByte(pixel.g);
        b = FloatToByte(pixel.b);
        return pixel;
      }
      default:
        r = 0;
        g = 0;
        b = 0;
        return Color.clear;
    }
  }

  /// <summary>
  /// Reads the length value from the center pixel of the first tile (at position 0, 0).
  /// The length is stored as a 24-bit integer (3 bytes) in the RGB channels.
  /// </summary>
  private int ReadLength(int tileSize, int tilesPerColumn)
  {
    // Get RGB bytes from the first tile (tile 0)
    // GetTileRGB handles NaN/infinity checks internally
    GetTileRGB(0, tileSize, tilesPerColumn, out byte r, out byte g, out byte b);

    // Convert bytes to int (little-endian, 24-bit)
    int length = r | (g << 8) | (b << 16);

    return length;
  }

  private bool HasPixelData()
  {
    switch (currentPixelFormat)
    {
      case PixelFormatColor32:
        return pixelDataBytes != null && pixelDataBytes.Length > 0;
      case PixelFormatFloat:
        return pixelDataFloat != null && pixelDataFloat.Length > 0;
      default:
        return false;
    }
  }

  private byte FloatToByte(float value)
  {
    return (byte)Mathf.Clamp(Mathf.RoundToInt(value * 255f), 0, 255);
  }

}