import hashlib
import re
from collections.abc import Iterable

from xdsl.dialects import arith, builtin, llvm, printf
from xdsl.ir import Attribute, MLContext, Operation, SSAValue
from xdsl.passes import ModulePass
from xdsl.pattern_rewriter import (
    PatternRewriter,
    PatternRewriteWalker,
    RewritePattern,
    op_type_rewrite_pattern,
)

i8 = builtin.IntegerType(8)


def legalize_str_for_symbol_name(val: str):
    """
    Takes any string and legalizes it to be a global llvm symbol.
    (for the strictest possible interpretation of this)

     - Replaces all whitespaces and dots with _
     - Deletes all non ascii alphanumerical characters
     - Strips all underscores from the start and end of the string

    The resulting string consists only of ascii letters, underscores and digits.

    This is a surjective mapping, meaning that multiple inputs will produce the same
    output. This function alone cannot be used to get a uniquely identifying global
    symbol name for a string!
    """
    val = re.sub(r"(\s+|\.)", "_", val)
    val = re.sub(r"[^A-Za-z0-9_]+", "", val).strip("_")
    return val


def _key_from_str(val: str) -> str:
    """
    Generate a symbol name from any given string.

    Takes the first ten letters of the string plus it's sha1 hash to create a
    (pretty much) globally unique symbol name.
    """
    h = hashlib.new("sha1")
    h.update(val.encode())
    return f"{legalize_str_for_symbol_name(val[:10])}_{h.hexdigest()}"


def _format_string_spec_from_print_op(
    op: printf.PrintFormatOp,
) -> Iterable[str | SSAValue]:
    """
    Translates the op:
    printf.print_format "val = {}, val2 = {}", %1 : i32, %2 : f32

    into this sequence:
    ["val = ", %1, ", val2 = ", %2]

    Empty string parts are omitted.
    """
    format_str = op.format_str.data.split("{}")
    args = iter(op.format_vals)

    for part in format_str[:-1]:
        if part != "":
            yield part
        yield next(args)
    if format_str[-1] != "":
        yield format_str[-1]


def _format_str_for_type(t: Attribute):
    match t:
        case builtin.f64:
            return "%f"
        case builtin.i32:
            return "%i"
        case builtin.i64:
            return "%li"
        case _:
            raise ValueError(f"Cannot find printf code for {t}")


class PrintlnOpToPrintfCall(RewritePattern):
    collected_global_symbs: dict[str, llvm.GlobalOp]

    def __init__(self):
        self.collected_global_symbs = dict()

    def _construct_global(self, val: str):
        """
        Constructs an llvm.global operation containing the string. Assigns a unique
        symbol name to the value that is derived from the string value.
        """
        data = val.encode() + b"\x00"

        t_type = builtin.TensorType(i8, [len(data)])

        return llvm.GlobalOp(
            llvm.LLVMArrayType.from_size_and_type(len(data), i8),
            _key_from_str(val),
            constant=True,
            linkage="internal",
            value=builtin.DenseIntOrFPElementsAttr.from_list(t_type, data),
        )

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: printf.PrintFormatOp, rewriter: PatternRewriter, /):
        format_str = ""
        args: list[SSAValue] = []
        casts: list[Operation] = []
        # make sure all arguments are in the format libc expects them to be
        # e.g. floats must be promoted to double before calling
        for part in _format_string_spec_from_print_op(op):
            if isinstance(part, str):
                format_str += part
            elif isinstance(part.type, builtin.IndexType):
                # index must be cast to fixed bitwidth before printing
                casts.append(new_val := arith.IndexCastOp(part, builtin.i64))
                args.append(new_val.result)
                format_str += "%li"
            elif part.type == builtin.f32:
                # f32 must be promoted to f64 before printing
                casts.append(new_val := arith.ExtFOp(part, builtin.f64))
                args.append(new_val.result)
                format_str += "%f"
            else:
                args.append(part)
                format_str += _format_str_for_type(part.type)

        globl = self._construct_global(format_str + "\n")
        self.collected_global_symbs[globl.sym_name.data] = globl

        rewriter.replace_matched_op(
            casts
            + [
                ptr := llvm.AddressOfOp(globl.sym_name, llvm.LLVMPointerType.opaque()),
                llvm.CallOp("printf", ptr.result, *args, variadic_args=len(args)),
            ]
        )


class PrintfToLLVM(ModulePass):
    name = "printf-to-llvm"

    def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
        add_printf_call = PrintlnOpToPrintfCall()

        PatternRewriteWalker(add_printf_call).rewrite_module(op)

        if not add_printf_call.collected_global_symbs:
            return

        op.body.block.add_ops(
            [
                llvm.FuncOp(
                    "printf",
                    llvm.LLVMFunctionType(
                        [llvm.LLVMPointerType.opaque()], is_variadic=True
                    ),
                    linkage=llvm.LinkageAttr("external"),
                ),
                *add_printf_call.collected_global_symbs.values(),
            ]
        )
