# python_tracer.py — line-level Python execution tracer for LeetTutor.
#
# SECURITY NOTE: execute_and_trace() uses exec() to run user-supplied code.
# This is intentional — LeetTutor is a personal tool for tracing YOUR OWN
# code locally. Do not host it as a public service without a sandbox.
import sys
import json
import inspect
import math
import ast


def serialize_value(val, depth=0, _seen=None):
    """Recursively serialize a Python value to a JSON-safe structure.

    - Tracks object ids in `_seen` to detect true reference cycles, not just
      depth, so long linked lists are not truncated to '...cycle...'.
    - Tuples/sets are serialized as tagged objects distinguishable from lists.
    - float inf/nan produce JSON-safe sentinel strings.
    - A hard depth cap of 50 still guards against pathological structures.
    """
    if _seen is None:
        _seen = set()

    if depth > 50:
        return "...depth limit..."

    if val is None:
        return None

    t_name = type(val).__name__

    if inspect.ismodule(val) or inspect.isclass(val) or inspect.isroutine(val) or callable(val):
        return str(val)

    if t_name == 'float':
        if val == float('inf'):
            return "Infinity"
        if val == float('-inf'):
            return "-Infinity"
        if math.isnan(val):
            return "NaN"
        return val

    if t_name in ('int', 'str', 'bool'):
        return val

    if isinstance(val, tuple):
        return {"__type__": "tuple", "items": [serialize_value(x, depth + 1, _seen) for x in val]}

    if isinstance(val, list):
        return [serialize_value(x, depth + 1, _seen) for x in val]

    if isinstance(val, dict):
        return {str(k): serialize_value(v, depth + 1, _seen) for k, v in val.items()}

    if isinstance(val, set):
        return {"__type__": "set", "items": [serialize_value(x, depth + 1, _seen) for x in val]}

    obj_id = id(val)
    if obj_id in _seen:
        return f"...cycle ({t_name})..."
    _seen = _seen | {obj_id}

    if hasattr(val, '__dict__'):
        obj_dict = {}
        for k, v in val.__dict__.items():
            if not k.startswith('_'):
                obj_dict[k] = serialize_value(v, depth + 1, _seen)
        obj_dict['__type__'] = t_name
        return obj_dict

    return str(val)


# ---------------------------------------------------------------------------
# Preamble prepended to user code before exec(). PREAMBLE_LINES is derived
# from this string automatically, and _PREAMBLE_NAMES is computed by
# actually running it once — both stay correct even if the preamble changes,
# with zero hardcoded magic numbers or name lists to maintain by hand.
# ---------------------------------------------------------------------------
_PREAMBLE = """\
from typing import *
import collections
import math
import heapq
import bisect

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

def to_list_node(arr):
    if not arr: return None
    head = ListNode(arr[0])
    curr = head
    for x in arr[1:]:
        curr.next = ListNode(x)
        curr = curr.next
    return head

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def to_tree_node(vals):
    if not vals: return None
    from collections import deque
    root = TreeNode(vals[0])
    q = deque([root])
    i = 1
    while q and i < len(vals):
        node = q.popleft()
        if i < len(vals) and vals[i] is not None:
            node.left = TreeNode(vals[i])
            q.append(node.left)
        i += 1
        if i < len(vals) and vals[i] is not None:
            node.right = TreeNode(vals[i])
            q.append(node.right)
        i += 1
    return root

"""

PREAMBLE_LINES = _PREAMBLE.count('\n')

# Every name the preamble binds (typing.* exports, stdlib modules, the
# ListNode/TreeNode/to_list_node/to_tree_node helpers). Anything with these
# names is "scaffolding", not the learner's own data, and is excluded from
# every variables panel. Computed by actually executing the preamble once,
# so it's automatically correct for whatever Python version this runs on
# (typing's exported names vary release to release) and automatically
# tracks any future edits to _PREAMBLE with no separate list to maintain.
_preamble_ns = {}
exec(_PREAMBLE, _preamble_ns)  # pylint: disable=exec-used
_PREAMBLE_NAMES = set(_preamble_ns.keys())
_PREAMBLE_NAMES.discard('__builtins__')


