jtree - Interactive JSON structure explorer with jq selectors

jtree.py

"""
jtree - Interactive JSON structure explorer with jq selectors

Displays JSON structure as a navigable table. Press Enter to preview extracted
data, 'p' to print and exit, 'y' to copy selector, 'q' to quit.
"""
import json
import sys
import re
import argparse
import os
import curses
import subprocess
from typing import Any, Dict, List, Set, Tuple, NamedTuple, Optional


class TableRow(NamedTuple):
    """Represents a row in the output table."""
    tree: str
    example: str
    selector: str


def is_valid_jq_identifier(key: str) -> bool:
    """Check if key can be used as bare .key in jq (no brackets needed)."""
    return bool(re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', key))


def build_jq_path(parent_path: str, key: str, is_array_element: bool = False) -> str:
    """Build jq selector path for a child key."""
    if is_array_element:
        if is_valid_jq_identifier(key):
            return f"{parent_path}.{key}"
        else:
            return f'{parent_path}["{key}"]'
    else:
        if is_valid_jq_identifier(key):
            return f"{parent_path}.{key}" if parent_path != "." else f".{key}"
        else:
            return f'{parent_path}["{key}"]'


def get_type_name(value: Any) -> str:
    """Get friendly type name for a value."""
    if value is None:
        return "null"
    elif isinstance(value, bool):
        return "bool"
    elif isinstance(value, int):
        return "int"
    elif isinstance(value, float):
        return "float"
    elif isinstance(value, str):
        return "string"
    elif isinstance(value, list):
        return "array"
    elif isinstance(value, dict):
        return "object"
    else:
        return type(value).__name__


def merge_array_schemas(array: List[Any]) -> Tuple[Dict[str, int], Set[str]]:
    """
    Merge schemas of all objects in an array.
    Returns: (key_counts dict, set of all types seen)
    """
    key_counts: Dict[str, int] = {}
    types_seen: Set[str] = set()

    for item in array:
        types_seen.add(get_type_name(item))
        if isinstance(item, dict):
            for key in item.keys():
                key_counts[key] = key_counts.get(key, 0) + 1

    return key_counts, types_seen


def format_scalar_preview(value: Any, max_len: int = 30) -> str:
    """Format a scalar value for preview."""
    if value is None:
        return "null"
    elif isinstance(value, bool):
        return str(value).lower()
    elif isinstance(value, (int, float)):
        s = str(value)
        return s[:max_len] + "..." if len(s) > max_len else s
    elif isinstance(value, str):
        preview = value[:max_len] + "..." if len(value) > max_len else value
        return preview
    else:
        return str(value)[:max_len]


def build_tree_rows(
    data: Any,
    jq_path: str = ".",
    prefix: str = "",
    is_last: bool = True,
    level: int = 0,
    max_level: int = None,
    is_array_element: bool = False
) -> List[TableRow]:
    """Build tree representation rows for table output."""
    rows = []

    if max_level is not None and level >= max_level:
        return rows

    # Tree drawing characters (always ASCII for simplicity in TUI)
    last_connector = "+-- "
    mid_connector = "+-- "
    last_extension = "    "
    mid_extension = "|   "

    if isinstance(data, dict):
        keys = list(data.keys())
        for i, key in enumerate(keys):
            is_last_key = (i == len(keys) - 1)
            connector = last_connector if is_last_key else mid_connector
            extension = last_extension if is_last_key else mid_extension
            child_value = data[key]

            if isinstance(child_value, list):
                # Array - fold into parent key
                child_path = build_jq_path(jq_path, key, is_array_element) + "[]"
                tree_label = f"{prefix}{connector}{key}"
                example = f"[{len(child_value)} items]"
                rows.append(TableRow(tree_label, example, child_path))

                # Recurse into array elements (union)
                if len(child_value) > 0 and (max_level is None or level + 1 < max_level):
                    key_counts, types_seen = merge_array_schemas(child_value)

                    if key_counts:
                        # Array of objects - show union schema
                        object_count = sum(1 for item in child_value if isinstance(item, dict))
                        union_keys = list(key_counts.keys())

                        for j, union_key in enumerate(union_keys):
                            is_last_union = (j == len(union_keys) - 1)
                            count = key_counts[union_key]

                            key_display = union_key
                            if count < object_count:
                                key_display = f"{union_key}?"

                            # Get sample value
                            sample_value = None
                            for item in child_value:
                                if isinstance(item, dict) and union_key in item:
                                    sample_value = item[union_key]
                                    break

                            child_jq = build_jq_path(child_path, union_key, is_array_element=True)
                            union_connector = last_connector if is_last_union else mid_connector
                            union_extension = last_extension if is_last_union else mid_extension

                            tree_label = f"{prefix}{extension}{union_connector}{key_display}"

                            if isinstance(sample_value, dict):
                                example = f"{{{len(sample_value)} keys}}"
                            elif isinstance(sample_value, list):
                                example = f"[{len(sample_value)} items]"
                            else:
                                example = format_scalar_preview(sample_value)

                            rows.append(TableRow(tree_label, example, child_jq))

                            # Recurse into nested structures
                            if (isinstance(sample_value, (dict, list)) and
                                (max_level is None or level + 2 < max_level)):
                                nested_prefix = prefix + extension + union_extension
                                nested_rows = build_tree_rows(
                                    sample_value,
                                    child_jq,
                                    nested_prefix,
                                    is_last_union,
                                    level + 2,
                                    max_level,
                                    is_array_element=True
                                )
                                rows.extend(nested_rows)

                    elif len(types_seen) > 0:
                        # Array of mixed/scalar types
                        type_list = "|".join(sorted(types_seen))
                        tree_label = f"{prefix}{extension}{last_connector}[*]"
                        rows.append(TableRow(tree_label, type_list, child_path))

            elif isinstance(child_value, dict):
                # Nested object
                child_path = build_jq_path(jq_path, key, is_array_element)
                tree_label = f"{prefix}{connector}{key}"
                example = f"{{{len(child_value)} keys}}"
                rows.append(TableRow(tree_label, example, child_path))

                if max_level is None or level + 1 < max_level:
                    child_rows = build_tree_rows(
                        child_value,
                        child_path,
                        prefix + extension,
                        is_last_key,
                        level + 1,
                        max_level,
                        is_array_element
                    )
                    rows.extend(child_rows)

            else:
                # Scalar value
                child_path = build_jq_path(jq_path, key, is_array_element)
                tree_label = f"{prefix}{connector}{key}"
                example = format_scalar_preview(child_value)
                rows.append(TableRow(tree_label, example, child_path))

    elif isinstance(data, list):
        # Bare array node
        if len(data) == 0:
            empty_conn = last_connector
            rows.append(TableRow(f"{prefix}{empty_conn}(empty)", "", jq_path))
        else:
            key_counts, types_seen = merge_array_schemas(data)

            if key_counts:
                object_count = sum(1 for item in data if isinstance(item, dict))
                union_keys = list(key_counts.keys())

                for i, union_key in enumerate(union_keys):
                    is_last_union = (i == len(union_keys) - 1)
                    union_connector = last_connector if is_last_union else mid_connector
                    union_extension = last_extension if is_last_union else mid_extension
                    count = key_counts[union_key]

                    key_display = union_key
                    if count < object_count:
                        key_display = f"{union_key}?"

                    sample_value = None
                    for item in data:
                        if isinstance(item, dict) and union_key in item:
                            sample_value = item[union_key]
                            break

                    child_jq = build_jq_path(jq_path, union_key, is_array_element=True)
                    tree_label = f"{prefix}{union_connector}{key_display}"

                    if isinstance(sample_value, dict):
                        example = f"{{{len(sample_value)} keys}}"
                    elif isinstance(sample_value, list):
                        example = f"[{len(sample_value)} items]"
                    else:
                        example = format_scalar_preview(sample_value)

                    rows.append(TableRow(tree_label, example, child_jq))

                    if (isinstance(sample_value, (dict, list)) and
                        (max_level is None or level + 1 < max_level)):
                        nested_rows = build_tree_rows(
                            sample_value,
                            child_jq,
                            prefix + union_extension,
                            is_last_union,
                            level + 1,
                            max_level,
                            is_array_element=True
                        )
                        rows.extend(nested_rows)
            else:
                type_list = "|".join(sorted(types_seen))
                rows.append(TableRow(f"{prefix}{last_connector}[*]", type_list, jq_path))

    return rows


def extract_with_selector(data: Any, selector: str, debug_file=None) -> Optional[str]:
    """
    Extract data using a jq-style selector path. Returns formatted JSON string.
    Handles: .key, .key1.key2, .key[], .key[].subkey, .["weird key"]
    """
    try:
        # Parse the selector into path segments
        parts = []
        i = 0
        s = selector

        # Skip leading dot
        if s.startswith('.'):
            s = s[1:]

        if not s:
            # Selector is just "." — return entire document
            return json.dumps(data, indent=2)

        while s:
            if s.startswith('['):
                # Could be ["key"] or []
                if s.startswith('[]'):
                    parts.append(('iterate', None))
                    s = s[2:]
                elif s.startswith('["'):
                    end = s.index('"]')
                    key = s[2:end]
                    parts.append(('key', key))
                    s = s[end + 2:]
                else:
                    # numeric index
                    end = s.index(']')
                    idx = int(s[1:end])
                    parts.append(('index', idx))
                    s = s[end + 1:]
            elif s.startswith('.'):
                s = s[1:]
            else:
                # Bare key - read until next . or [
                match = re.match(r'^([A-Za-z_][A-Za-z0-9_]*)', s)
                if match:
                    parts.append(('key', match.group(1)))
                    s = s[match.end():]
                else:
                    break

        if debug_file:
            with open(debug_file, 'a') as f:
                f.write(f"extract: selector={selector!r}, parts={parts}\n")

        # Walk the data
        def walk(current, parts_remaining):
            if not parts_remaining:
                return [current]

            part_type, part_val = parts_remaining[0]
            rest = parts_remaining[1:]

            if part_type == 'key':
                if isinstance(current, dict) and part_val in current:
                    return walk(current[part_val], rest)
                else:
                    return [None]
            elif part_type == 'index':
                if isinstance(current, list) and part_val < len(current):
                    return walk(current[part_val], rest)
                else:
                    return [None]
            elif part_type == 'iterate':
                if isinstance(current, list):
                    results = []
                    for item in current:
                        results.extend(walk(item, rest))
                    return results
                else:
                    return [None]
            return [None]

        results = walk(data, parts)

        if debug_file:
            with open(debug_file, 'a') as f:
                f.write(f"extract: got {len(results)} results\n")

        if not results:
            return None

        # Format output
        lines = []
        for r in results:
            lines.append(json.dumps(r, indent=2))

        return '\n'.join(lines)

    except Exception as e:
        if debug_file:
            with open(debug_file, 'a') as f:
                f.write(f"extract exception: {e}\n")
        return None


def copy_to_clipboard(text: str) -> bool:
    """Copy text to clipboard. Returns True on success."""
    commands = [
        ['pbcopy'],           # macOS
        ['xclip', '-selection', 'clipboard'],  # Linux (X11)
        ['xsel', '--clipboard', '--input'],    # Linux (X11)
        ['wl-copy'],          # Linux (Wayland)
        ['clip'],             # Windows
    ]

    for cmd in commands:
        try:
            proc = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )
            proc.communicate(input=text.encode())
            if proc.returncode == 0:
                return True
        except FileNotFoundError:
            continue

    return False


