#include <triton/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.h>
#include <triton/Dialect/TritonNvidiaGPU/IR/Dialect.h>
#include <triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h>
#include <triton/Tools/LayoutUtils.h>

#include "mlir/Dialect/Arith/IR/Arith.h"

namespace tt = mlir::triton;
namespace ttg = mlir::triton::gpu;

namespace mlir::triton::nvidia_gpu {

ttg::SharedEncodingTrait getEncodingFromDescriptor(Operation *op,
                                                   RankedTensorType tensorType,
                                                   Value desc) {
  auto descType = cast<TensorDescType>(desc.getType());
  Attribute encoding = descType.getSharedLayout();
  if (!encoding) {
    constexpr auto msg =
        "Internal Error: Tensor descriptor should have encoding set";
    if (op)
      op->emitError() << msg;
    llvm::report_fatal_error(msg);
  }
  auto sharedEnc = cast<ttg::SharedEncodingTrait>(encoding);
  if (descType.getShape() == tensorType.getShape())
    return sharedEnc;

  return ttg::updateEncodingForShape(op, sharedEnc, tensorType);
}

bool hasCGABroadcast(ttg::MemDescType memDescType) {
  auto kBlock = StringAttr::get(memDescType.getContext(), "block");
  return ttg::toLinearLayout(memDescType)
             .getFreeVariableMasks()
             .lookup(kBlock) != 0;
}

Value sextI16ToI32Indices(Value indices, OpBuilder &builder, Location loc) {
  auto indicesType = cast<RankedTensorType>(indices.getType());
  if (indicesType.getElementType().isInteger(32))
    return indices;
  return arith::ExtSIOp::create(
      builder, loc, indicesType.clone(builder.getI32Type()), indices);
}

FailureOr<int> getTMASwizzleMode(Location loc, tt::TensorDescInterface ty) {
  auto encoding = ty.getSharedLayout();
  auto mmaEncoding =
      dyn_cast_if_present<ttg::NVMMASharedEncodingAttr>(encoding);
  unsigned swizzleBytes = mmaEncoding ? mmaEncoding.getSwizzlingByteWidth() : 0;
  if (!mmaEncoding) {
    auto swizzledEnc = dyn_cast<ttg::SwizzledSharedEncodingAttr>(encoding);
    if (!swizzledEnc || swizzledEnc.getVec() != 1 ||
        swizzledEnc.getPerPhase() != 1 || swizzledEnc.getMaxPhase() != 1) {
      return emitError(loc)
             << "unhandled shared memory layout for TMA descriptor: "
             << encoding;
    }
  }

  bool fp4Padded = isFp4Padded(encoding);
  if (fp4Padded && swizzleBytes != 128) {
    return emitError(loc) << "fp4 padded operands (elem type .b4x16_p64) only "
                             "supports 128-byte swizzling, but got "
                          << swizzleBytes;
  }

  int32_t swizzleMode = 0;
  if (swizzleBytes == 128) {
    swizzleMode = 3;
  } else if (swizzleBytes == 64) {
    swizzleMode = 2;
  } else if (swizzleBytes == 32) {
    swizzleMode = 1;
  } else {
    assert(swizzleBytes == 0);
  }
  return swizzleMode;
}

enum TMA_ELEMENT_TYPES {
  TMA_U8 = 0,
  TMA_U16 = 1,
  TMA_U32 = 2,
  TMA_S32 = 3,
  TMA_U64 = 4,
  TMA_S64 = 5,
  TMA_F16 = 6,
  TMA_F32 = 7,
  TMA_F32_FTZ = 8,
  TMA_F64 = 9,
  TMA_BF16 = 10,
  TMA_TF32 = 11,
  TMA_TF32_FTZ = 12,
  TMA_B4X16 = 13,
  TMA_B4X16_P64 = 14,
  TMA_B6X16_P32 = 15,
  TMA_B6P2X16 = 15,
};

FailureOr<int> getTMAElementType(Location loc, tt::TensorDescInterface ty) {
  auto encoding = ty.getSharedLayout();
  bool fp4Padded = isFp4Padded(encoding);

  if (fp4Padded)
    return TMA_B4X16_P64;

  auto elemTy = ty.getElementType();
  if (elemTy.isBF16()) {
    return TMA_BF16;
  } else if (elemTy.isF16()) {
    return TMA_F16;
  } else if (elemTy.isF32()) {
    return TMA_F32;
  } else if (elemTy.isF64()) {
    return TMA_F64;
  }

  auto elemSize = elemTy.getIntOrFloatBitWidth() / 8;
  switch (elemSize) {
  case 1:
    return TMA_U8;
  case 2:
    return TMA_U16;
  case 4:
    return elemTy.isSignedInteger() ? TMA_S32 : TMA_U32;
  case 8:
    return elemTy.isSignedInteger() ? TMA_S64 : TMA_U64;
  default:
    break;
  }
  return emitError(loc)
         << "Tensor descriptor element type must have size 1, 2, or 4 but got "
         << elemSize;
}

LogicalResult createTMADesc(Value tmaPtr, MakeTensorDescOp op,
                            OpBuilder &builder) {
  using namespace mlir;
  auto loc = op.getLoc();
  auto mkI32Constant = [&](int32_t val) {
    return arith::ConstantOp::create(builder, loc, builder.getI32Type(),
                                     builder.getI32IntegerAttr(val));
  };

  auto elemType = op.getBase().getType().getPointeeType();
  auto elemSize = elemType.getIntOrFloatBitWidth() / 8;
  auto encoding = op.getType().getSharedLayout();
  auto mmaEncoding =
      llvm::dyn_cast_or_null<gpu::NVMMASharedEncodingAttr>(encoding);
  bool fp4Padded = mmaEncoding && mmaEncoding.getFp4Padded();

  auto shapePerCTA = gpu::getShapePerCTA(encoding, op.getType().getShape());
  // MakeTensorDescOp creates tiled descriptors (not im2col)
  auto blockShape = getTMABlockShape(encoding, shapePerCTA,
                                     /*packedSize=*/false, gpu::TMAMode::Tiled);
  auto contigDimSize = blockShape.back();

  llvm::SmallVector<Value> boxDim;
  if (fp4Padded && contigDimSize != 128) {
    return op->emitError(
        "FP4 padded loads require 128 elements or more in the last dim");
  }
  boxDim.push_back(mkI32Constant(contigDimSize));
  for (int k = shapePerCTA.size() - 2; k >= 0; --k)
    boxDim.push_back(mkI32Constant(blockShape[k]));

  if (!mmaEncoding) {
    auto swizzledEnc =
        dyn_cast_if_present<gpu::SwizzledSharedEncodingAttr>(encoding);
    if (!swizzledEnc || swizzledEnc.getVec() != 1 ||
        swizzledEnc.getPerPhase() != 1 || swizzledEnc.getMaxPhase() != 1) {
      op->emitError() << "Unhandled encoding type";
      return failure();
    }
  }

  auto maybeSwizzleMode = getTMASwizzleMode(loc, op.getType());
  if (failed(maybeSwizzleMode))
    return failure();
  auto swizzleMode = *maybeSwizzleMode;

  Value elemSizeVal = arith::ConstantOp::create(
      builder, loc, builder.getI64Type(), builder.getI64IntegerAttr(elemSize));

  SmallVector<Value> globalDim(llvm::reverse(op.getShape()));
  SmallVector<Value> globalStride;
  for (int k = op.getStrides().size() - 2; k >= 0; --k) {
    globalStride.push_back(op.getStrides()[k]);
  }

  if (fp4Padded) {
    // Convert number of bytes to number of mxfp4 elements
    globalDim[0] =
        arith::MulIOp::create(builder, loc, globalDim[0], mkI32Constant(2));
  }

  SmallVector<Value> elementStride(globalDim.size(), mkI32Constant(1));

  for (int i = 0; i < globalStride.size(); ++i)
    globalStride[i] =
        arith::MulIOp::create(builder, loc, globalStride[i], elemSizeVal);

  auto elemTypeEnum = getTMAElementType(loc, op.getType());
  if (failed(elemTypeEnum))
    return failure();

  auto fillMode = (op.getPadding() == triton::PaddingOption::PAD_NAN) ? 1 : 0;

  TensormapCreateOp::create(
      builder, loc,
      /*desc_ptr=*/tmaPtr,
      /*global_address=*/op.getBase(),
      /*box_dim=*/boxDim,
      /*global_dim=*/globalDim,
      /*global_stride=*/globalStride,
      /*element_strides=*/elementStride,
      /*elem_type*/ builder.getI32IntegerAttr(*elemTypeEnum),
      /*interleave_layout*/ builder.getI32IntegerAttr(0),
      /*swizzle_mode=*/builder.getI32IntegerAttr(swizzleMode),
      /*fill_mode=*/builder.getI32IntegerAttr(fillMode));
  return success();
}

} // namespace mlir::triton::nvidia_gpu
