// Copyright (c) 2022 The Khronos Group Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Validates ray query instructions from SPV_KHR_ray_query

#include <string>

#include "source/val/instruction.h"
#include "source/val/validate.h"
#include "source/val/validation_state.h"

namespace spvtools {
namespace val {

bool IsInterfaceVariable(ValidationState_t& _, const Instruction* inst,
                         spv::ExecutionModel model) {
  bool foundInterface = false;
  for (auto entry_point : _.entry_points()) {
    const auto* models = _.GetExecutionModels(entry_point);
    if (models->find(model) == models->end()) return false;
    for (const auto& desc : _.entry_point_descriptions(entry_point)) {
      for (auto interface : desc.interfaces) {
        if (inst->id() == interface) {
          foundInterface = true;
          break;
        }
      }
    }
  }
  return foundInterface;
}

spv_result_t ValidateEmitMeshTasks(ValidationState_t& _,
                                   const Instruction* inst) {
  _.function(inst->function()->id())
      ->RegisterExecutionModelLimitation([](spv::ExecutionModel model,
                                            std::string* message) {
        if (model != spv::ExecutionModel::TaskEXT) {
          if (message) {
            *message = "OpEmitMeshTasksEXT requires TaskEXT execution model";
          }
          return false;
        }
        return true;
      });

  const uint32_t group_count_x = _.GetOperandTypeId(inst, 0);
  if (!_.IsUnsignedIntScalarType(group_count_x) ||
      _.GetBitWidth(group_count_x) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Group Count X must be a 32-bit unsigned int scalar";
  }

  const uint32_t group_count_y = _.GetOperandTypeId(inst, 1);
  if (!_.IsUnsignedIntScalarType(group_count_y) ||
      _.GetBitWidth(group_count_y) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Group Count Y must be a 32-bit unsigned int scalar";
  }

  const uint32_t group_count_z = _.GetOperandTypeId(inst, 2);
  if (!_.IsUnsignedIntScalarType(group_count_z) ||
      _.GetBitWidth(group_count_z) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Group Count Z must be a 32-bit unsigned int scalar";
  }

  if (inst->operands().size() == 4) {
    const auto payload = _.FindDef(inst->GetOperandAs<uint32_t>(3));
    if (payload->opcode() != spv::Op::OpVariable) {
      return _.diag(SPV_ERROR_INVALID_DATA, inst)
             << "Payload must be the result of a OpVariable";
    }
    if (payload->GetOperandAs<spv::StorageClass>(2) !=
        spv::StorageClass::TaskPayloadWorkgroupEXT) {
      return _.diag(SPV_ERROR_INVALID_DATA, inst)
             << "Payload OpVariable must have a storage class of "
                "TaskPayloadWorkgroupEXT";
    }
  }
  return SPV_SUCCESS;
}

spv_result_t ValidateSetMeshOutputs(ValidationState_t& _,
                                    const Instruction* inst) {
  _.function(inst->function()->id())
      ->RegisterExecutionModelLimitation([](spv::ExecutionModel model,
                                            std::string* message) {
        if (model != spv::ExecutionModel::MeshEXT) {
          if (message) {
            *message = "OpSetMeshOutputsEXT requires MeshEXT execution model";
          }
          return false;
        }
        return true;
      });

  const uint32_t vertex_count = _.GetOperandTypeId(inst, 0);
  if (!_.IsUnsignedIntScalarType(vertex_count) ||
      _.GetBitWidth(vertex_count) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Vertex Count must be a 32-bit unsigned int scalar";
  }

  const uint32_t primitive_count = _.GetOperandTypeId(inst, 1);
  if (!_.IsUnsignedIntScalarType(primitive_count) ||
      _.GetBitWidth(primitive_count) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Primitive Count must be a 32-bit unsigned int scalar";
  }

  // Will only validate if constants are used (or spec constant frozen)
  uint64_t vertex_count_value = 0;
  if (_.EvalConstantValUint64(inst->GetOperandAs<uint32_t>(0),
                              &vertex_count_value)) {
    _.function(inst->function()->id())
        ->RegisterLimitation(
            [vertex_count_value](const ValidationState_t& state,
                                 const Function* entry_point,
                                 std::string* message) {
              const uint32_t output_vertices =
                  state.GetOutputVertices(entry_point->id());
              if (vertex_count_value > output_vertices) {
                *message =
                    "OpSetMeshOutputsEXT Vertex Count (" +
                    std::to_string(vertex_count_value) +
                    ") is larger than the OutputVertices in OpExecutionMode (" +
                    std::to_string(output_vertices) + ").";
                return false;
              }
              return true;
            });
  }
  uint64_t primitive_count_value = 0;
  if (_.EvalConstantValUint64(inst->GetOperandAs<uint32_t>(1),
                              &primitive_count_value)) {
    _.function(inst->function()->id())
        ->RegisterLimitation(
            [primitive_count_value](const ValidationState_t& state,
                                    const Function* entry_point,
                                    std::string* message) {
              const uint32_t output_primitives =
                  state.GetOutputPrimitivesEXT(entry_point->id());
              if (primitive_count_value > output_primitives) {
                *message = "OpSetMeshOutputsEXT Primitive Count (" +
                           std::to_string(primitive_count_value) +
                           ") is larger than the OutputPrimitivesEXT in "
                           "OpExecutionMode (" +
                           std::to_string(output_primitives) + ").";
                return false;
              }
              return true;
            });
  }

  return SPV_SUCCESS;
}

spv_result_t ValidateMeshVariable(ValidationState_t& _,
                                  const Instruction* inst) {
  if (!_.HasCapability(spv::Capability::MeshShadingEXT)) {
    return SPV_SUCCESS;
  }
  bool is_mesh_interface_var =
      IsInterfaceVariable(_, inst, spv::ExecutionModel::MeshEXT);
  bool is_frag_interface_var =
      IsInterfaceVariable(_, inst, spv::ExecutionModel::Fragment);

  const spv::StorageClass storage_class =
      inst->GetOperandAs<spv::StorageClass>(2);
  bool storage_output = (storage_class == spv::StorageClass::Output);
  bool storage_input = (storage_class == spv::StorageClass::Input);

  if (_.HasDecoration(inst->id(), spv::Decoration::PerPrimitiveEXT)) {
    if (is_frag_interface_var && !storage_input) {
      return _.diag(SPV_ERROR_INVALID_DATA, inst)
             << "PerPrimitiveEXT decoration must be applied only to "
                "variables in the Input Storage Class in the Fragment "
                "Execution Model.";
    }

    if (is_mesh_interface_var && !storage_output) {
      return _.diag(SPV_ERROR_INVALID_DATA, inst)
             << _.VkErrorID(4336)
             << "PerPrimitiveEXT decoration must be applied only to "
                "variables in the Output Storage Class in the "
                "Storage Class in the MeshEXT Execution Model.";
    }
  }

  // This only applies to user interface variables, not built-ins (they
  // are validated with the rest of the builtin)
  if (is_mesh_interface_var && storage_output &&
      !_.HasDecoration(inst->id(), spv::Decoration::BuiltIn)) {
    const Instruction* pointer_inst = _.FindDef(inst->type_id());
    if (pointer_inst->opcode() == spv::Op::OpTypePointer) {
      if (!_.IsArrayType(pointer_inst->word(3))) {
        return _.diag(SPV_ERROR_INVALID_DATA, inst)
               << "In the MeshEXT Execution Mode, all Output Variables "
                  "must contain an Array.";
      }
    }
  }

  return SPV_SUCCESS;
}

spv_result_t MeshShadingPass(ValidationState_t& _, const Instruction* inst) {
  const spv::Op opcode = inst->opcode();
  switch (opcode) {
    case spv::Op::OpEmitMeshTasksEXT:
      return ValidateEmitMeshTasks(_, inst);
    case spv::Op::OpSetMeshOutputsEXT:
      return ValidateSetMeshOutputs(_, inst);
    case spv::Op::OpVariable:
      return ValidateMeshVariable(_, inst);
    // No validation rules (for the moment).
    case spv::Op::OpWritePackedPrimitiveIndices4x8NV:
    default:
      break;
  }

  return SPV_SUCCESS;
}

}  // namespace val
}  // namespace spvtools