"""Solveit tools for inspecting the symbol table and importing modules"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_inspecttools.ipynb.

# %% auto #0
__all__ = ['importmodule', 'SymbolNotFound', 'resolve', 'symsrc', 'showsrc', 'gettype', 'getdir', 'getval', 'getnth', 'symlen',
           'symslice', 'symsearch', 'symset', 'run_code_interactive', 'tracetool', 'fmt_trace', 'inspect_tool_info']

# %% ../nbs/02_inspecttools.ipynb #5ae7ad05
import inspect, re, sys, ast, builtins, os, linecache
from importlib import import_module
from . import add_msg

# %% ../nbs/02_inspecttools.ipynb #9778fca8
def _find_frame_dict(var:str):
    "Find the dict (globals or locals) containing var"
    frame = inspect.currentframe().f_back
    while frame:
        if var in frame.f_globals: return frame.f_globals
        frame = frame.f_back
    raise ValueError(f"Could not find {var} in any scope")

# %% ../nbs/02_inspecttools.ipynb #7fd415fa
def importmodule(
    mod: str # The module to import (e.g. 'torch.nn.functional')
):
    """Import a module into the caller's global namespace so it's available for `symsrc`, `getval`, `getdir`, etc.
    Use this before inspecting or using symbols from modules not yet imported."""
    g = _find_frame_dict('__msg_id')
    import_module(mod)
    g[mod.split('.')[0]] = import_module(mod.split('.')[0])

# %% ../nbs/02_inspecttools.ipynb #e40f68c9
class SymbolNotFound(Exception): pass

_last = None

def resolve(
    sym: str  # Dotted symbol path, with optional [n] indexing, e.g. "module.attr.subattr[1]" or "_last" for previous result
):
    """Resolve a dotted symbol string to its Python object, with optional [n] indexing.
    Sets global `_last` to the resolved object for chaining.
    Pass `"_last"` to reference the result of the previous tool call.

    Examples:

    - `resolve("sympy.sets.sets.Interval")` -> `<class 'sympy.sets.sets.Interval'>`
    - `resolve("mylist[2]")` -> third element of mylist"""
    global _last
    if (sym := sym.strip()) == '_last': return _last
    g = _find_frame_dict('__msg_id')
    if match := re.match(r'^(\w+)\[(\d+)\]$', sym):
        attr, idx = match.groups()
        parts, _last = ['_last'], _last[int(idx)] if attr == '_last' else g[attr][int(idx)]
    else: parts = re.split(r'\.(?![^\[]*\])', sym)
    try: obj = _last if parts[0] == '_last' else g[parts[0]]
    except KeyError: raise SymbolNotFound(f"Symbol '{parts[0]}' not found. Consider using `importmodule` first.")
    for part in parts[1:]:
        match = re.match(r'(\w+)\[(\d+)\]$', part)
        if match:
            attr, idx = match.groups()
            obj = getattr(obj, attr)[int(idx)]
        else: obj = getattr(obj, part)
    _last = obj
    return obj

# %% ../nbs/02_inspecttools.ipynb #659bf879
def _src_from_lines(lines, start):
    "Extract a single definition from lines starting at start (0-indexed)"
    src = ''.join(lines[start:])
    try: tree = ast.parse(src)
    except SyntaxError: return None
    if not tree.body: return None
    return ''.join(lines[start:start + tree.body[0].end_lineno])

# %% ../nbs/02_inspecttools.ipynb #cee0fc8c
def symsrc(
    sym: str  # Dotted symbol path (e.g `Interval` or `sympy.sets.sets.Interval`) or "_last" for previous result
):
    """Get the source code for a symbol.

    Examples:

    - `symsrc("Interval")` -> source code of Interval class if it's already imported
    - `symsrc("sympy.sets.sets.Interval")` -> source code of Interval class
    - `symsrc("_last")` -> source of object from previous tool call
    - For dispatchers or registries of callables: `getnth("module.dispatcher.funcs", n) then symsrc("_last")`"""
    try: obj = resolve(sym)
    except SymbolNotFound as e: return str(e)
    if not (isinstance(obj, type) or callable(obj)): obj = obj.__class__
    try: fname = inspect.getfile(obj)
    except OSError: fname = "<session>"
    try: return f"File: {fname}\n\n{inspect.getsource(obj)}"
    except OSError: pass
    name = obj.__name__
    pat = rf'^(class|def)\s+{name}\b'
    for fname, (_, _, lines, _) in linecache.cache.items():
        src = ''.join(lines)
        if match := re.search(pat, src, re.MULTILINE):
            start = src[:match.start()].count('\n')
            if extracted := _src_from_lines(lines, start): return f"File: {fname}\n\n{extracted}"
    raise OSError(f"Source for {name} not found")

# %% ../nbs/02_inspecttools.ipynb #ef12c6ca
def showsrc(
    sym: str  # Dotted symbol path (e.g `Interval` or `sympy.sets.sets.Interval`) or "_last" for previous result
):
    "Create a note to show the user (and following LLM prompts) the source of `sym`, following `symsrc` rules"
    res = symsrc(sym)
    add_msg(f'```python\n{res}\n```')
    return {'success': 'Message has been added to dialog successfully'}

# %% ../nbs/02_inspecttools.ipynb #bbf67405
def gettype(
    sym: str  # Dotted symbol path (e.g `Interval` or `sympy.sets.sets.Interval`) or "_last" for previous result
):
    """Get the type of a symbol and set `_last`.

    Examples:

    - `gettype("sympy.sets.sets.Interval")` -> `<class 'type'>`
    - `gettype("_last")` -> type of previous result"""
    return type(resolve(sym))

# %% ../nbs/02_inspecttools.ipynb #1cd34596
def getdir(
    sym: str,  # Dotted symbol path (e.g `Interval` or `sympy.sets.sets.Interval`) or "_last" for previous result
    exclude_private: bool=False # Filter out attrs starting with "_"
):
    """Get dir() listing of a symbol's attributes and set `_last`. E.g: `getdir("sympy.Interval")` -> `['__add__', '__and__', ...]`"""
    res = dir(resolve(sym))
    if not exclude_private: return res
    return [o for o in res if o[0]!='_']

