#include "sm3_registers.h"

#include "sm3_converter.h"

#include "../ir/ir_utils.h"

namespace dxbc_spv::sm3 {

RegisterFile::RegisterFile(Converter &converter)
: m_converter(converter) {
}


RegisterFile::~RegisterFile() {

}


void RegisterFile::initialize(ir::Builder& builder) {
  uint32_t addressRegisterComponents = m_converter.getShaderInfo().getVersion().first >= 2u ? 4u : 1u;

  for (uint32_t i = 0u; i < addressRegisterComponents; i++) {
    m_a0Reg[i] = builder.add(ir::Op::DclTmp(ir::ScalarType::eI32, m_converter.getEntryPoint()));

    if (m_converter.getOptions().includeDebugNames) {
      std::string name = m_converter.makeRegisterDebugName(RegisterType::eAddr, 0u, util::componentBit(Component(i)));
      builder.add(ir::Op::DebugName(m_a0Reg[i], name.c_str()));
    }
  }

  m_aLReg = builder.add(ir::Op::DclTmp(ir::ScalarType::eI32, m_converter.getEntryPoint()));

  if (m_converter.getOptions().includeDebugNames) {
    std::string name = m_converter.makeRegisterDebugName(RegisterType::eLoop, 0u, ComponentBit::eAll);
    builder.add(ir::Op::DebugName(m_aLReg, name.c_str()));
  }

  for (uint32_t i = 0u; i < 4u; i++) {
    m_pReg[i] = builder.add(ir::Op::DclTmp(ir::ScalarType::eBool, m_converter.getEntryPoint()));

    if (m_converter.getOptions().includeDebugNames) {
      std::string name = m_converter.makeRegisterDebugName(RegisterType::ePredicate, 0u, util::componentBit(Component(i)));
      builder.add(ir::Op::DebugName(m_pReg[i], name.c_str()));
    }
  }
}


ir::SsaDef RegisterFile::getOrDeclareTemp(ir::Builder& builder, uint32_t index, Component component) {
  uint32_t tempIndex = 4u * index + uint8_t(component);

  if (tempIndex >= m_rRegs.size())
    m_rRegs.resize(tempIndex + 1);

  if (!m_rRegs[tempIndex]) {
    m_rRegs[tempIndex] = builder.add(ir::Op::DclTmp(ir::ScalarType::eUnknown, m_converter.getEntryPoint()));

    if (m_converter.getOptions().includeDebugNames) {
      std::string name = m_converter.makeRegisterDebugName(RegisterType::eTemp, index, util::componentBit(component));
      builder.add(ir::Op::DebugName(m_rRegs[tempIndex], name.c_str()));
    }
  }

  return m_rRegs[tempIndex];
}


ir::SsaDef RegisterFile::getOrDeclareTextureReg(ir::Builder& builder, uint32_t index, Component component) {
  uint32_t textureIndex = 4u * index + uint8_t(component);

  if (!m_textureRegs[textureIndex]) {
    m_textureRegs[textureIndex] = builder.add(ir::Op::DclTmp(ir::ScalarType::eUnknown, m_converter.getEntryPoint()));

    if (m_converter.getOptions().includeDebugNames) {
      std::string name = m_converter.makeRegisterDebugName(RegisterType::eTexture, index, util::componentBit(component));
      builder.add(ir::Op::DebugName(m_textureRegs[textureIndex], name.c_str()));
    }
  }

  return m_textureRegs[textureIndex];
}


ir::SsaDef RegisterFile::emitTempLoad(
  ir::Builder& builder,
  uint32_t regIndex,
  Swizzle swizzle,
  WriteMask componentMask,
  ir::ScalarType type
) {
  auto returnType = makeVectorType(type, componentMask);

  std::array<ir::SsaDef, 4u> components = { };

  for (auto c : swizzle.getReadMask(componentMask)) {
    auto component = componentFromBit(c);

    auto tmpReg = getOrDeclareTemp(builder, regIndex, component);
    auto scalar = builder.add(ir::Op::TmpLoad(ir::ScalarType::eUnknown, tmpReg));

    /* Convert to requested type */
    scalar = builder.add(ir::Op::ConsumeAs(type, scalar));

    components[uint8_t(component)] = scalar;
  }

  return composite(builder, returnType, components.data(), swizzle, componentMask);
}


ir::SsaDef RegisterFile::emitTextureRegLoad(
    ir::Builder& builder,
    uint32_t regIndex,
    Swizzle swizzle,
    WriteMask componentMask,
    ir::ScalarType type
  ) {
  auto returnType = makeVectorType(type, componentMask);

  std::array<ir::SsaDef, 4u> components = { };

  for (auto c : swizzle.getReadMask(componentMask)) {
    auto component = componentFromBit(c);

    auto texReg = getOrDeclareTextureReg(builder, regIndex, component);
    auto scalar = builder.add(ir::Op::TmpLoad(ir::ScalarType::eUnknown, texReg));

    /* Convert to requested type */
    scalar = builder.add(ir::Op::ConsumeAs(type, scalar));

    components[uint8_t(component)] = scalar;
  }

  return composite(builder, returnType, components.data(), swizzle, componentMask);
}


ir::SsaDef RegisterFile::emitPredicateLoad(
            ir::Builder&            builder,
            Swizzle                 swizzle,
            WriteMask               componentMask) {
  auto returnType = makeVectorType(ir::ScalarType::eBool, componentMask);

  std::array<ir::SsaDef, 4u> components = { };
  for (auto c : swizzle.getReadMask(componentMask)) {
    auto component = componentFromBit(c);
    components[uint8_t(component)] = builder.add(ir::Op::TmpLoad(ir::ScalarType::eBool, m_pReg[uint8_t(component)]));
  }

  return composite(builder, returnType, components.data(), swizzle, componentMask);
}


ir::SsaDef RegisterFile::emitAddressLoad(
            ir::Builder&            builder,
            RegisterType            registerType,
            Swizzle                 swizzle,
            ir::SsaDef              loopCounter) {
  Component component = swizzle.get(0u);

  if (registerType == RegisterType::eAddr) {
    /* The a0 register only has a single component on SM1.1 */
    dxbc_spv_assert(m_converter.getShaderInfo().getVersion().first >= 2u
        || component == Component::eX);

    dxbc_spv_assert(m_a0Reg[uint8_t(component)]);

    return builder.add(ir::Op::TmpLoad(ir::ScalarType::eI32, m_a0Reg[uint8_t(component)]));
  } else if (registerType == RegisterType::eLoop) {
    dxbc_spv_assert(component == Component::eX);

    if (!loopCounter) {
      Logger::err("No valid loop counter found in current context.");
      return builder.makeUndef(ir::ScalarType::eI32);
    }

    return builder.add(ir::Op::TmpLoad(ir::ScalarType::eI32, loopCounter));
  } else {
    ShaderInfo info = m_converter.getShaderInfo();

    Logger::err("Register type: ",
      UnambiguousRegisterType { registerType, info.getType(), info.getVersion().first },
      " not supported as a relative index");

    dxbc_spv_unreachable();
    return ir::SsaDef();
  }
}


bool RegisterFile::emitStore(
          ir::Builder&            builder,
    const Operand&                operand,
          WriteMask               writeMask,
          ir::SsaDef              predicateVec,
          ir::SsaDef              value) {
  const auto& valueDef = builder.getOp(value);
  auto valueType = valueDef.getType().getBaseType(0u);

  auto regIndex = operand.getIndex();

  uint32_t componentIndex = 0u;

  for (auto c : writeMask) {
    auto component = componentFromBit(c);

    /* Extract scalar and 'convert' to unknown type */
    auto scalar = extractFromVector(builder, value, componentIndex);

    ir::SsaDef reg;

    switch (operand.getRegisterType()) {
      case RegisterType::eTemp: {
        reg = getOrDeclareTemp(builder, regIndex, component);

        if (!valueType.isUnknownType())
          scalar = builder.add(ir::Op::ConsumeAs(ir::ScalarType::eUnknown, scalar));
      } break;

      case RegisterType::eAddr:
      /* case RegisterType::eTexture: Same Value */ {

        /* Vertex shaders use this register type value for address */
        if (m_converter.getShaderInfo().getType() == ShaderType::eVertex) {
          dxbc_spv_assert(m_converter.getShaderInfo().getVersion().first >= 2u || c == ComponentBit::eX);

          if (!valueType.isIntType()) {
            /* a0 can be written to using the mova instruction on SM2+
             * or the mov instruction on SM1.1. */

            scalar = builder.add(ir::Op::ConsumeAs(ir::ScalarType::eF32, scalar));
            scalar = builder.add(ir::Op::FRound(ir::ScalarType::eF32, scalar, ir::RoundMode::eNearestEven));
            scalar = builder.add(ir::Op::ConvertFtoI(ir::ScalarType::eI32, scalar));
          }

          reg = m_a0Reg[uint8_t(component)];
        } else {
          /* Texture registers act as special purpose temp registers in PS 1.1 - PS 1.3. */
          reg = getOrDeclareTextureReg(builder, regIndex, component);

          if (!valueType.isUnknownType())
            scalar = builder.add(ir::Op::ConsumeAs(ir::ScalarType::eUnknown, scalar));
        }
      } break;

      case RegisterType::ePredicate: {
        dxbc_spv_assert(valueType.isBoolType());
        reg = m_pReg[uint8_t(component)];
      } break;

      default:
        dxbc_spv_unreachable();
        return false;
    }

    if (predicateVec) {
      /* Check if the matching component of the predicate register vector is true first.
       * Pick the old value if not. */
      auto condComponent = extractFromVector(builder, predicateVec, componentIndex);

      ir::BasicType regType = builder.getOp(reg).getType().getBaseType(0u);

      auto oldValue = builder.add(ir::Op::TmpLoad(regType, reg));
      scalar = builder.add(ir::Op::Select(regType, condComponent, scalar, oldValue));
    }

    m_stores.emplace_back(reg, scalar);

    componentIndex++;
  }

  return true;
}


void RegisterFile::emitBufferedStores(ir::Builder& builder) {
  for (const auto& store : m_stores)
    builder.add(ir::Op::TmpStore(store.reg, store.value));

  m_stores.clear();
}


ir::SsaDef RegisterFile::getOrDeclareLabel(
        ir::Builder&            builder,
  const Operand&                operand) {
  dxbc_spv_assert(operand.getRegisterType() == RegisterType::eLabel);

  auto index = operand.getIndex();

  if (index >= m_labels.size())
    m_labels.resize(index + 1u);

  auto& fn = m_labels[index];

  if (!fn) {
    /* Declare function right after declaration section and begin/end
     * it immediately. We will move the cursor once the corresponding
     * label instruction is encountered. */
    auto ref = builder.getCode().first->getDef();

    fn = builder.addBefore(ref, ir::Op::Function(ir::ScalarType::eVoid));
    builder.addAfter(fn, ir::Op::FunctionEnd());
    builder.add(ir::Op::DebugName(fn, (std::string("l") + std::to_string(index)).c_str()));
  }

  return fn;
}

}
