# complexity_analyzer.py — heuristic time/space complexity estimation.
#
# IMPORTANT HONESTY NOTE: determining the exact asymptotic complexity of
# arbitrary code is undecidable in general. This module does NOT prove
# anything — it applies the same pattern-matching heuristics an experienced
# reviewer applies at a glance (loop nesting depth, recursion branching
# factor, divide-and-conquer signals, sort() calls, memoization, hidden
# O(n) checks like `x in a_list`). It is a *teaching aid*, not a verifier.
# Every result carries the reasoning that produced it so the learner can
# judge for themselves, and a confidence label that drops for patterns the
# heuristic is less sure about.
import ast

# Ordered worst-to-best-free (best/smallest first) so comparisons are simple
# index lookups instead of a bespoke comparator per pair.
_BIGO_ORDER = [
    'O(1)', 'O(log n)', 'O(n)', 'O(n log n)',
    'O(n^2)', 'O(n^2 log n)', 'O(n^3)', 'O(2^n)', 'O(n!)',
]


def _rank(cls):
    try:
        return _BIGO_ORDER.index(cls)
    except ValueError:
        return len(_BIGO_ORDER)  # unrecognized -> sorts as "worst/unknown"


def _bigo_max(a, b):
    return a if _rank(a) >= _rank(b) else b


def _call_name(node):
    """Best-effort function name for an ast.Call node (`foo()` or `self.foo()`)."""
    if isinstance(node.func, ast.Attribute):
        return node.func.attr
    if isinstance(node.func, ast.Name):
        return node.func.id
    return None


def _looks_like_halving(while_node):
    """Heuristic for while-loops that shrink their range each iteration
    (binary-search style), e.g. `hi = mid - 1`, `x //= 2`, `n = n // 2`."""
    for child in ast.walk(while_node):
        if isinstance(child, ast.AugAssign) and isinstance(child.op, ast.FloorDiv):
            return True
        if isinstance(child, ast.Assign) and isinstance(child.value, ast.BinOp) \
                and isinstance(child.value.op, (ast.FloorDiv, ast.Div)):
            return True
    return False


def _arg_looks_halved(node, halving_vars=frozenset()):
    """Heuristic for a call argument that looks like 'half the input' —
    a slice (`arr[mid:]`), a `// 2` expression, or a reference to a
    variable that was itself computed via halving earlier (`mid = (lo+hi)//2`
    then `search(mid + 1, hi)`) — the signature of divide-and-conquer or
    binary-search-style recursion."""
    if isinstance(node, ast.Subscript):
        return True
    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.FloorDiv):
        return True
    for n in ast.walk(node):
        if isinstance(n, ast.Name) and n.id in halving_vars:
            return True
    return False


def _find_halving_vars(func_node):
    """Names assigned via `x = <expr> // 2` (or `/ 2`) anywhere in the
    function — e.g. `mid = (lo + hi) // 2` — used so a later call like
    `search(mid + 1, hi)` is still recognized as operating on half the
    remaining range even though the division isn't in the call itself."""
    names = set()
    for node in ast.walk(func_node):
        if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
            val = node.value
            if isinstance(val, ast.BinOp) and isinstance(val.op, (ast.FloorDiv, ast.Div)):
                names.add(node.targets[0].id)
    return names


def _infer_var_kinds(func_node):
    """Best-effort map of local variable name -> 'dict' | 'set' | 'list',
    based on the literal/constructor used at assignment. Used to avoid
    false-positive 'hidden O(n) in-check' warnings against dicts/sets,
    which is by far the most common LeetCode idiom (`seen = {}` then
    `if x in seen`) and must not be mislabeled as O(n^2).
    """
    kinds = {}
    for node in ast.walk(func_node):
        if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
            name = node.targets[0].id
            val = node.value
            if isinstance(val, ast.Dict):
                kinds[name] = 'dict'
            elif isinstance(val, ast.Set):
                kinds[name] = 'set'
            elif isinstance(val, ast.List):
                kinds[name] = 'list'
            elif isinstance(val, ast.Call) and isinstance(val.func, ast.Name):
                if val.func.id == 'dict':
                    kinds[name] = 'dict'
                elif val.func.id == 'set':
                    kinds[name] = 'set'
                elif val.func.id == 'list':
                    kinds[name] = 'list'
    return kinds


def _always_exits(body):
    """True if the last statement in this block unconditionally leaves the
    function/loop (return/raise/continue/break) — the AST signature of the
    common 'guard clause' idiom (`if x: return y` followed by more code
    that's only reached when the guard was false)."""
    if not body:
        return False
    return isinstance(body[-1], (ast.Return, ast.Raise, ast.Continue, ast.Break))


