# Load this module into LLDB by running:
#     command script import /path/to/Halide/tools/lldbhalide.py
import functools

import lldb


def normalize(raw):
    return raw.lstrip('"').rstrip('"').replace(r'\n', ' ').replace('  ', ' ')


def summary_string(summary_fn):
    @functools.wraps(summary_fn)
    def wrapper(value, _):
        if value is None or not value.IsValid():
            return "<invalid>"

        try:
            return normalize(summary_fn(value).GetSummary())
        except Exception as e:
            return f"<error:{e},{value}>"

    return wrapper


@summary_string
def call_name(value):
    return value.EvaluateExpression("this->name()", lldb.SBExpressionOptions())


@summary_string
def call_lldb_string(value):
    return value.EvaluateExpression(f"Halide::Internal::lldb_string(*this)", lldb.SBExpressionOptions())


class ProxyChildrenProvider:
    def __init__(self, valobj, _):
        self.inner = valobj
        self.update()

    def update(self):
        pass

    def num_children(self):
        return self.inner.GetNumChildren()

    def get_child_index(self, name):
        return self.inner.GetIndexOfChildWithName(name)

    def get_child_at_index(self, index):
        return self.inner.GetChildAtIndex(index)


class IRChildrenProvider(ProxyChildrenProvider):
    def __init__(self, valobj, _):
        super().__init__(valobj.GetChildMemberWithName("ptr"), None)


class BoxChildrenProvider(IRChildrenProvider):
    def __init__(self, valobj, _):
        super().__init__(valobj.GetChildMemberWithName("contents"), None)


class FunctionChildrenProvider(ProxyChildrenProvider):
    def __init__(self, valobj, _):
        contents = valobj.EvaluateExpression("*this->contents.get()", lldb.SBExpressionOptions())
        print(contents)
        super().__init__(contents, None)


def __lldb_init_module(debugger, _):
    base_exprs = ["Add", "And", "Broadcast", "Call", "Cast", "Div", "EQ", "GE", "GT", "LE", "LT", "Let", "Load", "Max",
                  "Min", "Mod", "Mul", "NE", "Not", "Or", "Ramp", "Reinterpret", "Select", "Shuffle", "Sub", "Variable",
                  "VectorReduce"]

    for ty in base_exprs:
        debugger.HandleCommand(
            f"type summary add Halide::Internal::{ty} --python-function lldbhalide.call_lldb_string"
        )

    for ty in ('Expr', 'Internal::Stmt'):
        debugger.HandleCommand(
            f"type summary add Halide::{ty} --python-function lldbhalide.call_lldb_string"
        )
        debugger.HandleCommand(
            f'type synthetic add Halide::{ty} -l lldbhalide.IRChildrenProvider'
        )

    for ty in ("Definition", "FuncSchedule", "ReductionDomain", "StageSchedule"):
        debugger.HandleCommand(
            f"type synthetic add Halide::Internal::{ty} -l lldbhalide.BoxChildrenProvider"
        )

    debugger.HandleCommand(
        f'type synthetic add Halide::Internal::Function -l lldbhalide.FunctionChildrenProvider'
    )

    debugger.HandleCommand("type summary add Halide::Internal::Dim -s '${var.var%S}'")
    debugger.HandleCommand("type summary add Halide::RVar --python-function lldbhalide.call_name")
    debugger.HandleCommand("type summary add Halide::Var --python-function lldbhalide.call_name")

    debugger.HandleCommand("type summary add halide_type_t -s '${var.code%S} bits=${var.bits%u} lanes=${var.lanes%u}'")
    debugger.HandleCommand("type summary add Halide::Internal::RefCount -s ${var.count.Value%S}")
