#include "LumeLib.h"

#include <llvm/Support/CommandLine.h>
#include <llvm/Support/InitLLVM.h>
#include <llvm/Support/MemoryBuffer.h>
#include <llvm/Support/ToolOutputFile.h>

#include <mlir/Dialect/LLVMIR/LLVMDialect.h>
#include <mlir/Dialect/SPIRV/IR/SPIRVDialect.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/MLIRContext.h>
#include <mlir/InitAllTranslations.h>
#include <mlir/Support/FileUtilities.h>
#include <mlir/Target/SPIRV/Deserialization.h>

#include <memory>
#include <string>

namespace {
llvm::cl::opt<std::string> inputFilename(
    llvm::cl::Positional,
    llvm::cl::desc("<input file>"),
    llvm::cl::init("-"));

llvm::cl::opt<std::string> outputFilename(
    "o",
    llvm::cl::desc("Output file (default: stdout)"),
    llvm::cl::value_desc("filename"),
    llvm::cl::init("-"));

llvm::cl::opt<bool> dumpMLIR(
    "dump-mlir",
    llvm::cl::desc("Dump the SPIR-V MLIR dialect"),
    llvm::cl::init(false));

llvm::cl::opt<bool> dumpAll(
    "dump-all",
    llvm::cl::desc("Dump all dialects and LLVM IR"),
    llvm::cl::init(false));
}

int main(int argc, char **argv) {
  llvm::InitLLVM initLLVM(argc, argv);
  llvm::cl::ParseCommandLineOptions(argc, argv, "Lume Compiler - SPIR-V to LLVM");

  mlir::registerFromSPIRVTranslation();

  std::string errorMessage;
  auto input = mlir::openInputFile(inputFilename, &errorMessage);
  if (!input) {
    llvm::errs() << "Error: " << errorMessage << '\n';
    return 1;
  }

  mlir::MLIRContext context;
  context.loadDialect<mlir::spirv::SPIRVDialect>();
  context.loadDialect<mlir::LLVM::LLVMDialect>();

  mlir::spirv::DeserializationOptions options;
  auto spvModule = lume::CompileSPVToMLIR(std::move(input), context, options);
  
  if (!spvModule) {
    llvm::errs() << "Error: Failed to deserialize SPIR-V binary\n";
    return 1;
  }

  mlir::OpBuilder builder(&context);
  auto topModule = mlir::ModuleOp::create(builder.getUnknownLoc());
  topModule.getBody()->push_back(spvModule.release());

  if (dumpMLIR || dumpAll) {
    llvm::errs() << "=== MLIR SPIR-V Dialect ===\n";
    topModule.dump();
    llvm::errs() << "\n";
  }

  if (llvm::failed(lume::ConvertSPVDialectToLLVMDialect(context, &topModule))) {
    llvm::errs() << "Error: Failed to convert SPIR-V dialect to LLVM dialect\n";
    return 1;
  }

  if (dumpMLIR || dumpAll) {
    llvm::errs() << "=== MLIR LLVM Dialect ===\n";
    topModule.dump();
    llvm::errs() << "\n";
  }

  if (outputFilename != "-") {
    auto output = mlir::openOutputFile(outputFilename, &errorMessage);
    if (!output) {
      llvm::errs() << "Error: " << errorMessage << '\n';
      return 1;
    }
    
    topModule.print(output->os());
    output->keep();
  } else if (!dumpMLIR && !dumpAll) {
    topModule.print(llvm::outs());
  }

  return 0;
}