def _count_calls_flat(node, func_name):
    return sum(1 for n in ast.walk(node) if isinstance(n, ast.Call) and _call_name(n) == func_name)


def _max_simultaneous_calls(stmts, func_name):
    """The max number of self-recursive calls that can occur together
    within ONE execution of this statement sequence.

    Sequential statements that both unconditionally execute are summed
    (e.g. merge sort's two recursive calls in a row: branching factor 2).
    Mutually exclusive branches of the same if/else — including the guard-
    clause idiom, where a trailing `if: return ...` makes every statement
    after it an implicit 'else' — take the MAX instead of the sum, since
    only one path actually runs per invocation (e.g. binary-search-style
    recursion with a call in each arm: branching factor 1, not 2).
    """
    total = 0
    for i, stmt in enumerate(stmts):
        if isinstance(stmt, ast.If):
            body_calls = _max_simultaneous_calls(stmt.body, func_name)
            if _always_exits(stmt.body):
                # Everything after this if is an implicit else.
                rest = _max_simultaneous_calls(list(stmts[i + 1:]), func_name)
                orelse_calls = _max_simultaneous_calls(stmt.orelse, func_name) + rest if stmt.orelse else rest
                return total + max(body_calls, orelse_calls)
            orelse_calls = _max_simultaneous_calls(stmt.orelse, func_name) if stmt.orelse else 0
            total += max(body_calls, orelse_calls)
        elif isinstance(stmt, (ast.For, ast.While)):
            total += _max_simultaneous_calls(stmt.body, func_name)
        else:
            total += _count_calls_flat(stmt, func_name)
    return total


class _FunctionAnalyzer(ast.NodeVisitor):
    """Walks a single function's body (does not descend into other
    top-level function/class defs) collecting complexity signals."""

    def __init__(self, func_name, var_kinds=None, halving_vars=None):
        self.func_name = func_name
        self.var_kinds = var_kinds or {}
        self.halving_vars = halving_vars or set()
        self.cur_depth = 0
        self.max_depth = 0
        self.time_reasons = []
        self.space_reasons = []
        self.warnings = []

        self.has_sort = False
        self.log_loop = False

        self.self_call_count = 0
        self.halves_input = False
        self.has_memo = False
        self.recurses_on_substructure = False

        self.scaling_structures = 0     # e.g. dp = [0]*n
        self.nested_scaling_structures = 0  # e.g. dp = [[0]*n for _ in range(n)]
        self.appends_in_loop = False

    # ── loops ──
    def visit_For(self, node):
        self.cur_depth += 1
        self.max_depth = max(self.max_depth, self.cur_depth)
        self.time_reasons.append(f"line {node.lineno}: for-loop (nesting level {self.cur_depth})")
        self.generic_visit(node)
        self.cur_depth -= 1

    def visit_While(self, node):
        if _looks_like_halving(node):
            self.log_loop = True
            self.time_reasons.append(f"line {node.lineno}: while-loop that shrinks its range each pass (binary-search-like) → O(log n) factor")
            self.generic_visit(node)
        else:
            self.cur_depth += 1
            self.max_depth = max(self.max_depth, self.cur_depth)
            self.time_reasons.append(f"line {node.lineno}: while-loop (nesting level {self.cur_depth})")
            self.generic_visit(node)
            self.cur_depth -= 1

    # ── calls: sort(), recursion, memoization hints ──
    def visit_Call(self, node):
        name = _call_name(node)

        if name in ('sort', 'sorted'):
            self.has_sort = True
            self.time_reasons.append(f"line {node.lineno}: {name}() → at least O(n log n)")

        if name == self.func_name:
            self.self_call_count += 1
            self.time_reasons.append(f"line {node.lineno}: recursive call to {self.func_name}()")
            for arg in node.args:
                if _arg_looks_halved(arg, self.halving_vars):
                    self.halves_input = True
                if isinstance(arg, ast.Attribute):
                    # e.g. maxDepth(root.left) / maxDepth(root.right) —
                    # each call descends into a DISJOINT child, unlike
                    # fib(n-1)/fib(n-2) which both operate on overlapping
                    # shrinking ranges of the SAME sequence. A node can't
                    # be two different calls' subtree at once, so total
                    # work across all calls is bounded by structure size,
                    # not exponential in branching factor.
                    self.recurses_on_substructure = True

        if name in ('lru_cache', 'cache'):
            self.has_memo = True
            self.time_reasons.append(f"line {node.lineno}: @{name} decorator → memoized recursion")

        self.generic_visit(node)

    def visit_Compare(self, node):
        for op, comparator in zip(node.ops, node.comparators):
            if not isinstance(op, ast.In) or self.cur_depth == 0 or not isinstance(comparator, ast.Name):
                continue
            kind = self.var_kinds.get(comparator.id)
            if kind == 'list':
                self.warnings.append(
                    f"line {node.lineno}: '`in`' check inside a loop against `{comparator.id}`, "
                    f"which looks like a list — this adds another O(n) factor per check."
                )
            # dict/set membership is O(1) average case — never warn on those.
            # Unknown-kind names (parameters, etc.) are left unflagged too,
            # since a false "this is O(n^2)" is more misleading than a
            # missed warning.
        self.generic_visit(node)

    def visit_Return(self, node):
        self.generic_visit(node)

    def visit_Assign(self, node):
        var_name = node.targets[0].id if node.targets and isinstance(node.targets[0], ast.Name) else None
        if var_name and ('memo' in var_name.lower() or 'cache' in var_name.lower()):
            self.has_memo = True
        # memo[key] = ... style cache writes (target is a Subscript on a
        # memo/cache-named variable)
        if node.targets and isinstance(node.targets[0], ast.Subscript):
            base = node.targets[0].value
            if isinstance(base, ast.Name):
                base_kind = self.var_kinds.get(base.id)
                if 'memo' in base.id.lower() or 'cache' in base.id.lower():
                    self.has_memo = True
                # A dict/set (or unknown-but-subscripted) container being
                # written to inside a loop is the classic hash-map-grows-
                # with-input pattern (`seen[n] = i`) — O(n) space.
                if self.cur_depth > 0 and base_kind in ('dict', None):
                    self.scaling_structures += 1
                    self.space_reasons.append(
                        f"line {node.lineno}: `{base.id}[...]` written inside a loop → O(n) space (a hash map/array growing with the input)"
                    )

        # dp = [0] * n  /  dp = [x for x in ...]
        if isinstance(node.value, ast.BinOp) and isinstance(node.value.op, ast.Mult):
            self.scaling_structures += 1
            self.space_reasons.append(f"line {node.lineno}: list built with `* n` repetition → O(n) space")
        elif isinstance(node.value, ast.ListComp):
            elt = node.value.elt
            # Matches `[[0]*n for _ in range(n)]` (elt is itself a repeated
            # list) and `[[x for x in row] for row in grid]` (elt is itself
            # a comprehension) — both build a genuinely 2-D structure.
            is_2d = isinstance(elt, ast.ListComp) or (
                isinstance(elt, ast.BinOp) and isinstance(elt.op, ast.Mult)
            )
            if is_2d:
                self.nested_scaling_structures += 1
                self.space_reasons.append(f"line {node.lineno}: 2-D list comprehension → O(n^2) space")
            else:
                self.scaling_structures += 1
                self.space_reasons.append(f"line {node.lineno}: list comprehension → O(n) space")
        self.generic_visit(node)

    def visit_Expr(self, node):
        # foo.append(x) inside a loop -> a result list that grows with input
        if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Attribute):
            if node.value.func.attr in ('append', 'add') and self.cur_depth > 0:
                self.appends_in_loop = True
        self.generic_visit(node)


