from __future__ import annotations

import ast
from typing import Any, Callable

import pyccolo as pyc

from pipenb.utils import has_augmentations, is_outer_or_allowlisted


def parent_is_or_boolop(node_id: int) -> bool:
    parent = pyc.BaseTracer.containing_ast_by_id.get(node_id)
    return isinstance(parent, ast.BoolOp) and isinstance(parent.op, ast.Or)


def should_instrument_for_spec(
    spec: pyc.AugmentationSpec | set[pyc.AugmentationSpec] | str,
    attr: str | None = None,
) -> Callable[[ast.AST | int], bool]:
    if isinstance(spec, str):
        spec = getattr(NullishTracer, spec)
    return lambda node: is_outer_or_allowlisted(node) and has_augmentations(
        getattr(node, attr or "", node), spec
    )


def should_instrument_for_spec_on_parent(
    spec: pyc.AugmentationSpec | set[pyc.AugmentationSpec] | str,
    attr: str | None = None,
) -> Callable[[ast.AST | int], bool]:
    if isinstance(spec, str):
        spec = getattr(NullishTracer, spec)

    def should_instrument(node_or_id: ast.AST | int) -> bool:
        node_id = node_or_id if isinstance(node_or_id, int) else id(node_or_id)
        parent = pyc.BaseTracer.containing_ast_by_id.get(node_id)
        if parent is None or not is_outer_or_allowlisted(parent):
            return False
        return has_augmentations(getattr(parent, attr or "", parent), spec)

    return should_instrument


class NullishInstrumentationChainChecker(ast.NodeVisitor):
    def __init__(self) -> None:
        self._contains_nullish_instrumentation = False

    def __call__(self, node_id: int) -> bool:
        if not is_outer_or_allowlisted(node_id):
            return False
        node = pyc.BaseTracer.ast_node_by_id.get(node_id)
        if node is None:
            return False
        self._contains_nullish_instrumentation = False
        self.visit(node)
        return self._contains_nullish_instrumentation

    def visit_Attribute(self, node: ast.Attribute) -> None:
        if has_augmentations(
            node,
            {
                NullishTracer.optional_chaining_spec,
                NullishTracer.permissive_attr_dereference_spec,
            },
        ):
            self._contains_nullish_instrumentation = True
            return
        self.visit(node.value)

    def visit_Call(self, node: ast.Call) -> None:
        if has_augmentations(node, NullishTracer.call_optional_chaining_spec):
            self._contains_nullish_instrumentation = True
            return
        self.visit(node.func)


