#pragma once

#include "atomic.h"
#include <ck_tile/core.hpp>
#include <hip/amd_detail/amd_warp_functions.h>
#include <hip/hip_bf16.h>
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
#include <rocwmma/rocwmma.hpp>

#define HIPRT_INF_F __int_as_float(0x7f800000)
#define HIPRT_NEGINF_F __int_as_float(0xff800000)
#define HIPRT_NAN_F __int_as_float(0x7fffffff)
#define HIPRT_MIN_DENORM_F __int_as_float(0x00000001)
#define HIPRT_MAX_NORMAL_F __int_as_float(0x7f7fffff)
#define HIPRT_NEG_ZERO_F __int_as_float(0x80000000)
#define HIPRT_ZERO_F 0.0f
#define HIPRT_ONE_F 1.0f

/* double precision constants */
#define HIPRT_INF __hiloint2double(0x7ff00000, 0x00000000)
#define HIPRT_NAN __hiloint2double(0xfff80000, 0x00000000)

#define uint unsigned int
#define uchar unsigned char
#define ushort unsigned short

#define TL_DEVICE __forceinline__ __device__
#define TL_DEVICE_NOINLINE __noinline__ __device__

#define TILELANG_CHECK(stmt)                                                   \
  do {                                                                         \
    hipError_t __err = (stmt);                                                 \
    if (__err != hipSuccess) {                                                 \
      snprintf(error_buf, ERROR_BUF_SIZE, "%s:%d: %s - %s", __FILE__,          \
               __LINE__, hipGetErrorName(__err), hipGetErrorString(__err));    \
      return -1;                                                               \
    }                                                                          \
  } while (0)

#define TILELANG_CHECK_LAST_ERROR(kernel_name)                                 \
  do {                                                                         \
    hipError_t __err = hipGetLastError();                                      \
    if (__err != hipSuccess) {                                                 \
      snprintf(error_buf, ERROR_BUF_SIZE, "kernel_name: %s - %s",              \
               hipGetErrorName(__err), hipGetErrorString(__err));              \
      return -1;                                                               \
    }                                                                          \
  } while (0)

#define half _Float16
#define __float2half_rn(x) half(x)

#define hpow __ocml_pown_f16
#define hsqrt __ocml_sqrt_f16

using float16_t = _Float16;
using float16x2 =
    __attribute__((__vector_size__(2 * sizeof(float16_t)))) float16_t;
using float16x4 =
    __attribute__((__vector_size__(4 * sizeof(float16_t)))) float16_t;
using float16x8 =
    __attribute__((__vector_size__(8 * sizeof(float16_t)))) float16_t;
using float16x16 =
    __attribute__((__vector_size__(16 * sizeof(float16_t)))) float16_t;

using half_t = float16_t;

using bfloat16_t = hip_bfloat16;

struct bfloat16x2 {
  bfloat16_t x, y;
};

struct bfloat16x4 {
  bfloat16_t data[4];
};

struct bfloat16x8 {
  bfloat16_t data[8];
};

struct bfloat16x16 {
  bfloat16_t data[16];
};

typedef
    __attribute__((__vector_size__(4 * sizeof(short)))) short bfloat16x4_vec;
typedef
    __attribute__((__vector_size__(8 * sizeof(short)))) short bfloat16x8_vec;

using int32x4 = __attribute__((__vector_size__(4 * sizeof(int)))) int;
using int32x16 = __attribute__((__vector_size__(16 * sizeof(int)))) int;
using float32x4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
using float32x16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
using float32x32 = __attribute__((__vector_size__(32 * sizeof(float)))) float;

using int8x4 = __attribute__((__vector_size__(4 * sizeof(int8_t)))) int8_t;

// Pack two half_t values.
TL_DEVICE unsigned __pack_half2(const half_t x, const half_t y) {
  unsigned v0 = *((unsigned short *)&x);
  unsigned v1 = *((unsigned short *)&y);
  return (v1 << 16) | v0;
}

// Pack two bfloat16_t values.
TL_DEVICE unsigned __pack_bfloat162(const bfloat16_t x, const bfloat16_t y) {
  unsigned v0 = *((unsigned short *)&x);
  unsigned v1 = *((unsigned short *)&y);
  return (v1 << 16) | v0;
}

TL_DEVICE int tl_dp4a_fallback(const int a, const int b, int c) {
#pragma unroll
  for (int i = 0; i < 4; ++i) {
    const int ai = static_cast<int8_t>((a >> (8 * i)) & 0xff);
    const int bi = static_cast<int8_t>((b >> (8 * i)) & 0xff);
    c += ai * bi;
  }
  return c;
}

#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) ||    \
    defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__)
#define TL_AMDGPU_HAS_SUDOT4 1
#endif

#if defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) ||       \
    defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) ||       \
    defined(__gfx950__) || defined(__gfx1011__) || defined(__gfx1012__) ||     \
    defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) ||    \
    defined(__gfx1034__) || defined(__gfx1035__) ||                            \
    defined(TL_AMDGPU_HAS_SUDOT4)
