/*
 * Copyright 2026 Rive
 */

#include "utils/serialized_replay.hpp"
#include "utils/serialize_ops.hpp"
#include "rive/core/binary_reader.hpp"
#include "rive/math/mat2d.hpp"
#include "utils/no_op_renderer.hpp"
#include <unordered_map>
#include <vector>

using namespace rive;

namespace
{
// Absent ids mean a truncated or corrupt stream; callers fail instead of
// dereferencing a null resource.
template <typename T>
T* find(std::unordered_map<uint64_t, rcp<T>>& map, uint64_t id)
{
  auto it = map.find(id);
  return it != map.end() ? it->second.get() : nullptr;
}
} // namespace

bool rive::replaySerializedCommands(
    Span<const uint8_t> stream,
    Factory* factory,
    Renderer* renderer,
    const SerializedReplayHooks& hooks)
{
  BinaryReader reader(stream);
  if (reader.readByte() != 'S' || reader.readByte() != 'R' ||
      reader.readByte() != 'I' || reader.readByte() != 'V')
  {
    return false;
  }
  if (reader.readVarUint64() != 1)
  {
    return false;
  }

  std::unordered_map<uint64_t, rcp<RenderPath>> paths;
  std::unordered_map<uint64_t, rcp<RenderPaint>> paints;
  std::unordered_map<uint64_t, rcp<RenderShader>> shaders;
  std::unordered_map<uint64_t, rcp<RenderImage>> images;
  std::unordered_map<uint64_t, rcp<RenderBuffer>> buffers;
  std::unordered_map<uint64_t, rcp<ImageMeshInstances>> meshInstances;

  // Cache-as-bitmap. A canvas is declared by makeRenderCanvas, its content
  // arrives inline between canvasContentBegin/End, and the composite that
  // samples it is an ordinary drawImage of the same id -- so a canvas image
  // lands in `images` alongside the decoded ones.
  struct CanvasSize
  {
    uint32_t width;
    uint32_t height;
  };
  std::unordered_map<uint64_t, CanvasSize> canvasSizes;
  // Where the ops currently being read draw. The screen renderer at the
  // bottom, a host-supplied one for each open canvas, and the null sink for
  // a canvas the host declined -- its content still has to parse (later ops
  // reference resources declared inside it) but must not reach the screen.
  NoOpRenderer droppedContent;
  Renderer* active = renderer;
  // One entry per open canvas bracket. The id rides along with the renderer
  // it interrupted so an end can be checked against the begin that opened
  // it; without it the stack depth is the only thing validated, and a
  // mismatched end would close the wrong canvas and still report success.
  struct OpenCanvas
  {
    uint64_t id;
    Renderer* interruptedRenderer;
  };
  std::vector<OpenCanvas> interrupted;

  while (!reader.reachedEnd())
  {
    SerializeOp op = static_cast<SerializeOp>(reader.readVarUint64());
    if (reader.hasError())
      return false;
    switch (op)
    {
      case SerializeOp::makeRenderPath:
      {
        uint64_t id = reader.readVarUint64();
        // Geometry and fill rule arrive as later ops.
        paths[id] = factory->makeEmptyRenderPath();
        break;
      }
      case SerializeOp::makeRenderPaint:
      {
        uint64_t id = reader.readVarUint64();
        paints[id] = factory->makeRenderPaint();
        break;
      }
      case SerializeOp::rewind:
      {
        uint64_t id = reader.readVarUint64();
        RenderPath* path = find(paths, id);
        if (path == nullptr)
          return false;
        path->rewind();
        break;
      }
      case SerializeOp::fillRule:
      {
        uint64_t id = reader.readVarUint64();
        RenderPath* path = find(paths, id);
        if (path == nullptr)
          return false;
        path->fillRule(static_cast<FillRule>(reader.readVarUint64()));
        break;
      }
      case SerializeOp::addRawPath:
      {
        uint64_t id = reader.readVarUint64();
        RenderPath* path = find(paths, id);
        if (path == nullptr)
          return false;
        RawPath rp = deserializeRawPath(reader);
        path->addRawPath(rp);
        break;
      }
      case SerializeOp::color:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        paint->color(static_cast<unsigned int>(reader.readVarUint64()));
        break;
      }
      case SerializeOp::style:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        // The stream writes 0 for stroke and 1 for fill.
        bool stroked = reader.readVarUint64() == 0;
        paint->style(
            stroked ? RenderPaintStyle::stroke : RenderPaintStyle::fill);
        break;
      }
      case SerializeOp::thickness:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        paint->thickness(reader.readFloat32());
        break;
      }
      case SerializeOp::join:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        paint->join(static_cast<StrokeJoin>(reader.readVarUint64()));
        break;
      }
      case SerializeOp::cap:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        paint->cap(static_cast<StrokeCap>(reader.readVarUint64()));
        break;
      }
      case SerializeOp::feather:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        paint->feather(reader.readFloat32());
        break;
      }
      case SerializeOp::blendMode:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        paint->blendMode(static_cast<BlendMode>(reader.readVarUint64()));
        break;
      }
      case SerializeOp::additiveness:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        paint->additiveness(reader.readFloat32());
        break;
      }
      case SerializeOp::strokePosition:
      {
        uint64_t id = reader.readVarUint64();
        uint64_t position = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        // Renderers assert on a position they don't know.
        if (paint == nullptr || position > (uint64_t)StrokePosition::outside)
          return false;
        paint->strokePosition(static_cast<StrokePosition>(position));
        break;
      }
      case SerializeOp::shader:
      {
        uint64_t id = reader.readVarUint64();
        uint64_t shaderId = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        // The stream writes 0 for both nullptr and shader id 0. The map
        // resolves it since a missing entry yields a null rcp.
        paint->shader(shaders[shaderId]);
        break;
      }
      case SerializeOp::shaderTransform:
      {
        uint64_t id = reader.readVarUint64();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        Mat2D transform;
        for (int i = 0; i < 6; i++)
        {
          transform[i] = reader.readFloat32();
        }
        paint->shaderTransform(transform);
        break;
      }
      case SerializeOp::paintModulatedImage:
      {
        uint64_t id = reader.readVarUint64();
        uint64_t rawImageId = reader.readVarUint64();
        auto filter = static_cast<ImageFilter>(reader.readVarUint64());
        auto wrapX = static_cast<ImageWrap>(reader.readVarUint64());
        auto wrapY = static_cast<ImageWrap>(reader.readVarUint64());
        float xx = reader.readFloat32();
        float xy = reader.readFloat32();
        float yx = reader.readFloat32();
        float yy = reader.readFloat32();
        float tx = reader.readFloat32();
        float ty = reader.readFloat32();
        RenderPaint* paint = find(paints, id);
        if (paint == nullptr)
          return false;
        // Image id is offset by one; 0 means no image.
        RenderImage* image = rawImageId == 0 ? nullptr
                                             : images[rawImageId - 1].get();
        ImageSampler sampler;
        sampler.filter = filter;
        sampler.wrapX = wrapX;
        sampler.wrapY = wrapY;
        paint->modulatedImage(image, sampler, Mat2D(xx, xy, yx, yy, tx, ty));
        break;
      }
      case SerializeOp::makeLinearGradient:
      case SerializeOp::makeRadialGradient:
      {
        uint64_t id = reader.readVarUint64();
        size_t count = static_cast<size_t>(reader.readVarUint64());
        std::vector<ColorInt> colors(count);
        std::vector<float> stops(count);
        for (size_t i = 0; i < count; ++i)
        {
          colors[i] = static_cast<ColorInt>(reader.readVarUint64());
          stops[i] = reader.readFloat32();
        }
        float a = reader.readFloat32();
        float b = reader.readFloat32();
        float c = reader.readFloat32();
        if (op == SerializeOp::makeLinearGradient)
        {
          float d = reader.readFloat32();
          shaders[id] = factory->makeLinearGradient(
              a,
              b,
              c,
              d,
              colors.data(),
              stops.data(),
              count);
        }
        else
        {
          shaders[id] = factory->makeRadialGradient(
              a,
              b,
              c,
              colors.data(),
              stops.data(),
              count);
        }
        break;
      }
      case SerializeOp::makeAngularGradient:
      {
        uint64_t id = reader.readVarUint64();
        size_t count = static_cast<size_t>(reader.readVarUint64());
        std::vector<ColorInt> colors(count);
        std::vector<float> stops(count);
        for (size_t i = 0; i < count; ++i)
        {
          colors[i] = static_cast<ColorInt>(reader.readVarUint64());
          stops[i] = reader.readFloat32();
        }
        shaders[id] = factory->makeAngularGradient(
            colors.data(),
            stops.data(),
            count);
        break;
      }
      case SerializeOp::decodeImage:
      {
        uint64_t id = reader.readVarUint64();
        size_t size = static_cast<size_t>(reader.readVarUint64());
        Span<const uint8_t> data = reader.readBytes(size);
        images[id] = factory->decodeImage(data);
        break;
      }
      case SerializeOp::makeRenderBuffer:
      {
        uint64_t id = reader.readVarUint64();
        size_t size = static_cast<size_t>(reader.readVarUint64());
        auto type = static_cast<RenderBufferType>(reader.readVarUint64());
        auto flags = static_cast<RenderBufferFlags>(reader.readVarUint64());
        buffers[id] = factory->makeRenderBuffer(type, flags, size);
        break;
      }
      case SerializeOp::setVertexBufferData:
      case SerializeOp::setIndexBufferData:
      {
        uint64_t id = reader.readVarUint64();
        RenderBuffer* buf = find(buffers, id);
        if (buf == nullptr)
          return false;
        void* mapped = buf->map();
        if (op == SerializeOp::setVertexBufferData)
        {
          size_t n = buf->sizeInBytes() / sizeof(float);
          float* out = static_cast<float*>(mapped);
          for (size_t i = 0; i < n; ++i)
            out[i] = reader.readFloat32();
        }
        else
        {
          size_t n = buf->sizeInBytes() / sizeof(uint16_t);
          uint16_t* out = static_cast<uint16_t*>(mapped);
          for (size_t i = 0; i < n; ++i)
            out[i] = static_cast<uint16_t>(reader.readVarUint64());
        }
        buf->unmap();
        break;
      }
      case SerializeOp::save:
        active->save();
        break;
      case SerializeOp::restore:
        active->restore();
        break;
      case SerializeOp::transform:
      {
        float m[6];
        for (int i = 0; i < 6; ++i)
          m[i] = reader.readFloat32();
        active->transform(Mat2D(m[0], m[1], m[2], m[3], m[4], m[5]));
        break;
      }
      case SerializeOp::modulateOpacity:
        active->modulateOpacity(reader.readFloat32());
        break;
      case SerializeOp::modulateColor:
      {
        ColorInt color = (ColorInt)reader.readVarUint64();
        active->modulateColor(color, reader.readVarUint64() != 0);
        break;
      }
      case SerializeOp::drawPath:
      {
        uint64_t pathId = reader.readVarUint64();
        uint64_t paintId = reader.readVarUint64();
        RenderPath* path = find(paths, pathId);
        RenderPaint* paint = find(paints, paintId);
        if (path == nullptr || paint == nullptr)
          return false;
        active->drawPath(path, paint);
        break;
      }
      case SerializeOp::clipPath:
      {
        uint64_t pathId = reader.readVarUint64();
        RenderPath* path = find(paths, pathId);
        if (path == nullptr)
          return false;
        active->clipPath(path);
        break;
      }
      case SerializeOp::applyLayerMask:
      {
        uint64_t maskId = reader.readVarUint64();
        auto mode = static_cast<LayerMaskMode>(reader.readVarUint64());
        // Null when a decode failed or when the host declined the
        // canvas this id names. Skipping leaves the layer unmasked,
        // which is the same degradation the feature falls back to
        // everywhere else, and the rest of the frame still replays.
        if (RenderImage* mask = find(images, maskId))
        {
          active->applyLayerMask(mask, ImageSampler::LinearClamp(), mode);
        }
        break;
      }
      case SerializeOp::drawImage:
      case SerializeOp::drawImageAdditive:
      {
        uint64_t imageId = reader.readVarUint64();
        auto blend = static_cast<BlendMode>(reader.readVarUint64());
        float opacity = reader.readFloat32();
        // Only the additive variant carries the trailing float; the
        // plain op means an additiveness of 0.
        float additiveness = op == SerializeOp::drawImageAdditive
                                 ? reader.readFloat32()
                                 : 0.0f;
        // Null when a decode failed or when the host declined the
        // canvas this id names; either way there is nothing to
        // composite and the rest of the frame still replays.
        if (RenderImage* image = find(images, imageId))
        {
          active->drawImage(
              image,
              ImageSampler::LinearClamp(),
              blend,
              opacity,
              additiveness);
        }
        break;
      }
      case SerializeOp::drawImageMesh:
      case SerializeOp::drawImageMeshAdditive:
      {
        RenderImage* image = find(images, reader.readVarUint64());
        auto blend = static_cast<BlendMode>(reader.readVarUint64());
        float opacity = reader.readFloat32();
        RenderBuffer* pos = find(buffers, reader.readVarUint64());
        RenderBuffer* uvs = find(buffers, reader.readVarUint64());
        RenderBuffer* idx = find(buffers, reader.readVarUint64());
        float additiveness = op == SerializeOp::drawImageMeshAdditive
                                 ? reader.readFloat32()
                                 : 0.0f;
        // Draws referring to failed or invalid objects are dropped.
        if (image != nullptr && pos != nullptr && uvs != nullptr &&
            idx != nullptr)
        {
          uint32_t vertexCount = static_cast<uint32_t>(
              pos->sizeInBytes() / (2 * sizeof(float)));
          uint32_t indexCount = static_cast<uint32_t>(
              idx->sizeInBytes() / sizeof(uint16_t));
          active->drawImageMesh(
              image,
              ImageSampler::LinearClamp(),
              ref_rcp(pos),
              ref_rcp(uvs),
              ref_rcp(idx),
              vertexCount,
              indexCount,
              blend,
              opacity,
              additiveness);
        }
        break;
      }
      case SerializeOp::makeImageMeshInstances:
      {
        uint64_t id = reader.readVarUint64();
        size_t count = static_cast<size_t>(reader.readVarUint64());
        meshInstances[id] = factory->makeImageMeshInstances(count);
        break;
      }
      case SerializeOp::setImageMeshInstancesData:
      {
        uint64_t id = reader.readVarUint64();
        size_t count = static_cast<size_t>(reader.readVarUint64());
        // An unknown id is dropped, but its payload is still read
        // so the stream stays in step.
        ImageMeshInstances* mi = find(meshInstances, id);
        Span<ImageMeshInstanceData> dst = mi != nullptr
                                              ? mi->edit(count)
                                              : Span<ImageMeshInstanceData>();
        for (size_t i = 0; i < count; ++i)
        {
          ImageMeshInstanceData instance;
          float m[6];
          for (int j = 0; j < 6; ++j)
          {
            m[j] = reader.readFloat32();
          }
          instance.transform = Mat2D(m[0], m[1], m[2], m[3], m[4], m[5]);
          instance.uvTranslate.x = reader.readFloat32();
          instance.uvTranslate.y = reader.readFloat32();
          instance.uvScale.x = reader.readFloat32();
          instance.uvScale.y = reader.readFloat32();
          instance.opacity = reader.readFloat32();
          instance.additiveness = reader.readFloat32();
          if (i < dst.size())
          {
            dst[i] = instance;
          }
        }
        if (mi != nullptr)
        {
          mi->endEdit();
        }
        break;
      }
      case SerializeOp::drawImageMeshInstanced:
      {
        RenderImage* image = find(images, reader.readVarUint64());
        RenderBuffer* pos = find(buffers, reader.readVarUint64());
        RenderBuffer* uvs = find(buffers, reader.readVarUint64());
        RenderBuffer* idx = find(buffers, reader.readVarUint64());
        ImageMeshInstances* mi = find(meshInstances, reader.readVarUint64());
        if (image != nullptr && pos != nullptr && uvs != nullptr &&
            idx != nullptr && mi != nullptr)
        {
          uint32_t vertexCount = static_cast<uint32_t>(
              pos->sizeInBytes() / (2 * sizeof(float)));
          uint32_t indexCount = static_cast<uint32_t>(
              idx->sizeInBytes() / sizeof(uint16_t));
          active->drawImageMeshInstanced(
              image,
              ImageSampler::LinearClamp(),
              ref_rcp(pos),
              ref_rcp(uvs),
              ref_rcp(idx),
              vertexCount,
              indexCount,
              ref_rcp(mi));
        }
        break;
      }
      case SerializeOp::makeRenderCanvas:
      {
        uint64_t id = reader.readVarUint64();
        uint32_t w = static_cast<uint32_t>(reader.readVarUint64());
        uint32_t h = static_cast<uint32_t>(reader.readVarUint64());
        // No allocation here: the host mints the target when the
        // content actually opens, on whatever device it replays
        // against.
        canvasSizes[id] = {w, h};
        break;
      }
      case SerializeOp::canvasContentBegin:
      {
        uint64_t id = reader.readVarUint64();
        uint32_t clearColor = static_cast<uint32_t>(reader.readVarUint64());
        auto it = canvasSizes.find(id);
        if (it == canvasSizes.end())
        {
          return false; // content for a canvas never declared
        }
        interrupted.push_back({id, active});
        rcp<RenderImage> image;
        Renderer* content = hooks.onCanvasContentBegin
                                ? hooks.onCanvasContentBegin(
                                      id,
                                      it->second.width,
                                      it->second.height,
                                      clearColor,
                                      &image)
                                : nullptr;
        // A host with no offscreen target returns null. Its content is
        // parsed and dropped, and the composite that names this id
        // finds no image and draws nothing.
        if (content != nullptr)
        {
          images[id] = std::move(image);
        }
        else
        {
          // The raster was declined, so whatever the host may have
          // allocated holds no content for this bracket. Drop any
          // image it set, and any image a previous bracket left under
          // this id, so the composite draws nothing rather than stale
          // or uninitialized pixels.
          images.erase(id);
        }
        active = content != nullptr ? content : &droppedContent;
        break;
      }
      case SerializeOp::canvasContentEnd:
      {
        uint64_t id = reader.readVarUint64();
        if (interrupted.empty())
        {
          return false; // end without a matching begin
        }
        if (interrupted.back().id != id)
        {
          // Ending a canvas the innermost bracket never opened.
          // Honoring it would tell the host to close an unrelated
          // canvas and hand the wrong renderer back to the ops that
          // follow, so treat it as a malformed stream.
          return false;
        }
        if (hooks.onCanvasContentEnd)
        {
          hooks.onCanvasContentEnd(id);
        }
        active = interrupted.back().interruptedRenderer;
        interrupted.pop_back();
        break;
      }
      case SerializeOp::frame:
        if (hooks.onFrame)
          hooks.onFrame();
        break;
      case SerializeOp::frameSize:
      {
        uint32_t w = static_cast<uint32_t>(reader.readVarUint64());
        uint32_t h = static_cast<uint32_t>(reader.readVarUint64());
        if (hooks.onFrameSize)
          hooks.onFrameSize(w, h);
        break;
      }
      default:
        return false; // unknown opcode
    }
    if (reader.hasError())
      return false;
  }
  // A canvas left open means the stream was cut mid-bracket.
  return interrupted.empty();
}