class NullishTracer(pyc.BaseTracer):
    global_guards_enabled = False

    class ResolvesToNone:
        def __init__(self, eventually: bool) -> None:
            self.__eventually = eventually

        def __getattr__(self, _item: str):
            if self.__eventually:
                return self
            else:
                return None

        def __call__(self, *_, **__):
            return self

    resolves_to_none_eventually = ResolvesToNone(eventually=True)
    resolves_to_none_immediately = ResolvesToNone(eventually=False)

    call_optional_chaining_spec = pyc.AugmentationSpec(
        aug_type=pyc.AugmentationType.call, token="?.(", replacement="("
    )

    optional_chaining_spec = pyc.AugmentationSpec(
        aug_type=pyc.AugmentationType.dot_suffix, token="?.", replacement="."
    )

    permissive_attr_dereference_spec = pyc.AugmentationSpec(
        aug_type=pyc.AugmentationType.dot_suffix, token=".?", replacement="."
    )

    nullish_coalescing_spec = pyc.AugmentationSpec(
        aug_type=pyc.AugmentationType.boolop, token="??", replacement=" or "
    )

    chain_checker = NullishInstrumentationChainChecker()

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._saved_ret_expr = None
        self.lexical_call_stack: pyc.TraceStack = self.make_stack()
        with self.lexical_call_stack.register_stack_state():
            # TODO: pop this the right number of times if an exception occurs
            self.cur_call_is_none_resolver: bool = False
        self.lexical_nullish_stack: pyc.TraceStack = self.make_stack()
        with self.lexical_nullish_stack.register_stack_state():
            self.cur_boolop_has_nullish_coalescer = False
            self.coalesced_value: Any | None = None

    def clear_stacks(self):
        # will be registered as a post_run_cell event
        while len(self.lexical_nullish_stack) > 0:
            self.lexical_nullish_stack.pop()

    @pyc.register_handler(
        pyc.before_attribute_load,
        when=should_instrument_for_spec(
            {optional_chaining_spec, permissive_attr_dereference_spec}
        ),
    )
    def handle_before_attr(self, obj, node: ast.Attribute, *_, **__):
        this_node_augmentations = self.get_augmentations(id(node))
        if self.optional_chaining_spec in this_node_augmentations and obj is None:
            return self.resolves_to_none_eventually
        elif (
            self.permissive_attr_dereference_spec in this_node_augmentations
            and not hasattr(obj, node.attr)
        ):
            return self.resolves_to_none_immediately
        else:
            return obj

    @pyc.register_handler(
        pyc.before_call,
        when=should_instrument_for_spec(call_optional_chaining_spec),
    )
    def handle_before_call(self, func, *_, **__):
        if func is None:
            func = self.resolves_to_none_eventually
        with self.lexical_call_stack.push():
            self.cur_call_is_none_resolver = func is self.resolves_to_none_eventually
        return func

    @pyc.register_raw_handler(
        pyc.before_argument,
        when=should_instrument_for_spec_on_parent(call_optional_chaining_spec),
    )
    def handle_before_arg(self, arg_lambda, *_, **__):
        if self.cur_call_is_none_resolver:
            return lambda: None
        else:
            return arg_lambda

    @pyc.register_raw_handler(
        pyc.after_call,
        when=should_instrument_for_spec(call_optional_chaining_spec),
    )
    def handle_after_call(self, *_, **__):
        self.lexical_call_stack.pop()

    @pyc.register_raw_handler(pyc.after_load_complex_symbol, when=chain_checker)
    def handle_after_load_complex_symbol(self, ret, *_, **__):
        if isinstance(ret, self.ResolvesToNone):
            return pyc.Null
        else:
            return ret

    @pyc.register_handler(
        pyc.before_boolop,
        when=lambda node: isinstance(node.op, ast.Or)
        and should_instrument_for_spec("nullish_coalescing_spec", attr="values")(node),
    )
    def before_or_boolop(self, ret, node: ast.BoolOp, *_, **__):
        with self.lexical_nullish_stack.push():
            self.cur_boolop_has_nullish_coalescer = any(
                self.nullish_coalescing_spec in self.get_augmentations(id(val))
                for val in node.values
            )
            self.coalesced_value = None
        return ret

    @pyc.register_handler(
        pyc.after_boolop,
        when=lambda node: isinstance(node.op, ast.Or)
        and should_instrument_for_spec("nullish_coalescing_spec", attr="values")(node),
    )
    def after_or_boolop(self, *_, **__):
        self.lexical_nullish_stack.pop()

    def _maybe_compute_nullish_coalesced_value(self, ret, node_id: int) -> None:
        if self.coalesced_value is not None:
            return
        val = ret()
        if self.nullish_coalescing_spec in self.get_augmentations(node_id):
            self.coalesced_value = None if val is None else val
        else:
            self.coalesced_value = val or None

    @pyc.register_raw_handler(
        pyc.before_boolop_arg,
        when=lambda node_id: parent_is_or_boolop(node_id)
        and should_instrument_for_spec_on_parent(
            "nullish_coalescing_spec", attr="values"
        )(node_id),
    )
    def before_or_boolup(self, ret, node_id: int, *_, is_last: bool, **__):
        if self.cur_boolop_has_nullish_coalescer:
            if is_last:
                if self.coalesced_value is None:
                    return ret
                else:
                    return lambda: self.coalesced_value
            else:
                self._maybe_compute_nullish_coalesced_value(ret, node_id)
                return lambda: None
        else:
            return ret
