#include "rive/assets/shader_asset.hpp"
#include "rive/signed_content_header.hpp"
// The script module's copy only reads bytes its host already admitted.
#if defined(WITH_RIVE_SCRIPTING) && !defined(RIVE_WASM_MODULE)
#define RIVE_SHADER_SIGNATURE_GATE
#include "rive/file.hpp"
#include "rive/importers/text_asset_importer.hpp"
#endif
#include <cstdio>

using namespace rive;

bool rive::ShaderAsset::decode(Span<const uint8_t> data, Factory* factory)
{
  // `data` is always a SignedContentHeader envelope
  // (`[flags:1][sig:64?][RSTB...]`) shared with ScriptAsset. Raw-RSTB
  // callers (editor live-preview) must prepend a single 0x00 flag byte.
  SignedContentHeader envelope(std::move(data));
  if (!envelope.isValid())
  {
    return false;
  }
  auto inner = envelope.content();
#ifdef RIVE_SHADER_SIGNATURE_GATE
  // Referenced shaders carry their own signature; in-band ones verify as a
  // group once the file reaches its signature.
  m_verified = envelope.isSigned() &&
               verifiesContentSignature(envelope.signature(), inner);
  m_refusalLogged = false;
#endif
  m_bytes = SimpleArray<uint8_t>(inner.data(), inner.size());
  m_indexed = false;
  bool indexed = admit();
#ifdef WITH_RIVE_SCRIPTING
  // Loaded after its file, so its own signature was the last to check.
  if (m_importFinished)
  {
    logRefusal();
  }
#endif
  return indexed;
}

bool ShaderAsset::accepted() const
{
#ifdef RIVE_SHADER_SIGNATURE_GATE
  return !m_imported || File::acceptsScript(m_verified, m_requireSignedScripts);
#else
  return true;
#endif
}

bool ShaderAsset::admit()
{
  if (m_indexed || !accepted())
  {
    return true;
  }
  m_indexed = index();
  return m_indexed;
}

#ifdef WITH_RIVE_SCRIPTING
void ShaderAsset::finishImport()
{
  m_importFinished = true;
  logRefusal();
}

void ShaderAsset::logRefusal()
{
  if (m_indexed || m_refusalLogged || m_bytes.empty() || accepted())
  {
    return;
  }
  m_refusalLogged = true;
  fprintf(
      stderr,
      "Shader '%s' is unavailable: its signature did not verify.\n",
      name().c_str());
}
#endif

bool ShaderAsset::index()
{
  m_index.clear();
  m_pairs.clear();
  m_entryBindingsOffset = 0;
  m_entryBindingsSize = 0;

  const uint8_t* p = m_bytes.data();
  const uint8_t* end = p + m_bytes.size();

  // Header: magic(u32LE) + version(u16LE) + variant_count(u8) +
  //         section_count(u8)
  if (m_bytes.size() < 8)
  {
    return false;
  }
  uint32_t magic = (uint32_t)p[0] | ((uint32_t)p[1] << 8) |
                   ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
  if (magic != 0x52535442u) // "RSTB"
  {
    return false;
  }
  uint16_t version = (uint16_t)(p[4] | (p[5] << 8));
  if (version != 4)
  {
    return false;
  }
  uint8_t variantCount = p[6];
  uint8_t sectionCount = p[7];
  p += 8;

  // VariantDescriptor: target(u8) + blob_offset(u32) + blob_size(u32) = 9
  constexpr size_t kVariantDescriptorSize = 9;
  if (p + (size_t)variantCount * kVariantDescriptorSize > end)
  {
    return false;
  }
  for (uint8_t v = 0; v < variantCount; v++)
  {
    uint8_t target = p[0];
    uint32_t blobOffset = (uint32_t)p[1] | ((uint32_t)p[2] << 8) |
                          ((uint32_t)p[3] << 16) | ((uint32_t)p[4] << 24);
    uint32_t blobSize = (uint32_t)p[5] | ((uint32_t)p[6] << 8) |
                        ((uint32_t)p[7] << 16) | ((uint32_t)p[8] << 24);
    p += kVariantDescriptorSize;
    m_index[target] = {blobOffset, blobSize};
  }

  // Tagged metadata sections.
  for (uint8_t sec = 0; sec < sectionCount; sec++)
  {
    if (p + 3 > end)
    {
      return false;
    }
    uint8_t tag = p[0];
    uint16_t length = (uint16_t)(p[1] | (p[2] << 8));
    p += 3;
    if (p + length > end)
    {
      return false;
    }

    if (tag == 1 && length >= 1)
    {
      // Tag 1: texture-sampler pairs.
      uint8_t pairCount = p[0];
      if (length >= 1 + (size_t)pairCount * 4)
      {
        const uint8_t* pd = p + 1;
        for (uint8_t i = 0; i < pairCount; i++)
        {
          m_pairs.push_back({pd[0], pd[1], pd[2], pd[3]});
          pd += 4;
        }
      }
    }
    else if (tag == 2 && length >= 1)
    {
      // Tag 2: entry point bindings, which ore decodes.
      m_entryBindingsOffset = (size_t)(p - m_bytes.data());
      m_entryBindingsSize = length;
    }
    // Unknown tags: skip.
    p += length;
  }

  // p now points to the start of the blob data section.
  size_t blobDataStart = (size_t)(p - m_bytes.data());

  // Fix up offsets to absolute and validate that each variant's range
  // stays within m_bytes. Arithmetic widened to size_t so an adversarial
  // (offset, size) pair can't overflow uint32_t past the bounds check.
  for (auto& kv : m_index)
  {
    size_t absOffset = (size_t)kv.second.offset + blobDataStart;
    if (absOffset > m_bytes.size() ||
        (size_t)kv.second.size > m_bytes.size() - absOffset)
    {
      m_index.clear();
      return false;
    }
    kv.second.offset = (uint32_t)absOffset;
  }

  return true;
}

Span<const uint8_t> ShaderAsset::findShader(uint8_t target) const
{
  if (!m_indexed)
  {
    return {};
  }
  auto it = m_index.find(target);
  if (it == m_index.end())
  {
    return {};
  }
  const auto& v = it->second;
  return Span<const uint8_t>(m_bytes.data() + v.offset, v.size);
}