# %% ../nbs/02_inspecttools.ipynb #0af80ded
def getval(
    sym: str  # Dotted symbol path (e.g `Interval` or `sympy.sets.sets.Interval`) or "_last" for previous result
):
    """Get repr of a symbol's value and set `_last`.

    Examples:
    
    - `getval("sympy.sets.sets.Interval")` -> `<class 'sympy.sets.sets.Interval'>`
    - `getval("some_dict.keys")` -> `dict_keys([...])`"""
    return repr(resolve(sym))

# %% ../nbs/02_inspecttools.ipynb #9542de0b
def getnth(
    sym: str,  # Dotted symbol path to a dict or object with .values()
    n: int     # Index into the values (0-based)
):
    """Get the nth value from a dict (or any object with .values()). Sets `_last` so you can chain with `symsrc("_last")` etc.

    Examples:
    
    - `getnth("dispatcher.funcs", 12)` -> 13th registered function
    - `getnth("dispatcher.funcs", 0); symsrc("_last")` -> source of first handler"""
    global _last
    _last = list(resolve(sym).values())[n]
    return _last

# %% ../nbs/02_inspecttools.ipynb #4ac6ca2d
def symlen(
    sym: str  # Dotted symbol path or "_last" for previous result
):
    "Returns the length of the given symbol"
    return len(resolve(sym))

# %% ../nbs/02_inspecttools.ipynb #a5dfbe8f
def symslice(
    sym: str,   # Dotted symbol path or "_last" for previous result
    start: int, # Starting index for slice
    end: int    # Ending index for slice
):
    "Returns the contents of the symbol from the given start to the end."
    try: return resolve(sym)[start:end]
    except Exception as e: return f'Error: {e}'

# %% ../nbs/02_inspecttools.ipynb #5fca4d70
def symsearch(
    sym:str,      # Dotted symbol path or "_last" for previous result
    term:str,     # Search term (exact string or regex pattern)
    regex:bool=True,  # If True, regex search; if False, exact match
    flags:int=0   # Regex flags (e.g., re.IGNORECASE)
):
    """Search contents of symbol, which is assumed to be str for regex, or iterable for non-regex.
    Regex mode returns (match, start, end) tuples; otherwise returns (item, index) tuples"""
    if regex: return str([(m.group(), m.start(), m.end()) for m in re.finditer(term, resolve(sym), flags)])
    else: return str([(x, i) for i, x in enumerate(resolve(sym)) if x == term])