def _find_function(tree, func_name):
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
            return node
    return None



def analyze(code, function_name):
    """Analyze `code` (a string of Python source) for the named function
    and return a dict describing an estimated time/space complexity with
    supporting reasoning. Never raises — falls back to an 'unknown' result
    with an explanatory warning if the function can't be located or parsed.
    """
    try:
        tree = ast.parse(code)
    except SyntaxError as e:
        return _unknown_result(f"Could not parse code: {e}")

    func_node = _find_function(tree, function_name)
    if func_node is None:
        return _unknown_result(f"Could not find a function named '{function_name}' to analyze.")

    fa = _FunctionAnalyzer(function_name, var_kinds=_infer_var_kinds(func_node), halving_vars=_find_halving_vars(func_node))
    fa.generic_visit(func_node)

    # A parameter literally named memo/cache (e.g. `def fib(n, memo={}):`)
    # is a strong memoization signal that a body-only scan can miss, since
    # default argument values aren't part of the traced body statements.
    param_names = [a.arg for a in func_node.args.args]
    if any('memo' in p.lower() or 'cache' in p.lower() for p in param_names):
        fa.has_memo = True

    # ── Time complexity ──
    depth = fa.max_depth
    hidden_linear_factor = any('adds another O(n) factor' in w for w in fa.warnings)
    effective_depth = depth + (1 if hidden_linear_factor else 0)

    if effective_depth == 0:
        time_class = 'O(log n)' if fa.log_loop else 'O(1)'
    elif effective_depth == 1:
        time_class = 'O(n)'
    elif effective_depth == 2:
        time_class = 'O(n^2)'
    elif effective_depth == 3:
        time_class = 'O(n^3)'
    else:
        time_class = f'O(n^{effective_depth})'

    if fa.log_loop and effective_depth >= 1:
        time_class = _bigo_max(time_class, 'O(n log n)')

    if fa.has_sort:
        time_class = _bigo_max(time_class, 'O(n log n)')

    confidence = 'medium'

    if fa.self_call_count > 0:
        branching = _max_simultaneous_calls(func_node.body, function_name) or 1
        if fa.has_memo:
            rec_class = 'O(n)'
            fa.time_reasons.append("recursion + a memo/cache detected → treated as roughly O(n) (revisits avoided); actual bound depends on the state space")
            confidence = 'low'
        elif fa.recurses_on_substructure:
            rec_class = 'O(n)'
            fa.time_reasons.append("recursive calls descend into child attributes (e.g. `.left`/`.right`/`.next`) — each node is visited once, so total work is linear in structure size even though there are multiple calls per node")
        elif branching >= 2 and not fa.halves_input:
            rec_class = 'O(2^n)'
            fa.time_reasons.append(f"{branching} recursive calls occur together in a single invocation, no halving or memoization detected → exponential branching")
        elif branching >= 2 and fa.halves_input:
            rec_class = 'O(n log n)'
            fa.time_reasons.append(f"{branching} recursive calls per invocation, each on roughly half the input → divide-and-conquer shape")
        elif fa.halves_input:
            rec_class = 'O(log n)'
            fa.time_reasons.append("recursive call(s) on roughly half the input, only one path taken per invocation → binary-search-like recursion")
        else:
            rec_class = 'O(n)'
            fa.time_reasons.append("recursive call(s) with only one path taken per invocation, input shrinks by a constant amount → linear recursion depth")
        time_class = _bigo_max(time_class, rec_class)

    # ── Space complexity ──
    space_class = 'O(1)'
    if fa.nested_scaling_structures > 0:
        space_class = 'O(n^2)'
    elif fa.scaling_structures > 0 or fa.appends_in_loop:
        space_class = 'O(n)'
        if fa.appends_in_loop and not any('append' in r for r in fa.space_reasons):
            fa.space_reasons.append("a list is grown (`.append`/`.add`) inside a loop → O(n) output/auxiliary space")

    if fa.self_call_count > 0:
        rec_space = 'O(log n)' if fa.halves_input else 'O(n)'
        if fa.recurses_on_substructure:
            fa.space_reasons.append("recursion adds O(n) call-stack depth in the worst case (a skewed/degenerate structure); O(log n) if the structure stays balanced")
        else:
            fa.space_reasons.append(f"recursion adds {rec_space} call-stack depth")
        space_class = _bigo_max(space_class, rec_space)

    if not fa.time_reasons:
        fa.time_reasons.append("no loops or recursion detected → constant-time body")
    if not fa.space_reasons:
        fa.space_reasons.append("no data structures whose size scales with input, and no recursion → constant extra space")

    return {
        'function': function_name,
        'found': True,
        'time': {'class': time_class, 'confidence': confidence, 'reasons': fa.time_reasons},
        'space': {'class': space_class, 'confidence': confidence, 'reasons': fa.space_reasons},
        'warnings': fa.warnings,
    }


