#include "Lexer.h"
#include "Parser.h"

#include <llvm/Support/CommandLine.h>
#include <llvm/Support/ErrorOr.h>
#include <llvm/Support/MemoryBuffer.h>
#include <llvm/Support/SourceMgr.h>
#include <mlir/IR/BuiltinOps.h.inc>
#include <mlir/IR/BuiltinOps.h.inc>
#include <mlir/IR/OwningOpRef.h>
#include <mlir/Parser/Parser.h>

#include "Dialect.h"
#include "MLIRGen.h"

namespace mlir
{
    class ModuleOp;
}

static llvm::cl::opt<std::string> inputFilename(llvm::cl::Positional,
                                                llvm::cl::desc("<input hello file>"),
                                                llvm::cl::init("-"),
                                                llvm::cl::value_desc("filename"));

namespace
{
    enum Action { None, DumpSyntaxNode, DumpMLIR };

    enum InputType { Hello, MLIR };
}

static llvm::cl::opt<InputType> inputType("x", llvm::cl::init(Hello),
                                          llvm::cl::desc("Decided the kind of input desired."),
                                          llvm::cl::values(
                                              clEnumValN(Hello, "hello", "load the input file as a hello source.")),
                                          llvm::cl::values(
                                              clEnumValN(MLIR, "mlir", "load the input file as a mlir source.")));

static llvm::cl::opt<Action> emitAction("emit", llvm::cl::desc("Select the kind of output desired"),
                                        llvm::cl::values(clEnumValN(DumpSyntaxNode, "ast", "Dump syntax node")),
                                        llvm::cl::values(clEnumValN(DumpMLIR, "mlir", "Dump mlir code")));

std::unique_ptr<hello::Module> parseInputFile(llvm::StringRef filename)
{
    llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr =
        llvm::MemoryBuffer::getFileOrSTDIN(filename);
    if (std::error_code ec = fileOrErr.getError())
    {
        llvm::errs() << "Could not open input file: " << ec.message() << "\n";
        return nullptr;
    }
    auto buffer = fileOrErr.get()->getBuffer();
    hello::LexerBuffer lexer(buffer.begin(), buffer.end(), std::string(filename));
    hello::Parser parser(lexer);
    return parser.parseModule();
}

int dumpMLIR()
{
    mlir::MLIRContext context;
    context.getOrLoadDialect<mlir::hello::HelloDialect>();

    if (inputType != MLIR && !llvm::StringRef(inputFilename).ends_with(".mlir"))
    {
        auto module = parseInputFile(inputFilename);
        if (module == nullptr)
        {
            return 1;
        }
        mlir::OwningOpRef<mlir::ModuleOp> mlirModule = hello::mlirGen(context, *module);

        if (!mlirModule)
        {
            return 1;
        }

        mlirModule->dump();
        return 0;
    }

    // Then the input file is mlir
    llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> fileOrErr =
        llvm::MemoryBuffer::getFileOrSTDIN(inputFilename);
    if (std::error_code ec = fileOrErr.getError())
    {
        llvm::errs() << "Could not open input file: " << ec.message() << "\n";
        return 1;
    }

    // Parse the input mlir.
    llvm::SourceMgr sourceMgr;
    sourceMgr.AddNewSourceBuffer(std::move(*fileOrErr), llvm::SMLoc());
    mlir::OwningOpRef<mlir::ModuleOp> module =
        mlir::parseSourceFile<mlir::ModuleOp>(sourceMgr, &context);
    if (!module)
    {
        llvm::errs() << "Error can't load file " << inputFilename << "\n";
        return 1;
    }

    module->dump();
    return 0;
}

int main(int argc, char** argv)
{
    llvm::cl::ParseCommandLineOptions(argc, argv, "Hello MLIR Compiler\n");

    auto module = parseInputFile(inputFilename);

    if (!module)
    {
        return 1;
    }

    switch (emitAction)
    {
    case DumpSyntaxNode:
        module->dump();
        return 0;
    case DumpMLIR:
        dumpMLIR();
        return 0;
    default:
        llvm::errs() << "Unrecognized action\n";
        return 1;
    }
}