#include <willow/IR/BasicBlock.h>
#include <willow/IR/Diagnostic.h>
#include <willow/IR/DiagnosticEngine.h>
#include <willow/IR/Module.h>
#include <willow/IR/Verifier.h>

namespace willow {

/// Verify that an instruction defines an SSA result
LogicalResult verifyResult(const Instruction &inst, DiagnosticEngine &diags);

LogicalResult verifyNumOperands(const Instruction &inst,
                                DiagnosticEngine &diags, std::size_t expected);

LogicalResult verifyModule(const Module &module, DiagnosticEngine &diags) {
  bool any_failure = false;

  for (auto &func : module.getFunctions()) {
    std::vector<Diagnostic> collected;
    DiagnosticEngine eng(
        [&](Diagnostic d) { collected.push_back(std::move(d)); });

    auto r = verifyFunction(func, eng);

    if (succeeded(r))
      continue;

    any_failure = true;

    auto diag = emit(diags, Severity::Error, std::nullopt);
    diag << "verification failed for function: '" << func.getName() << "'";

    for (auto &d : collected)
      diag.note(std::move(d));
  }

  return any_failure ? failure() : success();
}

// LogicalResult verifyFunction(const Function&function, DiagnosticEngine
// &diags) {
//   bool any_failure = false;
// }

LogicalResult verifyBasicBlock(const BasicBlock &BB, DiagnosticEngine &diags) {
  bool any_failure = false;

  if (BB.empty())
    return emit(diags, Severity::Error, BB.getLoc())
           << "Basic block '" << BB.getName() << "' has an empty body";

  auto *trailer = BB.trailer();

  if (!trailer->isTerminator()) {
  }

  if (!BB.trailer()->isTerminator())
    // TODO: terminator check

    for (auto &inst : BB.getBody()) {
      // verify inst
    }
}

// TODO: better instruction printing
LogicalResult verifyInst(const Instruction &inst, DiagnosticEngine &diags) {
  using enum Instruction::Opcode;
  switch (inst.opcode()) {
  case Add:
  case Mul:
  case Sub:
  case Div:
  case Mod:
  case Shl:
  case Shr:
  case Ashl:
  case Ashr:
  case And:
  case Or:
    return verifyBinaryInst(inst, diags);
  case Eq:
  case Lt:
  case Gt:
  case Le:
  case Ge: {
    Type ty = inst.getType();
    if (!ty.isInt() || ty.getNumBits() != 1)
      return emit(diags, Severity::Error, inst.getLoc())
             << std::format("unexpected result type '{}': compare "
                            "instructions return i1",
                            ty);

    size_t num_operands = inst.getNumOperands();
    if (num_operands != 2)
      return emit(diags, Severity::Error, inst.getLoc())
             << std::format("expected 2 operands, found {}", num_operands);

    const Value *lhs = inst.getOperand(0);
    const Value *rhs = inst.getOperand(1);

    Type lty = lhs->getType(), rty = rhs->getType();

    if (!lty.isInt())
      return emit(diags, Severity::Error, inst.getLoc()) << std::format(
                 "invalid operand type '{}': expected integral type", lty);

    if (!rty.isInt())
      return emit(diags, Severity::Error, inst.getLoc()) << std::format(
                 "invalid operand type '{}': expected integral type", rty);

    if (lty != rty)
      return emit(diags, Severity::Error, inst.getLoc())
             << "mismatched operand types";
  }
  case Not: {
    Type ty = inst.getType();

    if (failed(verifyResult(inst, diags)))
      return failure();

    const Value *operand = inst.getOperand(0);
    if (!operand)
      return emit(diags, Severity::Error, inst.getLoc())
             << "instruction 'not' requires 1 operand";

    Type oty = operand->getType();
    if (ty != oty)
      return emit(diags, Severity::Error, inst.getLoc())
             << std::format("expected argument type '{}', got '{}'", ty, oty);
  }
  case Jmp:
  case Br:
  case Call:
  case Ret:
  case Phi:
  case Alloca:
  }
}

// TODO: better naming?
LogicalResult verifyBinaryInst(const Instruction &inst,
                               DiagnosticEngine &diags) {
  Type ty = inst.getType();

  // TODO non scalars
  if (!ty.isInt())
    return emit(diags, Severity::Error, inst.getLoc())
           << "invalid instruction '" << inst << "': "
           << "expected an integral type, got '" << ty << "'";

  auto *lhs = inst.getOperand(0);
  auto *rhs = inst.getOperand(1);

  assert(lhs && "Binary op lacks LHS");
  assert(rhs && "Binary op needs RHS");

  if (lhs->getType() != ty) {
    return emit(diags, Severity::Error, inst.getLoc()) << std::format(
               "expected operand type '{}' got '{}'", ty, lhs->getType());
  }

  if (rhs->getType() != ty) {
    return emit(diags, Severity::Error, inst.getLoc()) << std::format(
               "expected operand type '{}' got '{}'", ty, rhs->getType());
  }
}

LogicalResult verifyResult(const Instruction &inst, DiagnosticEngine &diags) {
  if (inst.hasName())
    return success();

  return emit(diags, Severity::Error, inst.getLoc()) << "expected ssa result";
}

LogicalResult verifyNumOperands(const Instruction &inst,
                                DiagnosticEngine &diags, std::size_t expected) {
  if (inst.getNumOperands() != expected)
    return emitE
}

} // namespace willow