def _classify_line(text):
    """Cheap, dependency-free heuristic tag for what a line of code *does*.
    Used to give each step a short human label (Loop / Branch / Return /
    ...) in the UI without needing an LLM call per step.
    """
    t = text.strip()
    if not t:
        return None
    if t.startswith(('for ', 'while ')):
        return 'loop'
    if t.startswith(('if ', 'elif ', 'else')):
        return 'branch'
    if t.startswith('return'):
        return 'return'
    if t.startswith(('def ', 'class ')):
        return 'define'
    if t.startswith(('break', 'continue')):
        return 'control'
    if any(op in t for op in ('+=', '-=', '*=', '/=', '//=', '%=')):
        return 'update'
    # Plain assignment, but not a comparison (==, <=, >=, !=)
    if '=' in t and not any(op in t for op in ('==', '<=', '>=', '!=', ':=')):
        return 'assign'
    return 'call'


def _longest_parseable_prefix(code):
    """If `code` has a trailing syntax error — e.g. it was cut off
    mid-statement by a flaky extraction, or the learner is still in the
    middle of typing it — progressively drop lines from the END until
    what remains parses cleanly, rather than refusing to trace anything
    at all. Returns (trimmed_code, lines_dropped): dropped is 0 if the
    original code was already valid, and trimmed_code is None if not
    even a single line parses.
    """
    lines = code.split('\n')
    try:
        ast.parse(code)
        return code, 0
    except SyntaxError:
        pass

    for trim in range(1, len(lines) + 1):
        candidate = '\n'.join(lines[:len(lines) - trim])
        if not candidate.strip():
            break
        try:
            ast.parse(candidate)
            return candidate, trim
        except SyntaxError:
            continue
    return None, -1


