import argparse
import os
import sys
from collections.abc import Callable
from typing import IO

from xdsl.dialects.builtin import ModuleOp
from xdsl.ir import Dialect, MLContext
from xdsl.parser import Parser
from xdsl.passes import ModulePass
from xdsl.utils.exceptions import ParseError
from xdsl.utils.lexer import Span


def get_all_dialects() -> dict[str, Callable[[], Dialect]]:
    """Returns all available dialects."""

    def get_accfg():
        from xdsl.dialects.accfg import ACCFG

        return ACCFG

    def get_affine():
        from xdsl.dialects.affine import Affine

        return Affine

    def get_aie():
        from xdsl.dialects.experimental.aie import AIE

        return AIE

    def get_air():
        from xdsl.dialects.experimental.air import AIR

        return AIR

    def get_arith():
        from xdsl.dialects.arith import Arith

        return Arith

    def get_bufferization():
        from xdsl.dialects.bufferization import Bufferization

        return Bufferization

    def get_builtin():
        from xdsl.dialects.builtin import Builtin

        return Builtin

    def get_cf():
        from xdsl.dialects.cf import Cf

        return Cf

    def get_cmath():
        from xdsl.dialects.cmath import CMath

        return CMath

    def get_comb():
        from xdsl.dialects.comb import Comb

        return Comb

    def get_csl():
        from xdsl.dialects.csl import CSL

        return CSL

    def get_dmp():
        from xdsl.dialects.experimental.dmp import DMP

        return DMP

    def get_fir():
        from xdsl.dialects.experimental.fir import FIR

        return FIR

    def get_fsm():
        from xdsl.dialects.fsm import FSM

        return FSM

    def get_func():
        from xdsl.dialects.func import Func

        return Func

    def get_gpu():
        from xdsl.dialects.gpu import GPU

        return GPU

    def get_hlfir():
        from xdsl.dialects.experimental.hlfir import HLFIR

        return HLFIR

    def get_hls():
        from xdsl.dialects.experimental.hls import HLS

        return HLS

    def get_hw():
        from xdsl.dialects.hw import HW

        return HW

    def get_linalg():
        from xdsl.dialects.linalg import Linalg

        return Linalg

    def get_irdl():
        from xdsl.dialects.irdl.irdl import IRDL

        return IRDL

    def get_llvm():
        from xdsl.dialects.llvm import LLVM

        return LLVM

    def get_ltl():
        from xdsl.dialects.ltl import LTL

        return LTL

    def get_math():
        from xdsl.dialects.experimental.math import Math

        return Math

    def get_memref():
        from xdsl.dialects.memref import MemRef

        return MemRef

    def get_memref_stream():
        from xdsl.dialects.memref_stream import MemrefStream

        return MemrefStream

    def get_ml_program():
        from xdsl.dialects.ml_program import MLProgram

        return MLProgram

    def get_mpi():
        from xdsl.dialects.mpi import MPI

        return MPI

    def get_omp():
        from xdsl.dialects.omp import OMP

        return OMP

    def get_onnx():
        from xdsl.dialects.onnx import ONNX

        return ONNX

    def get_pdl():
        from xdsl.dialects.pdl import PDL

        return PDL

    def get_printf():
        from xdsl.dialects.printf import Printf

        return Printf

    def get_riscv_debug():
        from xdsl.dialects.riscv_debug import RISCV_Debug

        return RISCV_Debug

    def get_riscv():
        from xdsl.dialects.riscv import RISCV

        return RISCV

    def get_riscv_func():
        from xdsl.dialects.riscv_func import RISCV_Func

        return RISCV_Func

    def get_riscv_scf():
        from xdsl.dialects.riscv_scf import RISCV_Scf

        return RISCV_Scf

    def get_riscv_cf():
        from xdsl.dialects.riscv_cf import RISCV_Cf

        return RISCV_Cf

    def get_riscv_snitch():
        from xdsl.dialects.riscv_snitch import RISCV_Snitch

        return RISCV_Snitch

    def get_scf():
        from xdsl.dialects.scf import Scf

        return Scf

    def get_seq():
        from xdsl.dialects.seq import Seq

        return Seq

    def get_snitch():
        from xdsl.dialects.snitch import Snitch

        return Snitch

    def get_snitch_runtime():
        from xdsl.dialects.snitch_runtime import SnitchRuntime

        return SnitchRuntime

    def get_snitch_stream():
        from xdsl.dialects.snitch_stream import SnitchStream

        return SnitchStream

    def get_stencil():
        from xdsl.dialects.stencil import Stencil

        return Stencil

    def get_stream():
        from xdsl.dialects.stream import Stream

        return Stream

    def get_symref():
        from xdsl.frontend.symref import Symref

        return Symref

    def get_tensor():
        from xdsl.dialects.tensor import Tensor

        return Tensor

    def get_test():
        from xdsl.dialects.test import Test

        return Test

    def get_vector():
        from xdsl.dialects.vector import Vector

        return Vector

    def get_x86():
        from xdsl.dialects.x86 import X86

        return X86

    return {
        "accfg": get_accfg,
        "affine": get_affine,
        "aie": get_aie,
        "air": get_air,
        "arith": get_arith,
        "bufferization": get_bufferization,
        "builtin": get_builtin,
        "cf": get_cf,
        "cmath": get_cmath,
        "comb": get_comb,
        "csl": get_csl,
        "dmp": get_dmp,
        "fir": get_fir,
        "fsm": get_fsm,
        "func": get_func,
        "gpu": get_gpu,
        "hlfir": get_hlfir,
        "hls": get_hls,
        "hw": get_hw,
        "linalg": get_linalg,
        "irdl": get_irdl,
        "llvm": get_llvm,
        "ltl": get_ltl,
        "math": get_math,
        "memref": get_memref,
        "memref_stream": get_memref_stream,
        "ml_program": get_ml_program,
        "mpi": get_mpi,
        "omp": get_omp,
        "onnx": get_onnx,
        "pdl": get_pdl,
        "printf": get_printf,
        "riscv": get_riscv,
        "riscv_debug": get_riscv_debug,
        "riscv_func": get_riscv_func,
        "riscv_scf": get_riscv_scf,
        "riscv_cf": get_riscv_cf,
        "riscv_snitch": get_riscv_snitch,
        "scf": get_scf,
        "seq": get_seq,
        "snitch": get_snitch,
        "snrt": get_snitch_runtime,
        "snitch_stream": get_snitch_stream,
        "stencil": get_stencil,
        "stream": get_stream,
        "symref": get_symref,
        "tensor": get_tensor,
        "test": get_test,
        "vector": get_vector,
        "x86": get_x86,
    }