def supports_unicode_curses() -> bool:
    """Check if terminal supports Unicode in curses context."""
    # Check encoding
    import locale
    try:
        encoding = locale.getpreferredencoding()
        if 'UTF-8' in encoding.upper() or 'UTF8' in encoding.upper():
            return True
    except:
        pass

    # Check environment variables
    for var in ['LANG', 'LC_ALL', 'LC_CTYPE']:
        val = os.environ.get(var, '')
        if 'UTF-8' in val.upper() or 'UTF8' in val.upper():
            return True

    return False


def draw_table(stdscr, rows: List[TableRow], selected_idx: int, scroll_offset: int,
               preview_data: Optional[str] = None, status_msg: str = ""):
    """Draw the table in the curses window."""
    height, width = stdscr.getmaxyx()

    # Reserve space for preview panel and status bar
    preview_height = min(10, height // 3) if preview_data else 0
    status_height = 1
    table_height = height - preview_height - status_height - 4  # 4 for borders/header

    stdscr.clear()

    # Calculate column widths
    tree_width = min(max(len(row.tree) for row in rows) + 2, width // 3)
    example_width = min(max(len(row.example) for row in rows) + 2, width // 4)
    selector_width = width - tree_width - example_width - 4

    # Choose characters based on Unicode support
    use_unicode = supports_unicode_curses()
    h_line = "─" if use_unicode else "-"
    v_line = "│" if use_unicode else "|"

    # Draw header
    try:
        stdscr.addstr(0, 0, h_line * (width - 1), curses.A_DIM)
        header = f" Tree{' ' * (tree_width - 5)}{v_line} Example{' ' * (example_width - 8)}{v_line} jq Selector"
        stdscr.addstr(1, 0, header[:width - 1])
        stdscr.addstr(2, 0, h_line * (width - 1), curses.A_DIM)
    except curses.error:
        pass

    # Draw rows
    visible_rows = rows[scroll_offset:scroll_offset + table_height]
    for i, row in enumerate(visible_rows):
        row_idx = scroll_offset + i
        y = 3 + i

        if y >= height - preview_height - status_height - 1:
            break

        # Highlight selected row
        attr = curses.A_REVERSE if row_idx == selected_idx else curses.A_NORMAL

        # Truncate fields to fit
        tree_text = row.tree[:tree_width - 1].ljust(tree_width - 1)
        example_text = row.example[:example_width - 1].ljust(example_width - 1)
        selector_text = row.selector[:selector_width - 1]

        line = f"{tree_text}{v_line}{example_text}{v_line}{selector_text}"
        try:
            stdscr.addstr(y, 0, line[:width - 1], attr)
        except curses.error:
            pass

    # Draw preview panel if active
    if preview_data:
        preview_start = height - preview_height - status_height - 1
        try:
            stdscr.addstr(preview_start, 0, h_line * (width - 1), curses.A_DIM)
            stdscr.addstr(preview_start + 1, 0, " Preview:", curses.A_BOLD)
        except curses.error:
            pass

        preview_lines = preview_data.split('\n')
        for i, line in enumerate(preview_lines[:preview_height - 2]):
            try:
                stdscr.addstr(preview_start + 2 + i, 0, line[:width - 1])
            except curses.error:
                pass

    # Draw status bar
    status_y = height - 1
    try:
        stdscr.addstr(status_y, 0, h_line * (width - 1), curses.A_DIM)
    except curses.error:
        pass

    if status_msg:
        try:
            stdscr.addstr(status_y, 0, f" {status_msg}"[:width - 1], curses.A_BOLD)
        except curses.error:
            pass
    else:
        help_text = " Up/Down: Navigate | Enter: Preview | p: Print & Exit | y: Copy Selector | q: Quit"
        try:
            stdscr.addstr(status_y, 0, help_text[:width - 1], curses.A_DIM)
        except curses.error:
            pass

    stdscr.refresh()


def run_tui(stdscr, rows: List[TableRow], data: Any, debug_file=None):
    """Run the interactive TUI."""
    def debug_log(msg):
        if debug_file:
            with open(debug_file, 'a') as f:
                import datetime
                f.write(f"[{datetime.datetime.now()}] {msg}\n")
                f.flush()

    debug_log("TUI started")

    try:
        curses.curs_set(0)  # Hide cursor
        debug_log("Cursor hidden")
    except Exception as e:
        debug_log(f"Could not hide cursor: {e}")
        pass

    stdscr.keypad(True)  # Enable arrow keys
    debug_log("Keypad enabled")

    curses.use_default_colors()
    stdscr.nodelay(False)  # Block on getch() - wait for key press
    debug_log("Nodelay set to False (blocking mode)")

    selected_idx = 0
    scroll_offset = 0
    preview_data = None
    status_msg = ""

    debug_log(f"Starting main loop with {len(rows)} rows")

    while True:
        height, width = stdscr.getmaxyx()
        preview_height = min(10, height // 3) if preview_data else 0
        table_height = height - preview_height - 5

        # Auto-scroll to keep selected row visible
        if selected_idx < scroll_offset:
            scroll_offset = selected_idx
        elif selected_idx >= scroll_offset + table_height:
            scroll_offset = selected_idx - table_height + 1

        draw_table(stdscr, rows, selected_idx, scroll_offset, preview_data, status_msg)

        debug_log("About to call getch()")
        key = stdscr.getch()
        debug_log(f"Received key: {key} (chr: {chr(key) if 0 < key < 128 else 'N/A'})")

        # Clear status message on next keypress
        if status_msg:
            status_msg = ""

        if key == ord('q') or key == ord('Q'):
            debug_log("Quit key pressed")
            break
            break
        elif key == curses.KEY_UP or key == ord('k'):
            debug_log(f"Up key pressed (KEY_UP={curses.KEY_UP})")
            if selected_idx > 0:
                selected_idx -= 1
                preview_data = None  # Clear preview when navigating
        elif key == curses.KEY_DOWN or key == ord('j'):
            debug_log(f"Down key pressed (KEY_DOWN={curses.KEY_DOWN})")
            if selected_idx < len(rows) - 1:
                selected_idx += 1
                preview_data = None  # Clear preview when navigating
        elif key == curses.KEY_PPAGE:  # Page Up
            selected_idx = max(0, selected_idx - table_height)
            preview_data = None
        elif key == curses.KEY_NPAGE:  # Page Down
            selected_idx = min(len(rows) - 1, selected_idx + table_height)
            preview_data = None
        elif key == ord('g'):  # Go to top
            selected_idx = 0
            scroll_offset = 0
            preview_data = None
        elif key == ord('G'):  # Go to bottom
            selected_idx = len(rows) - 1
            preview_data = None
        elif key == ord('\n') or key == curses.KEY_ENTER or key == 10 or key == 13:
            debug_log(f"Enter key pressed (KEY_ENTER={curses.KEY_ENTER})")
            # Preview selected selector
            selector = rows[selected_idx].selector
            result = extract_with_selector(data, selector, debug_file)
            if result is not None:
                preview_data = result
                status_msg = "Preview loaded"
                debug_log("Preview generated successfully")
            else:
                status_msg = "Error: Could not extract data"
                debug_log("Failed to extract data with jq")
        elif key == ord('p') or key == ord('P'):
            # Print and exit
            selector = rows[selected_idx].selector
            result = extract_with_selector(data, selector, debug_file)
            if result is not None:
                return ('print', result)
            else:
                status_msg = "Error: Could not extract data"
        elif key == ord('y') or key == ord('Y'):
            debug_log("Copy key pressed")
            # Copy selector to clipboard
            selector = rows[selected_idx].selector
            if copy_to_clipboard(selector):
                status_msg = f"Copied: {selector}"
                debug_log(f"Copied to clipboard: {selector}")
            else:
                status_msg = "Error: Clipboard unavailable"
                debug_log("Clipboard copy failed")
        else:
            debug_log(f"Unhandled key: {key}")

    debug_log("Exiting TUI")
    return None


def main():
    parser = argparse.ArgumentParser(
        description="Interactive JSON structure explorer with jq selectors",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Interactive Mode (default):
  ↑/↓ or j/k    Navigate rows
  Enter         Preview extracted data
  p             Print to stdout and exit
  y             Copy jq selector to clipboard
  q             Quit

Non-Interactive Mode:
  --print       Print table and exit (no TUI)

Examples:
  curl https://api.example.com/data | jtree
  cat response.json | jtree -L 3
  curl https://api.example.com/data | jtree --print | grep users
"""
    )
    parser.add_argument(
        '-L', '--level',
        type=int,
        metavar='N',
        help='Maximum depth to display'
    )
    parser.add_argument(
        '--print',
        action='store_true',
        help='Print table and exit (no TUI)'
    )
    parser.add_argument(
        '--debug',
        metavar='FILE',
        type=str,
        help='Write debug log to FILE (e.g., /tmp/jtree.log)'
    )
    parser.add_argument(
        'file',
        nargs='?',
        type=argparse.FileType('r'),
        default=sys.stdin,
        help='JSON file to read (default: stdin)'
    )

    args = parser.parse_args()

    try:
        data = json.load(args.file)

        # Build rows
        rows = []
        if isinstance(data, dict):
            count = len(data)
            plural = "key" if count == 1 else "keys"
            rows.append(TableRow(".", f"{{{count} {plural}}}", "."))
            rows.extend(build_tree_rows(
                data,
                jq_path=".",
                prefix="",
                level=0,
                max_level=args.level
            ))
        elif isinstance(data, list):
            count = len(data)
            plural = "item" if count == 1 else "items"
            rows.append(TableRow(".", f"[{count} {plural}]", ".[]"))
            if count > 0:
                rows.extend(build_tree_rows(
                    data,
                    jq_path=".[]",
                    prefix="",
                    level=0,
                    max_level=args.level,
                    is_array_element=True
                ))
        else:
            example = format_scalar_preview(data)
            rows.append(TableRow(".", example, "."))

        # Print mode - just show the table
        if args.print:
            # Simple table output for non-interactive mode
            tree_width = max(len(row.tree) for row in rows) + 2
            example_width = max(len(row.example) for row in rows) + 2
            selector_width = max(len(row.selector) for row in rows) + 2

            # Check if we should use Unicode for print mode
            import locale
            use_unicode = False
            try:
                encoding = locale.getpreferredencoding()
                if 'UTF-8' in encoding.upper() or 'UTF8' in encoding.upper():
                    use_unicode = True
            except:
                pass

            h_line = "─" if use_unicode else "-"
            v_line = "│" if use_unicode else "|"

            print(h_line * (tree_width + example_width + selector_width + 2))
            print(f" Tree{' ' * (tree_width - 5)}{v_line} Example{' ' * (example_width - 8)}{v_line} jq Selector")
            print(h_line * (tree_width + example_width + selector_width + 2))

            for row in rows:
                tree_text = row.tree.ljust(tree_width - 1)
                example_text = row.example.ljust(example_width - 1)
                selector_text = row.selector
                print(f"{tree_text}{v_line}{example_text}{v_line}{selector_text}")

            print(h_line * (tree_width + example_width + selector_width + 2))
        else:
            # Interactive TUI mode
            if args.debug:
                with open(args.debug, 'w') as f:
                    f.write("=== jtree debug log ===\n")
                    f.write(f"About to start TUI with {len(rows)} rows\n")
                    f.flush()

            # Curses needs fd 0 to be a TTY. When stdin is piped,
            # swap the C-level fd 0 to point at /dev/tty.
            saved_fd = None
            if not os.isatty(0):
                if args.debug:
                    with open(args.debug, 'a') as f:
                        f.write("fd 0 is not a TTY, swapping to /dev/tty\n")
                try:
                    tty_fd = os.open('/dev/tty', os.O_RDWR)
                    saved_fd = os.dup(0)
                    os.dup2(tty_fd, 0)
                    os.close(tty_fd)
                except OSError as e:
                    print("Error: Cannot open /dev/tty for interactive mode.", file=sys.stderr)
                    print("Use --print for non-interactive output.", file=sys.stderr)
                    sys.exit(1)

            try:
                result = curses.wrapper(run_tui, rows, data, args.debug)
            except Exception as e:
                if args.debug:
                    with open(args.debug, 'a') as f:
                        f.write(f"ERROR in curses.wrapper: {e}\n")
                        import traceback
                        f.write(traceback.format_exc())
                raise
            finally:
                # Restore original fd 0
                if saved_fd is not None:
                    os.dup2(saved_fd, 0)
                    os.close(saved_fd)

            if result and result[0] == 'print':
                print(result[1])

    except json.JSONDecodeError as e:
        print(f"Error: Invalid JSON: {e}", file=sys.stderr)
        sys.exit(1)
    except BrokenPipeError:
        sys.exit(0)
    except KeyboardInterrupt:
        sys.exit(0)


if __name__ == "__main__":
    main()