class PythonCodeTracer:
    def __init__(self):
        self.steps = []

    def _make_trace(self, line_text_map, user_start):
        """Returns a sys.settrace callable that only records lines from the
        user's own code (never the preamble, never stdlib frames pulled in
        by `import heapq` etc.), and only meaningful local variables
        (never the ~80 names `from typing import *` dumps into scope).
        """

        def filtered_locals(frame):
            out = {}
            for name, val in frame.f_locals.items():
                if name.startswith('__') or name == 'self':
                    continue
                if name in _PREAMBLE_NAMES:
                    continue
                if inspect.ismodule(val) or inspect.isclass(val) or inspect.isroutine(val) or callable(val):
                    continue  # code structure (a def/class), not learner data
                out[name] = serialize_value(val)
            return out

        def build_call_stack(top_frame):
            """Walk f_back to collect every active frame that belongs to
            the exec'd code (skips frames once we leave '<string>' or drop
            below the user's own code region), outermost first — mirroring
            how Python Tutor renders the call stack.
            """
            frames = []
            f = top_frame
            while f is not None and f.f_code.co_filename == '<string>':
                frames.append(f)
                f = f.f_back
            frames.reverse()

            stack = []
            for fr in frames:
                stack.append({
                    "function": "Global" if fr.f_code.co_name == '<module>' else fr.f_code.co_name,
                    "line": fr.f_lineno,
                    "variables": filtered_locals(fr),
                })
            return stack

        def trace_calls(frame, event, arg):
            if event == 'line':
                line_no = frame.f_lineno
                if frame.f_code.co_filename != '<string>':
                    return trace_calls  # stdlib/site-packages frame — skip
                if line_no < user_start:
                    return trace_calls  # preamble line — skip

                current_locals = filtered_locals(frame)
                line_text = line_text_map.get(line_no, "")
                action = _classify_line(line_text)

                # Skip pure declaration lines (def/class headers) that have
                # no learner-meaningful state yet — this happens at module
                # scope (top-level `def foo(...):`) AND one level deeper,
                # inside the transient frame Python creates to execute a
                # class body (e.g. `class Solution:` briefly runs its own
                # frame just to define methods). Both were the source of
                # the original noise: dozens of near-empty "steps" for
                # every class/def header before real execution even starts.
                if not current_locals and action == 'define':
                    return trace_calls

                is_module_frame = frame.f_code.co_name == '<module>'
                self.steps.append({
                    "line": line_no,
                    "lineText": line_text,
                    "action": _classify_line(line_text),
                    "function": "Global" if is_module_frame else frame.f_code.co_name,
                    "stackDepth": len(inspect.stack()),
                    "callStack": build_call_stack(frame),
                    "isResult": is_module_frame and '_result' in current_locals,
                })
            return trace_calls

        return trace_calls

    def execute_and_trace(self, code_to_run, function_call):
        self.steps = []
        original_line_count = len(code_to_run.split('\n'))

        trimmed_code, dropped = _longest_parseable_prefix(code_to_run)

        if trimmed_code is None:
            # Not even a single line parses on its own — nothing to salvage.
            # Report the real syntax error against the original code so the
            # line number/message are meaningful.
            try:
                ast.parse(code_to_run)
            except SyntaxError as e:
                self.steps.append({
                    "error": str(e),
                    "error_type": "SyntaxError",
                    "line": (e.lineno or 1) + PREAMBLE_LINES,
                    "lineText": code_to_run.split('\n')[(e.lineno or 1) - 1] if e.lineno else "",
                })
            return json.dumps({"steps": self.steps, "truncated": None})

        truncated_info = None
        if dropped > 0:
            traced_lines = original_line_count - dropped
            truncated_info = {
                "droppedLines": dropped,
                "totalLines": original_line_count,
                "tracedLines": traced_lines,
                "message": (
                    f"Only tracing the first {traced_lines} of {original_line_count} lines — "
                    f"the rest doesn't parse as valid Python (a syntax error starts around there), "
                    f"so it was left out rather than blocking the trace entirely."
                ),
            }

        code_lines = trimmed_code.split('\n')
        user_start = PREAMBLE_LINES + 1
        call_line_no = user_start + len(code_lines) + 1  # one blank line, then the call

        line_text_map = {user_start + i: line for i, line in enumerate(code_lines)}
        line_text_map[call_line_no] = f"_result = ({function_call})"

        execution_block = _PREAMBLE + trimmed_code + f"\n\n_result = ({function_call})"

        tracer = self._make_trace(line_text_map, user_start)

        # Single shared namespace for both globals and locals. If these were
        # two different dicts, top-level classes like ListNode/TreeNode
        # would be stored via STORE_NAME into the locals dict, but nested
        # helper functions (to_list_node, to_tree_node, and the user's own
        # Solution methods) resolve free names via LOAD_GLOBAL against
        # func.__globals__ — the *other* dict — causing a NameError even
        # though the class is defined right above it.
        ns = {}
        sys.settrace(tracer)
        try:
            exec(execution_block, ns)  # pylint: disable=exec-used
            # sys.settrace's 'line' event fires *before* a line executes, so
            # there is no trace event after the final `_result = (...)`
            # assignment completes — the program just ends. Append the
            # result explicitly so the UI always has a definitive final step.
            if '_result' in ns:
                self.steps.append({
                    "line": call_line_no,
                    "lineText": line_text_map.get(call_line_no, ""),
                    "action": "return",
                    "function": "Global",
                    "stackDepth": 1,
                    "callStack": [{
                        "function": "Global",
                        "line": call_line_no,
                        "variables": {"_result": serialize_value(ns["_result"])},
                    }],
                    "isResult": True,
                })
        except Exception as e:
            tb = sys.exc_info()[2]
            while tb and tb.tb_next:
                tb = tb.tb_next
            raw_line = tb.tb_lineno if tb else -1
            self.steps.append({
                "error": str(e),
                "error_type": type(e).__name__,
                "line": raw_line,
                "lineText": line_text_map.get(raw_line, ""),
            })
        finally:
            sys.settrace(None)

        return json.dumps({"steps": self.steps, "truncated": truncated_info})


tracer = PythonCodeTracer()

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

    def js_execute_and_trace(code_js, call_js):
        return tracer.execute_and_trace(str(code_js), str(call_js))

    window.executePythonTracer = create_proxy(js_execute_and_trace)
    window.TRACER_PREAMBLE_LINES = PREAMBLE_LINES

    if hasattr(window, "onPythonLoaded"):
        window.onPythonLoaded()
except ImportError:
    pass