def get_all_passes() -> dict[str, Callable[[], type[ModulePass]]]:
    """Return the list of all available passes."""

    def get_arith_add_fastmath():
        from xdsl.transforms import arith_add_fastmath

        return arith_add_fastmath.AddArithFastMathFlagsPass

    def get_loop_hoist_memref():
        from xdsl.transforms import loop_hoist_memref

        return loop_hoist_memref.LoopHoistMemrefPass

    def get_canonicalize():
        from xdsl.transforms import canonicalize

        return canonicalize.CanonicalizePass

    def get_canonicalize_dmp():
        from xdsl.transforms import canonicalize_dmp

        return canonicalize_dmp.CanonicalizeDmpPass

    def get_convert_linalg_to_memref_stream():
        from xdsl.transforms import convert_linalg_to_memref_stream

        return convert_linalg_to_memref_stream.ConvertLinalgToMemrefStreamPass

    def get_convert_linalg_to_loops():
        from xdsl.transforms import convert_linalg_to_loops

        return convert_linalg_to_loops.ConvertLinalgToLoopsPass

    def get_convert_ml_program_to_memref():
        from xdsl.transforms import convert_ml_program_to_memref

        return convert_ml_program_to_memref.ConvertMlProgramToMemrefPass

    def get_convert_riscv_scf_for_to_frep():
        from xdsl.transforms import convert_riscv_scf_for_to_frep

        return convert_riscv_scf_for_to_frep.ConvertRiscvScfForToFrepPass

    def get_convert_scf_to_openmp():
        from xdsl.transforms import convert_scf_to_openmp

        return convert_scf_to_openmp.ConvertScfToOpenMPPass

    def get_convert_snitch_stream_to_snitch():
        from xdsl.backend.riscv.lowering import convert_snitch_stream_to_snitch

        return convert_snitch_stream_to_snitch.ConvertSnitchStreamToSnitch

    def get_constant_fold_interp():
        from xdsl.transforms import constant_fold_interp

        return constant_fold_interp.ConstantFoldInterpPass

    def get_convert_snrt_to_riscv():
        from xdsl.transforms import convert_snrt_to_riscv

        return convert_snrt_to_riscv.ConvertSnrtToRISCV

    def get_convert_stencil_to_ll_mlir():
        from xdsl.transforms.experimental import convert_stencil_to_ll_mlir

        return convert_stencil_to_ll_mlir.ConvertStencilToLLMLIRPass

    def get_convert_riscv_scf_to_riscv_cf():
        from xdsl.backend.riscv.lowering import convert_riscv_scf_to_riscv_cf

        return convert_riscv_scf_to_riscv_cf.ConvertRiscvScfToRiscvCfPass

    def get_cse():
        from xdsl.transforms import common_subexpression_elimination

        return common_subexpression_elimination.CommonSubexpressionElimination

    def get_dce():
        from xdsl.transforms import dead_code_elimination

        return dead_code_elimination.DeadCodeElimination

    def get_desymrefy():
        from xdsl.frontend.passes.desymref import DesymrefyPass

        return DesymrefyPass

    def get_gpu_allocs():
        from xdsl.transforms import gpu_allocs

        return gpu_allocs.MemrefToGPUPass

    def get_gpu_map_parallel_loops():
        from xdsl.transforms import gpu_map_parallel_loops

        return gpu_map_parallel_loops.GpuMapParallelLoopsPass

    def get_distribute_stencil():
        from xdsl.transforms.experimental.dmp import stencil_global_to_local

        return stencil_global_to_local.DistributeStencilPass

    def get_lower_halo_to_mpi():
        from xdsl.transforms.experimental.dmp import stencil_global_to_local

        return stencil_global_to_local.LowerHaloToMPI

    def get_individual_rewrite():
        from xdsl.transforms.individual_rewrite import IndividualRewrite

        return IndividualRewrite

    def get_lower_affine():
        from xdsl.transforms import lower_affine

        return lower_affine.LowerAffinePass

    def get_lower_mpi():
        from xdsl.transforms import lower_mpi

        return lower_mpi.LowerMPIPass

    def get_lower_riscv_func():
        from xdsl.transforms import lower_riscv_func

        return lower_riscv_func.LowerRISCVFunc

    def get_lower_snitch():
        from xdsl.transforms import lower_snitch

        return lower_snitch.LowerSnitchPass

    def get_memref_streamify():
        from xdsl.transforms import memref_streamify

        return memref_streamify.MemrefStreamifyPass

    def get_memref_stream_unnest_out_parameters():
        from xdsl.transforms import memref_stream_unnest_out_parameters

        return memref_stream_unnest_out_parameters.MemrefStreamUnnestOutParametersPass

    def get_mlir_opt():
        from xdsl.transforms import mlir_opt

        return mlir_opt.MLIROptPass

    def get_printf_to_llvm():
        from xdsl.transforms import printf_to_llvm

        return printf_to_llvm.PrintfToLLVM

    def get_printf_to_putchar():
        from xdsl.transforms import printf_to_putchar

        return printf_to_putchar.PrintfToPutcharPass

    def get_riscv_cse():
        from xdsl.transforms import riscv_cse

        return riscv_cse.RiscvCommonSubexpressionElimination

    def get_riscv_register_allocation():
        from xdsl.transforms import riscv_register_allocation

        return riscv_register_allocation.RISCVRegisterAllocation

    def get_scf_for_loop_flatten():
        from xdsl.transforms import scf_for_loop_flatten

        return scf_for_loop_flatten.ScfForLoopFlattenPass

    def get_riscv_scf_loop_range_folding():
        from xdsl.transforms import riscv_scf_loop_range_folding

        return riscv_scf_loop_range_folding.RiscvScfLoopRangeFoldingPass

    def get_snitch_register_allocation():
        from xdsl.transforms import snitch_register_allocation

        return snitch_register_allocation.SnitchRegisterAllocation

    def get_convert_arith_to_riscv():
        from xdsl.backend.riscv.lowering import convert_arith_to_riscv

        return convert_arith_to_riscv.ConvertArithToRiscvPass

    def get_convert_func_to_riscv_func():
        from xdsl.backend.riscv.lowering import convert_func_to_riscv_func

        return convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass

    def get_convert_memref_to_riscv():
        from xdsl.backend.riscv.lowering import convert_memref_to_riscv

        return convert_memref_to_riscv.ConvertMemrefToRiscvPass

    def get_convert_memref_stream_to_loops():
        from xdsl.transforms import convert_memref_stream_to_loops

        return convert_memref_stream_to_loops.ConvertMemrefStreamToLoopsPass

    def get_convert_onnx_to_linalg():
        from xdsl.transforms import convert_onnx_to_linalg

        return convert_onnx_to_linalg.ConvertOnnxToLinalgPass

    def get_convert_memref_stream_to_snitch():
        from xdsl.transforms import convert_memref_stream_to_snitch_stream

        return convert_memref_stream_to_snitch_stream.ConvertMemrefStreamToSnitch

    def get_convert_print_format_to_riscv_debug():
        from xdsl.backend.riscv.lowering import convert_print_format_to_riscv_debug

        return convert_print_format_to_riscv_debug.ConvertPrintFormatToRiscvDebugPass

    def get_scf_parallel_loop_tiling():
        from xdsl.transforms import scf_parallel_loop_tiling

        return scf_parallel_loop_tiling.ScfParallelLoopTilingPass

    def get_convert_scf_to_riscv_scf():
        from xdsl.backend.riscv.lowering import convert_scf_to_riscv_scf

        return convert_scf_to_riscv_scf.ConvertScfToRiscvPass

    def get_function_constant_pinning():
        from xdsl.transforms.experimental.function_constant_pinning import (
            FunctionConstantPinningPass,
        )

        return FunctionConstantPinningPass

    def get_lower_scf_for_to_labels():
        from xdsl.backend.riscv import riscv_scf_to_asm

        return riscv_scf_to_asm.LowerScfForToLabels

    def get_stencil_shape_inference():
        from xdsl.transforms.experimental import stencil_shape_inference

        return stencil_shape_inference.StencilShapeInferencePass

    def get_stencil_storage_materialization():
        from xdsl.transforms.experimental import stencil_storage_materialization

        return stencil_storage_materialization.StencilStorageMaterializationPass

    def get_reconcile_unrealized_casts():
        from xdsl.transforms import reconcile_unrealized_casts

        return reconcile_unrealized_casts.ReconcileUnrealizedCastsPass

    def get_hls_convert_stencil_to_ll_mlir():
        from xdsl.transforms.experimental import hls_convert_stencil_to_ll_mlir

        return hls_convert_stencil_to_ll_mlir.HLSConvertStencilToLLMLIRPass

    def get_lower_hls():
        from xdsl.transforms.experimental import lower_hls

        return lower_hls.LowerHLSPass

    def get_replace_incompatible_fpga():
        from xdsl.transforms.experimental import replace_incompatible_fpga

        return replace_incompatible_fpga.ReplaceIncompatibleFPGA

    def get_stencil_tensorize_z_dimension():
        from xdsl.transforms.experimental import stencil_tensorize_z_dimension

        return stencil_tensorize_z_dimension.StencilTensorizeZDimension

    def get_stencil_unroll():
        from xdsl.transforms import stencil_unroll

        return stencil_unroll.StencilUnrollPass

    def get_test_lower_snitch_stream_to_asm():
        from xdsl.transforms import test_lower_snitch_stream_to_asm

        return test_lower_snitch_stream_to_asm.TestLowerSnitchStreamToAsm

    return {
        "arith-add-fastmath": get_arith_add_fastmath,
        "loop-hoist-memref": get_loop_hoist_memref,
        "canonicalize-dmp": get_canonicalize_dmp,
        "canonicalize": get_canonicalize,
        "constant-fold-interp": get_constant_fold_interp,
        "convert-arith-to-riscv": get_convert_arith_to_riscv,
        "convert-func-to-riscv-func": get_convert_func_to_riscv_func,
        "convert-linalg-to-memref-stream": get_convert_linalg_to_memref_stream,
        "convert-linalg-to-loops": get_convert_linalg_to_loops,
        "convert-memref-stream-to-loops": get_convert_memref_stream_to_loops,
        "convert-memref-to-riscv": get_convert_memref_to_riscv,
        "convert-ml-program-to-memref": get_convert_ml_program_to_memref,
        "convert-onnx-to-linalg": get_convert_onnx_to_linalg,
        "convert-memref-stream-to-snitch": get_convert_memref_stream_to_snitch,
        "convert-print-format-to-riscv-debug": get_convert_print_format_to_riscv_debug,
        "convert-riscv-scf-for-to-frep": get_convert_riscv_scf_for_to_frep,
        "convert-riscv-scf-to-riscv-cf": get_convert_riscv_scf_to_riscv_cf,
        "convert-scf-to-openmp": get_convert_scf_to_openmp,
        "convert-scf-to-riscv-scf": get_convert_scf_to_riscv_scf,
        "convert-snitch-stream-to-snitch": get_convert_snitch_stream_to_snitch,
        "convert-snrt-to-riscv": get_convert_snrt_to_riscv,
        "convert-stencil-to-ll-mlir": get_convert_stencil_to_ll_mlir,
        "cse": get_cse,
        "dce": get_dce,
        "distribute-stencil": get_distribute_stencil,
        "dmp-to-mpi": get_lower_halo_to_mpi,
        "frontend-desymrefy": get_desymrefy,
        "function-constant-pinning": get_function_constant_pinning,
        "memref-to-gpu": get_gpu_allocs,
        "gpu-map-parallel-loops": get_gpu_map_parallel_loops,
        "hls-convert-stencil-to-ll-mlir": get_hls_convert_stencil_to_ll_mlir,
        "apply-individual-rewrite": get_individual_rewrite,
        "lower-affine": get_lower_affine,
        "lower-hls": get_lower_hls,
        "lower-mpi": get_lower_mpi,
        "lower-riscv-func": get_lower_riscv_func,
        "lower-riscv-scf-to-labels": get_lower_scf_for_to_labels,
        "lower-snitch": get_lower_snitch,
        "memref-streamify": get_memref_streamify,
        "memref-stream-unnest-out-parameters": get_memref_stream_unnest_out_parameters,
        "mlir-opt": get_mlir_opt,
        "printf-to-llvm": get_printf_to_llvm,
        "printf-to-putchar": get_printf_to_putchar,
        "reconcile-unrealized-casts": get_reconcile_unrealized_casts,
        "replace-incompatible-fpga": get_replace_incompatible_fpga,
        "riscv-allocate-registers": get_riscv_register_allocation,
        "riscv-cse": get_riscv_cse,
        "scf-for-loop-flatten": get_scf_for_loop_flatten,
        "riscv-scf-loop-range-folding": get_riscv_scf_loop_range_folding,
        "scf-parallel-loop-tiling": get_scf_parallel_loop_tiling,
        "snitch-allocate-registers": get_snitch_register_allocation,
        "stencil-shape-inference": get_stencil_shape_inference,
        "stencil-storage-materialization": get_stencil_storage_materialization,
        "stencil-tensorize-z-dimension": get_stencil_tensorize_z_dimension,
        "stencil-unroll": get_stencil_unroll,
        "test-lower-snitch-stream-to-asm": get_test_lower_snitch_stream_to_asm,
    }