#define TL_AMDGPU_HAS_SDOT4 1
#endif

TL_DEVICE int tl_dp4a(const int a, const int b, const int c) {
#if defined(TL_AMDGPU_HAS_SUDOT4)
  return __builtin_amdgcn_sudot4(true, a, true, b, c, false);
#elif defined(TL_AMDGPU_HAS_SDOT4)
  return __builtin_amdgcn_sdot4(a, b, c, false);
#else
  return tl_dp4a_fallback(a, b, c);
#endif
}

template <typename InDatatype, typename OutDatatype>
TL_DEVICE void DP4A(const InDatatype *a, const InDatatype *b, OutDatatype *c) {
  static_assert(sizeof(InDatatype) == 1,
                "DP4A expects a pointer to packed int8 lanes");
  static_assert(sizeof(OutDatatype) == sizeof(int),
                "DP4A expects 4-byte accumulator/output type");
  int a_int;
  int b_int;
  int c_int;
  __builtin_memcpy(&a_int, a, sizeof(a_int));
  __builtin_memcpy(&b_int, b, sizeof(b_int));
  __builtin_memcpy(&c_int, c, sizeof(c_int));
  const int out = tl_dp4a(a_int, b_int, c_int);
  __builtin_memcpy(c, &out, sizeof(out));
}

// __habs overloads for hip_bfloat16 and float16_t to resolve ambiguity on ROCm.
// hip_bfloat16 != __hip_bfloat16, and float16_t != __half, so the standard
// __habs overloads don't match exactly, causing ambiguous overload errors.
// Use __builtin_memcpy instead of reinterpret_cast to avoid strict-aliasing UB.
__device__ __forceinline__ hip_bfloat16 __habs(hip_bfloat16 a) {
  uint16_t bits;
  __builtin_memcpy(&bits, &a, sizeof(bits));
  bits &= 0x7FFFu;
  hip_bfloat16 result;
  __builtin_memcpy(&result, &bits, sizeof(result));
  return result;
}
__device__ __forceinline__ float16_t __habs(float16_t a) {
  uint16_t bits;
  __builtin_memcpy(&bits, &a, sizeof(bits));
  bits &= 0x7FFFu;
  float16_t result;
  __builtin_memcpy(&result, &bits, sizeof(result));
  return result;
}

