#include "rive/semantic/semantic_provider.hpp"

#include "rive/artboard.hpp"
#include "rive/component.hpp"
#include "rive/container_component.hpp"
#include "rive/math/mat2d.hpp"
#include "rive/node.hpp"
#include "rive/semantic/semantic_data.hpp"
#include "semantic_inference_registry.hpp"

using namespace rive;

static bool tryNodeWorldBounds(Node* node, AABB& outBounds)
{
  if (node == nullptr)
  {
    return false;
  }

  const auto localBounds = node->localBounds();
  if (localBounds.isEmptyOrNaN())
  {
    return false;
  }
  const auto worldTransform = node->worldTransform();
  const auto worldBounds = worldTransform.mapBoundingBox(localBounds);
  if (worldBounds.isEmptyOrNaN())
  {
    return false;
  }
  outBounds = worldBounds;
  return true;
}

AABB SemanticProvider::rootTransformAABB(Artboard* ab, const AABB& bounds)
{
  AABB out = AABB::forExpansion();
  AABB::expandTo(out, ab->rootTransform({bounds.minX, bounds.minY}));
  AABB::expandTo(out, ab->rootTransform({bounds.maxX, bounds.minY}));
  AABB::expandTo(out, ab->rootTransform({bounds.maxX, bounds.maxY}));
  AABB::expandTo(out, ab->rootTransform({bounds.minX, bounds.maxY}));
  return out;
}

bool SemanticProvider::canInferSemantics(Component* component)
{
  return supportsInferredSemantics(component);
}

ResolvedSemanticData SemanticProvider::resolveSemanticData(Component* component)
{
  ResolvedSemanticData out;
  if (component == nullptr || !component->is<Node>())
  {
    return out;
  }

  if (auto* sd = component->as<Node>()->firstChild<SemanticData>())
  {
    out.hasSemantics = true;
    out.role = sd->role();
    out.label = sd->label();
    return out;
  }

  // Inference fallback (Text -> label)
  resolveInferredSemantics(component, out);

  return out;
}

// Computes world-space bounds for a component. Three strategies:
// 1. Drawable nodes: localBounds() → worldTransform → rootTransform
// 2. Container nodes: merge bounds of all direct Node children
// 3. Fallback: world transform translation (degenerate point bounds)
AABB SemanticProvider::semanticBounds(Component* component)
{
  if (component == nullptr || !component->is<Node>())
  {
    return AABB();
  }

  auto* node = component->as<Node>();

  AABB worldBounds;
  if (tryNodeWorldBounds(node, worldBounds))
  {
    auto* ab = node->artboard();
    if (ab != nullptr)
    {
      return rootTransformAABB(ab, worldBounds);
    }
    return worldBounds;
  }

  // If this semantic node is a non-drawable container/group, derive bounds
  // from visual descendant nodes. ContainerComponent::forEachChild walks
  // the whole subtree (not just direct children): the predicate runs for
  // each descendant, and returning `true` tells forEachChild to keep
  // descending into that descendant's children. We always return true so
  // every Node descendant contributes bounds; non-Node descendants and
  // Nodes without valid bounds are skipped but their subtrees are still
  // walked.
  if (component->is<ContainerComponent>())
  {
    AABB merged = AABB::forExpansion();
    bool hasDescendantBounds = false;
    component->as<ContainerComponent>()->forEachChild(
        [&](Component* child) -> bool {
          if (child == nullptr || !child->is<Node>())
          {
            return true;
          }

          AABB childBounds;
          if (!tryNodeWorldBounds(child->as<Node>(), childBounds))
          {
            return true;
          }

          merged.expand(childBounds);
          hasDescendantBounds = true;
          return true;
        });

    if (hasDescendantBounds)
    {
      auto* ab = node->artboard();
      if (ab != nullptr)
      {
        return rootTransformAABB(ab, merged);
      }
      return merged;
    }
  }

  // Fallback for nodes/components without local bounds.
  const auto worldTransform = node->worldTransform();
  const float x = worldTransform[4];
  const float y = worldTransform[5];
  auto* ab = node->artboard();
  if (ab != nullptr)
  {
    Vec2D pos = ab->rootTransform(Vec2D(x, y));
    return AABB(pos.x, pos.y, pos.x, pos.y);
  }
  return AABB(x, y, x, y);
}