# 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.

alias _TargetType = __mlir_type.`!kgen.target`

@always_inline("nodebug")
fn _current_target() -> _TargetType:
  return __mlir_attr.`#kgen.param.expr<current_target> : !kgen.target`

@register_passable("trivial")
struct Target[value: _TargetType  = _current_target()]:
  @always_inline("nodebug")
  @staticmethod
  fn _get_field[name: StringLiteral, type: AnyTrivialRegType = StringLiteral]() -> type:
    return __mlir_attr[
      `#kgen.param.expr<target_get_field,`,
      Self.value,
      `, `,
      name.value,
      `> :`,
      type,
    ]

  @always_inline("nodebug")
  @staticmethod
  fn _has_feature[name: StringLiteral]() -> Bool:
    return __mlir_attr[
      `#kgen.param.expr<target_has_feature,`,
      Self.value,
      `,`,
      name.value,
      `> : i1`,
    ]

  @always_inline("nodebug")
  @staticmethod
  fn _size_of[mlir_type: AnyTrivialRegType]() -> Int:
    return __mlir_attr[
      `#kgen.param.expr<get_sizeof, #kgen.type<`,
      mlir_type,
      `> : !kgen.type,`,
      Self.value,
      `> : index`,
    ]

  @always_inline("nodebug")
  @staticmethod
  fn _align_of[mlir_type: AnyTrivialRegType]() -> Int:
    return __mlir_attr[
      `#kgen.param.expr<get_alignof, #kgen.type<`,
      mlir_type,
      `> : !kgen.type,`,
      Self.value,
      `> : index`,
    ]

  @always_inline("nodebug")
  @staticmethod
  fn os() -> StringLiteral:
    return Self._get_field["os"]()

  @always_inline("nodebug")
  @staticmethod
  fn triple() -> StringLiteral:
    return Self._get_field["triple"]()

  # Features

  @staticmethod
  fn arch() -> StringLiteral:
    return Self._get_field["arch"]()

  @staticmethod
  fn has_sse4() -> Bool:
    return Self._has_feature["sse4"]()

  @staticmethod
  fn has_avx() -> Bool:
    return Self._has_feature["avx"]()

  @staticmethod
  fn has_avx2() -> Bool:
    return Self._has_feature["avx2"]()

  @staticmethod
  fn has_avx512f() -> Bool:
    return Self._has_feature["avx512f"]()

  @staticmethod
  fn has_intel_amx() -> Bool:
    return Self._has_feature["amx-tile"]()

  @staticmethod
  fn has_fma() -> Bool:
    return Self._has_feature["fma"]()

  @staticmethod
  fn has_vnni() -> Bool:
    alias res = Self._has_feature["avx512vnni"]() or Self._has_feature["avxvnni"]()
    return res

  @staticmethod
  fn has_neon() -> Bool:
    return Self._has_feature["neon"]()

  @staticmethod
  fn has_neon_int8_dotprod() -> Bool:
    return Self._has_feature["dotprod"]()

  @staticmethod
  fn has_neon_int8_matmul() -> Bool:
    return Self._has_feature["i8mm"]()

  # Platforms

  @staticmethod
  fn is_x86() -> Bool:
    return Self.has_sse4()

  @staticmethod
  fn is_apple_silicon() -> Bool:
    return Self.is_apple_m1() or Self.is_apple_m2() or Self.is_apple_m3()

  @staticmethod
  fn is_apple_m1() -> Bool:
    return Self.arch() == "apple-m1"

  @staticmethod
  fn is_apple_m2() -> Bool:
    return Self.arch() == "apple-m2"

  @staticmethod
  fn is_apple_m3() -> Bool:
    return Self.arch() == "apple-m3"

  @staticmethod
  fn is_neoverse_n1() -> Bool:
    return Self.arch() == "neoverse-n1"

  # OS

  @staticmethod
  fn os_is_linux() -> Bool:
    return Self.os() == "linux"

  @staticmethod
  fn os_is_macos() -> Bool:
    return Self.os() in ("macos", "darwin")

  @staticmethod
  fn os_is_windows() -> Bool:
    return Self.os() == "windows"

  # Endianness

  @staticmethod
  fn is_big_endian() -> Bool:
    return Self._get_field["endianness"]() == "big"

  @staticmethod
  fn is_little_endian() -> Bool:
    return Self._get_field["endianness"]() == "little"

  # Pointer size

  @staticmethod
  fn is_32bit() -> Bool:
    return Self._get_field["pointer-size"]() == "32"

  @staticmethod
  fn is_64bit() -> Bool:
    return Self._get_field["pointer-size"]() == "64"

  @staticmethod
  fn size_of[type: AnyType]() -> Int:
    alias mlir_type = __mlir_attr[
        `#kgen.param.expr<rebind, #kgen.type<!kgen.param<`,
        type,
        `>> : `,
        AnyType,
        `> : !kgen.type`,
    ]
    return Self._size_of[mlir_type]()

  @staticmethod
  fn align_of[type: AnyType]() -> Int:
    alias mlir_type = __mlir_attr[
        `#kgen.param.expr<rebind, #kgen.type<!kgen.param<`,
        type,
        `>> : `,
        AnyType,
        `> : !kgen.type`,
    ]
    return Self._align_of[mlir_type]()

  # SIMD

  @staticmethod
  fn simd_bitwidth() -> Int:
    return Self._get_field["simd_bit_width", Int]()

  @staticmethod
  fn simd_bytewidth() -> Int:
    return Self.simd_bitwidth() // 8


from ir_utils import dump_ir

fn main():
  print(Target.simd_bitwidth())
  print(Target.size_of[Int]())