namespace tl {

// Packed x2 element-wise math helpers (HIP scalar fallbacks)
//
// HIP does not expose packed-FP32x2 instructions, so we provide per-lane
// scalar fallbacks to keep the TileLang language surface portable.

TL_DEVICE float2 add2(float2 a, float2 b) {
  float2 out;
  out.x = a.x + b.x;
  out.y = a.y + b.y;
  return out;
}

TL_DEVICE float2 sub2(float2 a, float2 b) {
  float2 out;
  out.x = a.x - b.x;
  out.y = a.y - b.y;
  return out;
}

TL_DEVICE float2 mul2(float2 a, float2 b) {
  float2 out;
  out.x = a.x * b.x;
  out.y = a.y * b.y;
  return out;
}

TL_DEVICE float2 fma2(float2 a, float2 b, float2 c) {
  float2 out;
  out.x = a.x * b.x + c.x;
  out.y = a.y * b.y + c.y;
  return out;
}

TL_DEVICE float2 max2(float2 a, float2 b) {
  float2 out;
  out.x = (a.x > b.x) ? a.x : b.x;
  out.y = (a.y > b.y) ? a.y : b.y;
  return out;
}

TL_DEVICE float2 min2(float2 a, float2 b) {
  float2 out;
  out.x = (a.x < b.x) ? a.x : b.x;
  out.y = (a.y < b.y) ? a.y : b.y;
  return out;
}

TL_DEVICE float2 abs2(float2 a) {
  float2 out;
  out.x = (a.x >= 0.0f) ? a.x : -a.x;
  out.y = (a.y >= 0.0f) ? a.y : -a.y;
  return out;
}

// Packed bfloat16x2 overloads for uint1 carrier.
// On HIP, uint1 = HIP_vector_type<unsigned int, 1> (32-bit), already defined
// by ROCm via amd_hip_vector_types.h — no additional typedef needed.
// A packed bfloat16x2 word layout:
//   bits [15: 0] = first  bfloat16 (sign at bit 15)
//   bits [31:16] = second bfloat16 (sign at bit 31)
// These overloads are required by the HIP codegen's ShuffleNode packing path
// (VisitExpr_ ShuffleNode emits uint1{__pack_bfloat162(a, b)}).
TL_DEVICE uint1 abs2(uint1 val) {
  // Clear both sign bits simultaneously.
  return uint1{val.x & 0x7FFF7FFFu};
}
TL_DEVICE uint1 max2(uint1 a, uint1 b) {
  bfloat16_t a0, a1, b0, b1;
  __builtin_memcpy(&a0, &a.x, sizeof(a0));
  __builtin_memcpy(&a1, (char *)&a.x + 2, sizeof(a1));
  __builtin_memcpy(&b0, &b.x, sizeof(b0));
  __builtin_memcpy(&b1, (char *)&b.x + 2, sizeof(b1));
  bfloat16_t r0 = (float)a0 > (float)b0 ? a0 : b0;
  bfloat16_t r1 = (float)a1 > (float)b1 ? a1 : b1;
  return uint1{__pack_bfloat162(r0, r1)};
}
TL_DEVICE uint1 min2(uint1 a, uint1 b) {
  bfloat16_t a0, a1, b0, b1;
  __builtin_memcpy(&a0, &a.x, sizeof(a0));
  __builtin_memcpy(&a1, (char *)&a.x + 2, sizeof(a1));
  __builtin_memcpy(&b0, &b.x, sizeof(b0));
  __builtin_memcpy(&b1, (char *)&b.x + 2, sizeof(b1));
  bfloat16_t r0 = (float)a0 < (float)b0 ? a0 : b0;
  bfloat16_t r1 = (float)a1 < (float)b1 ? a1 : b1;
  return uint1{__pack_bfloat162(r0, r1)};
}
TL_DEVICE uint1 add2(uint1 a, uint1 b) {
  bfloat16_t a0, a1, b0, b1;
  __builtin_memcpy(&a0, &a.x, sizeof(a0));
  __builtin_memcpy(&a1, (char *)&a.x + 2, sizeof(a1));
  __builtin_memcpy(&b0, &b.x, sizeof(b0));
  __builtin_memcpy(&b1, (char *)&b.x + 2, sizeof(b1));
  bfloat16_t r0 = (bfloat16_t)((float)a0 + (float)b0);
  bfloat16_t r1 = (bfloat16_t)((float)a1 + (float)b1);
  return uint1{__pack_bfloat162(r0, r1)};
}
TL_DEVICE uint1 mul2(uint1 a, uint1 b) {
  bfloat16_t a0, a1, b0, b1;
  __builtin_memcpy(&a0, &a.x, sizeof(a0));
  __builtin_memcpy(&a1, (char *)&a.x + 2, sizeof(a1));
  __builtin_memcpy(&b0, &b.x, sizeof(b0));
  __builtin_memcpy(&b1, (char *)&b.x + 2, sizeof(b1));
  bfloat16_t r0 = (bfloat16_t)((float)a0 * (float)b0);
  bfloat16_t r1 = (bfloat16_t)((float)a1 * (float)b1);
  return uint1{__pack_bfloat162(r0, r1)};
}

// Any
template <typename T> TL_DEVICE bool Any(T *a, int size) {
  for (int i = 0; i < size; i++) {
    if (a[i]) {
      return true;
    }
  }
  return false;
}

// All
template <typename T> TL_DEVICE bool All(T *a, int size) {
  for (int i = 0; i < size; i++) {
    if (!a[i]) {
      return false;
    }
  }
  return true;
}

// TODO(gong): support shfl_sync(rocm 7.1.1 provide shfl_sync)
// shfl_sync func
template <typename T> TL_DEVICE T shfl_xor(T val, int delta) {
  return __shfl_xor(val, delta);
}

template <typename T> TL_DEVICE T shfl_down(T val, int delta) {
  return __shfl_down(val, delta);
}

template <typename T> TL_DEVICE T shfl_up(T val, int delta) {
  return __shfl_up(val, delta);
}

template <typename T> TL_DEVICE T shfl(T val, int srcLane) {
  return __shfl(val, srcLane);
}

// specialize half_t
template <> TL_DEVICE half_t shfl_xor(half_t val, int delta) {
  float f = static_cast<float>(val);
  float r = __shfl_xor(f, delta);
  return half_t(r);
}

template <> TL_DEVICE half_t shfl_down(half_t val, int delta) {
  float f = static_cast<float>(val);
  float r = __shfl_down(f, delta);
  return half_t(r);
}

template <> TL_DEVICE half_t shfl_up(half_t val, int delta) {
  float f = static_cast<float>(val);
  float r = __shfl_up(f, delta);
  return half_t(r);
}

template <> TL_DEVICE half_t shfl(half_t val, int srcLane) {
  float f = static_cast<float>(val);
  float r = __shfl(f, srcLane);
  return half_t(r);
}

// specialize bfloat16_t
template <> TL_DEVICE bfloat16_t shfl_xor(bfloat16_t val, int laneMask) {
  float f = static_cast<float>(val);
  float r = __shfl_xor(f, laneMask);
  return bfloat16_t(r);
}

template <> TL_DEVICE bfloat16_t shfl_down(bfloat16_t val, int delta) {
  float f = static_cast<float>(val);
  float r = __shfl_down(f, delta);
  return bfloat16_t(r);
}

template <> TL_DEVICE bfloat16_t shfl_up(bfloat16_t val, int delta) {
  float f = static_cast<float>(val);
  float r = __shfl_up(f, delta);
  return bfloat16_t(r);
}

template <> TL_DEVICE bfloat16_t shfl(bfloat16_t val, int srcLane) {
  float f = static_cast<float>(val);
  float r = __shfl(f, srcLane);
  return bfloat16_t(r);
}

} // namespace tl