def _unknown_result(message):
    return {
        'function': None,
        'found': False,
        'time': {'class': 'unknown', 'confidence': 'none', 'reasons': [message]},
        'space': {'class': 'unknown', 'confidence': 'none', 'reasons': [message]},
        'warnings': [],
    }


def extract_call_function_name(call_string):
    """Parse a call expression like `Solution().twoSum([2,7], 9)` and
    return the OUTERMOST function/method name being called — 'twoSum'
    here, not 'to_tree_node' if that happens to appear nested inside an
    argument (`Solution().maxDepth(to_tree_node([...]))`). A proper AST
    parse handles nesting correctly where a regex scan can't.
    """
    try:
        tree = ast.parse(call_string.strip(), mode='eval')
        node = tree.body
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Attribute):
                return node.func.attr
            if isinstance(node.func, ast.Name):
                return node.func.id
    except SyntaxError:
        pass
    return None


# --- FFI JS Binding (only active inside Pyodide / PyScript) ---
try:
    from pyodide.ffi import create_proxy
    from js import window
    import json as _json

    def js_analyze_complexity(code_js, function_name_js):
        return _json.dumps(analyze(str(code_js), str(function_name_js)))

    def js_extract_call_function_name(call_js):
        return extract_call_function_name(str(call_js)) or ""

    window.analyzeComplexity = create_proxy(js_analyze_complexity)
    window.extractCallFunctionName = create_proxy(js_extract_call_function_name)
except ImportError:
    pass