class CommandLineTool:
    ctx: MLContext
    args: argparse.Namespace
    """
    The argument parsers namespace which holds the parsed commandline
    attributes.
    """

    available_frontends: dict[str, Callable[[IO[str]], ModuleOp]]
    """
    A mapping from file extension to a frontend that can handle this
    file type.
    """

    def register_all_arguments(self, arg_parser: argparse.ArgumentParser):
        arg_parser.add_argument(
            "input_file", type=str, nargs="?", help="path to input file"
        )

        frontends = [name for name in self.available_frontends]
        arg_parser.add_argument(
            "-f",
            "--frontend",
            type=str,
            required=False,
            choices=frontends,
            help="Frontend to be used for the input. If not set, "
            "the xdsl frontend or the one for the file extension "
            "is used.",
        )
        arg_parser.add_argument("--disable-verify", default=False, action="store_true")

        arg_parser.add_argument(
            "--allow-unregistered-dialect",
            default=False,
            action="store_true",
            help="Allow the parsing of unregistered dialects.",
        )

        arg_parser.add_argument(
            "--no-implicit-module",
            default=False,
            action="store_true",
            help="Disable implicit addition of a top-level module op during parsing.",
        )

    def get_input_stream(self) -> tuple[IO[str], str]:
        """
        Get the input stream to parse from, along with the file extension.
        """
        if self.args.input_file is None:
            f = sys.stdin
            file_extension = "mlir"
        else:
            f = open(self.args.input_file)
            _, file_extension = os.path.splitext(self.args.input_file)
            file_extension = file_extension.replace(".", "")
        return f, file_extension

    def get_input_name(self):
        return self.args.input_file or "stdin"

    def register_all_dialects(self):
        """
        Register all dialects that can be used.

        Add other/additional dialects by overloading this function.
        """
        for dialect_name, dialect_factory in get_all_dialects().items():
            self.ctx.register_dialect(dialect_name, dialect_factory)

    def register_all_frontends(self):
        """
        Register all frontends that can be used.

        Add other/additional frontends by overloading this function.
        """

        def parse_mlir(io: IO[str]):
            return Parser(
                self.ctx,
                io.read(),
                self.get_input_name(),
            ).parse_module(not self.args.no_implicit_module)

        self.available_frontends["mlir"] = parse_mlir

    def parse_chunk(
        self, chunk: IO[str], file_extension: str, start_offset: int = 0
    ) -> ModuleOp | None:
        """
        Parse the input file by invoking the parser specified by the `parser`
        argument. If not set, the parser registered for this file extension
        is used.
        """

        try:
            return self.available_frontends[file_extension](chunk)
        except ParseError as e:
            s = e.span
            e.span = Span(s.start, s.end, s.input, start_offset)
            if "parsing_diagnostics" in self.args and self.args.parsing_diagnostics:
                print(e.with_context())
            else:
                raise Exception("Failed to parse:\n" + e.with_context()) from e
        finally:
            chunk.close()
