Modules/_tree_utils.py

"""
Shared tree rendering utilities for filesystem-like output.

Provides functions to build and render tree structures with line connectors
(├──, └──, │) for visual hierarchy display.
"""

from __future__ import annotations

import os
from datetime import datetime
from typing import Callable, Optional


def build_tree(entries: list[tuple[str, dict]]) -> dict:
    """
    Build a nested tree structure from flat path entries.
    
    Args:
        entries: List of (path, metadata) tuples where path uses forward slashes.
                 Paths ending with '/' are treated as directories.
    
    Returns:
        Nested dict with "__files__" key for files at each level.
    
    Example:
        entries = [
            ("scripts/utils.py", {"size": 1234}),
            ("docs/readme.md", {"size": 500}),
        ]
        tree = build_tree(entries)
    """
    root: dict = {"__files__": []}
    
    for path, metadata in entries:
        parts = path.rstrip("/").split("/")
        is_dir = path.endswith("/")
        
        node = root
        for i, part in enumerate(parts[:-1]):
            if part not in node:
                node[part] = {"__files__": []}
            node = node[part]
        
        final = parts[-1]
        if is_dir:
            # Ensure directory exists
            if final not in node:
                node[final] = {"__files__": []}
            # Store directory metadata if provided
            if metadata:
                node[final]["__meta__"] = metadata
        else:
            # Add file
            node["__files__"].append((final, metadata))
    
    return root


def render_tree(
    node: dict,
    prefix: str = "",
    format_entry: Optional[Callable[[str, dict, bool], str]] = None,
) -> list[str]:
    """
    Render a tree with line connectors.
    
    Args:
        node: Nested dict from build_tree()
        prefix: Current line prefix for indentation
        format_entry: Optional callback to format each entry.
                      Signature: (name, metadata, is_dir) -> str
                      If None, uses default formatting.
    
    Returns:
        List of formatted lines.
    """
    result = []
    
    # Default formatter
    def default_format(name: str, meta: dict, is_dir: bool) -> str:
        if is_dir:
            return f"{name}/"
        size = meta.get("size")
        if size is not None:
            return f"{name}  ({_fmt_size(size)})"
        return name
    
    fmt = format_entry or default_format
    
    # Collect entries: subdirs first, then files
    entries = []
    subdirs = sorted(k for k in node.keys() if k not in ("__files__", "__meta__"))
    files_here = sorted(node.get("__files__", []), key=lambda x: x[0])
    
    for dirname in subdirs:
        dir_meta = node[dirname].get("__meta__", {})
        entries.append(("dir", dirname, node[dirname], dir_meta))
    for fname, fmeta in files_here:
        entries.append(("file", fname, None, fmeta))
    
    for i, entry in enumerate(entries):
        is_last = (i == len(entries) - 1)
        connector = "└── " if is_last else "├── "
        child_prefix = prefix + ("    " if is_last else "│   ")
        
        etype, name, subtree, meta = entry
        
        if etype == "dir":
            result.append(f"{prefix}{connector}{fmt(name, meta, True)}")
            result.extend(render_tree(subtree, child_prefix, format_entry))
        else:
            result.append(f"{prefix}{connector}{fmt(name, meta, False)}")
    
    return result


def _fmt_size(num_bytes: int) -> str:
    """Format byte size as human-readable string."""
    units = ["B", "KB", "MB", "GB"]
    size = float(num_bytes)
    for unit in units:
        if size < 1024.0:
            return f"{size:.1f} {unit}"
        size /= 1024.0
    return f"{size:.1f} TB"


def walk_and_build_tree(
    abs_path: str,
    *,
    show_hidden: bool = False,
    recursive: bool = False,
    max_entries: int = 100,
) -> tuple[dict, int, bool]:
    """
    Walk a directory and build a tree structure.
    
    Args:
        abs_path: Absolute path to directory
        show_hidden: Include hidden files/dirs (starting with '.')
        recursive: Recurse into subdirectories
        max_entries: Maximum entries before truncation
    
    Returns:
        (tree, total_entries, truncated)
    """
    entries: list[tuple[str, dict]] = []
    total = 0
    truncated = False
    
    for root, dirs, files in os.walk(abs_path):
        # Filter hidden
        if not show_hidden:
            dirs[:] = [d for d in dirs if not d.startswith('.')]
            files = [f for f in files if not f.startswith('.')]
        
        dirs.sort()
        files.sort()
        
        # Compute relative path from the listing root
        try:
            rel_root = os.path.relpath(root, abs_path)
        except Exception:
            rel_root = ""
        prefix = "" if rel_root == "." else rel_root.replace("\\", "/") + "/"
        
        # Add directories (with trailing slash to indicate dir)
        for d in dirs:
            p = os.path.join(root, d)
            try:
                mtime = datetime.fromtimestamp(os.path.getmtime(p)).strftime("%Y-%m-%d %H:%M")
            except Exception:
                mtime = "?"
            entries.append((f"{prefix}{d}/", {"mtime": mtime}))
            total += 1
            if total >= max_entries:
                truncated = True
                break
        
        if truncated:
            break
        
        # Add files
        for f in files:
            p = os.path.join(root, f)
            try:
                size = os.path.getsize(p)
                mtime = datetime.fromtimestamp(os.path.getmtime(p)).strftime("%Y-%m-%d %H:%M")
            except Exception:
                size, mtime = 0, "?"
            entries.append((f"{prefix}{f}", {"size": size, "mtime": mtime}))
            total += 1
            if total >= max_entries:
                truncated = True
                break
        
        if truncated:
            break
        
        if not recursive:
            break
    
    return build_tree(entries), total, truncated


def format_dir_listing(
    abs_path: str,
    display_path: str,
    *,
    show_hidden: bool = False,
    recursive: bool = False,
    max_entries: int = 100,
    fmt_size_fn: Optional[Callable[[int], str]] = None,
) -> str:
    """
    Format a directory listing as a visual tree.
    
    Args:
        abs_path: Absolute path to directory
        display_path: User-friendly path to show in header
        show_hidden: Include hidden files/dirs
        recursive: Recurse into subdirectories
        max_entries: Maximum entries before truncation
        fmt_size_fn: Optional custom size formatter (defaults to _fmt_size)
    
    Returns:
        Formatted string with tree output.
    """
    fmt_size = fmt_size_fn or _fmt_size
    
    tree, total, truncated = walk_and_build_tree(
        abs_path,
        show_hidden=show_hidden,
        recursive=recursive,
        max_entries=max_entries,
    )
    
    # Formatter with size + date
    def format_entry(name: str, meta: dict, is_dir: bool) -> str:
        mtime = meta.get("mtime", "")
        if is_dir:
            return f"{name}/  ({mtime})"
        size = meta.get("size", 0)
        return f"{name}  ({fmt_size(size)}, {mtime})"
    
    tree_lines = render_tree(tree, "    ", format_entry)
    
    header = f"Listing of {display_path}\nRoot: /\nEntries: {total}"
    if truncated:
        header += f"\n… Truncated at {max_entries} entries."
    
    lines = [header, "", "└── /"]
    lines.extend(tree_lines)
    
    return "\n".join(lines).strip()


__all__ = [
    "build_tree",
    "render_tree",
    "_fmt_size",
    "walk_and_build_tree",
    "format_dir_listing",
]