# %% ../nbs/02_inspecttools.ipynb #02c09e1a
def symset(
    val: str  # Value to assign to _ai_sym
):
    "Set _ai_sym to the given value"
    _find_frame_dict('__msg_id')['_ai_sym'] = val

# %% ../nbs/02_inspecttools.ipynb #a8e90aa9
def run_code_interactive(
    code:str # Code to have user run
):
    """Insert code into user's dialog and request for the user to run it. Use other tools where possible, 
    but if they can not find needed information, *ALWAYS* use this instead of guessing or giving up.
    IMPORTANT: This tool is TERMINAL - after calling it, you MUST stop all tool usage 
    and wait for user response. Never call additional tools after this one."""
    add_msg('# Please run this:\n'+code, msg_type='code')
    return {'success': "CRITICAL: Message added to user dialog. STOP IMMEDIATELY. Do NOT call any more tools. Wait for user to run code and respond."}

# %% ../nbs/02_inspecttools.ipynb #a1a11745
_builtins = set(dir(builtins))

def _collapse(v): return v[0] if len(v) > 0 and all(x == v[0] for x in v) else v
def _process(vars): return {k: _collapse(v) for k,v in vars.items() if k not in _builtins}

# %% ../nbs/02_inspecttools.ipynb #ba9f1666
def tracetool(
    sym: str,  # Dotted symbol path of callable to run
    args: list=None,  # Positional args for callable (JSON values passed directly)
    kwargs: dict=None,  # Keyword args for callable (JSON values passed directly)
    target_func: str=None  # Dotted symbol path of function to trace (defaults to sym)
)->list[tuple[str, dict[str, tuple[int, dict[str, tuple|list]]]]]:  # List of (stack_str,trace_dict); trace_dict maps source snippets to (hit_count, variables), unchanged vars collapsed to single tuple
    """
    Trace execution using sys.monitoring (Python 3.12+), returning a list of per-call traces.

    Return:
    - list of length <= 10
    - one element per call to `target_func` (including recursion)
    - each element is: (stack_str, trace_dict)
        stack_str: call stack string (filtered so `fn` is the shallowest frame shown)
        trace_dict: {
        "<source snippet for AST-line>": ( hit_count, { "var": [ (type_name, truncated_repr), ... up to 10 ], ... } ),
        ...}

    Semantics:
    - "Line" means an AST-level line: separate statements (even if on one physical line via `;`).
    - Compound statements are keyed by their header only.
    - Unchanged variables → `('type', 'repr')` tuple, changed variables → `[('type', 'repr'), ...]` list.
    - Comprehensions are treated as a line node and are monitored, including inside the comprehension frame, with per-iteration snapshots.
    - Snapshots are recorded after each line finishes, so assignments show updated values.
    """
    from tracefunc import tracefunc
    target = resolve(target_func) if target_func else None
    traces = tracefunc(resolve(sym), *(args or []), target_func=target, **(kwargs or {}))
    return [(stack, {src: (hits, _process(vars)) for src, (hits, vars) in d.items()}) for stack, d in traces]

# %% ../nbs/02_inspecttools.ipynb #7c97567d
def fmt_trace(traces):
    "Format tracetool output as markdown"
    def fmt_val(val):
        if isinstance(val, tuple): return f"`{val[1]}`"
        return " → ".join(f"`{v[1]}`" for v in val)
    def fmt_vars(vars):
        if not vars: return ''
        return '<br>'.join(f"`{k}`: {fmt_val(v)}" for k, v in vars.items())
    def fmt_line(src, hits, vars):
        src_short = src.split('\n')[0][:60]
        if len(src.split('\n')[0]) > 60: src_short += '...'
        return f"| `{src_short}` | {hits} | {fmt_vars(vars)} |"
    parts = []
    for stack, trace in traces:
        if stack: parts.append(f"**Stack:**\n```\n{stack}\n```\n")
        parts.append("| Source | Hits | Variables |\n|--------|------|-----------|")
        parts.extend(fmt_line(src, hits, vars) for src, (hits, vars) in trace.items() if hits > 0)
        parts.append("")
    return '\n'.join(parts)

# %% ../nbs/02_inspecttools.ipynb #78c24a04
def inspect_tool_info():
    add_msg('Tools available from inspecttools: &`[symsrc,showsrc,gettype,getdir,doimport,getval,getnth,symlen,symslice,symsearch,symset,run_code_interactive,tracetool]`')
