# BVH, May 2026.
# Generic logging / pretty-print helpers (non-a4d-specific).

import numpy as np
import torch
from rich import print as rich_print
from rich.panel import Panel

try:
    import lovely_tensors
    lovely_tensors.monkey_patch()
except Exception:
    pass


def _is_rank0():
    if torch.distributed.is_initialized():
        return torch.distributed.get_rank() == 0
    return True


def _is_string_leaf(val):
    """Treat strings and (nested) string lists as leaves, not containers to recurse into."""
    if isinstance(val, str):
        return True
    if isinstance(val, (list, tuple)) and len(val) > 0:
        return _is_string_leaf(val[0])
    return False


def _summarize_value(v, max_str=80):
    """Compact one-line summary of a value. Defensive: any per-value failure
    falls back to a type name so this is safe to call from error-inspection paths."""
    try:
        if isinstance(v, torch.Tensor):
            # `.v` is added by lovely_tensors.monkey_patch(); fall back if unpatched.
            if hasattr(v, 'v') and v.numel() <= 256:
                return str(v.v).split('\n')[0]
            return str(v)
        if isinstance(v, np.ndarray):
            from lovely_numpy import lo
            return str(lo(v).v).split('\n')[0] if v.size <= 256 else str(lo(v))
        if isinstance(v, (list, tuple)):
            t = type(v).__name__
            if len(v) == 0:
                return f'{t}(0)'
            sample = _summarize_value(v[0])
            return f'{t}({len(v)}) [{sample}, ...]'
        if isinstance(v, str):
            s = repr(v)
            return s if len(s) <= max_str else s[:max_str] + '...'
        if isinstance(v, (int, float, bool, type(None))):
            return repr(v)
        return type(v).__name__
    except Exception as e:
        return f'<{type(v).__name__} summary failed: {e!r}>'


def summarize_dict(d, prefix='', depth=0, max_depth=4, rich=True):
    """Recursively summarize a nested dict into compact lines.
    :param rich: if True, embed rich markup (e.g. [cyan]...[/]); if False, emit
        plain text (safe for loguru / file sinks).
    Returns a list of strings."""
    k_open, k_close = ('[cyan]', '[/]') if rich else ('', '')
    d_open, d_close = ('[dim]', '[/]') if rich else ('', '')
    lines = []
    if not isinstance(d, dict) or depth >= max_depth:
        lines.append(f'{prefix}{_summarize_value(d)}')
        return lines
    indent = '  ' * depth
    for k, v in d.items():
        if isinstance(v, dict) and len(v) > 0:
            vals = list(v.values())
            if len(vals) > 3 and all(type(x) == type(vals[0]) for x in vals):
                sample_key = next(iter(v))
                summary = _summarize_value(vals[0])
                lines.append(f'{indent}{prefix}{k_open}{k}{k_close}: '
                             f'{d_open}dict({len(v)}){d_close} '
                             f'{{{sample_key}: {summary}, ...}}')
            else:
                lines.append(f'{indent}{prefix}{k_open}{k}{k_close}:')
                lines.extend(summarize_dict(v, prefix='', depth=depth + 1,
                                            max_depth=max_depth, rich=rich))
        else:
            lines.append(f'{indent}{prefix}{k_open}{k}{k_close}: {_summarize_value(v)}')
    return lines


def format_batch_summary(data_batch, label='data_batch'):
    """Return a plain-text summary of a data_batch dict, without rich markup."""
    body = '\n'.join(summarize_dict(data_batch, rich=False))
    return f'=== {label} ===\n{body}'


def print_batch_summary(data_batch, label='data_batch'):
    """Print a compact overview of a data_batch dict. Rank 0 only."""
    if not _is_rank0():
        return
    lines = summarize_dict(data_batch, rich=True)
    rich_print(Panel('\n'.join(lines), title=f'[bold]{label}[/]',
                     border_style='blue', expand=False))


def fmt_path(key_path):
    """Render a key chain as 'k1->k2->k3' for diagnostic messages."""
    if not key_path:
        return '(root)'
    return '->'.join(repr(k) for k in key_path)


def safe_repr(v, max_len=240):
    try:
        r = repr(v)
    except Exception as e:
        return f'<repr failed: {type(e).__name__}: {e}>'
    if len(r) > max_len:
        return r[:max_len] + f'... <truncated, total len={len(r)}>'
    return r


def to_json_safe(obj, max_numel=4096, max_keys=64):
    '''
    Recursively convert a (possibly nested) structure to JSON-serializable form: small tensors/arrays
    become lists, large ones become a {'_tensor': shape, 'dtype': ...} stub; dicts with more than
    max_keys entries are summarized (e.g. per-(t,cam) camera/rgb dicts hold hundreds of entries, which
    would bloat the file) to {'_dict_len': N, '_sample': {first few}}; Camera-like objects expose their
    Tcw (+ K); anything else falls back to repr(). Pairs with io.save_json for dumping batch dicts.
    '''
    if torch.is_tensor(obj):
        return obj.detach().cpu().tolist() if obj.numel() <= max_numel \
            else {'_tensor': list(obj.shape), 'dtype': str(obj.dtype)}
    if isinstance(obj, np.ndarray):
        return obj.tolist() if obj.size <= max_numel \
            else {'_tensor': list(obj.shape), 'dtype': str(obj.dtype)}
    if isinstance(obj, np.integer):
        return int(obj)
    if isinstance(obj, np.floating):
        return float(obj)
    if isinstance(obj, dict):
        if len(obj) > max_keys:
            keys = list(obj.keys())
            return {'_dict_len': len(obj),
                    '_sample': {str(k): to_json_safe(obj[k], max_numel, max_keys) for k in keys[:3]}}
        return {str(k): to_json_safe(v, max_numel, max_keys) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [to_json_safe(v, max_numel, max_keys) for v in obj]
    if isinstance(obj, (str, int, float, bool)) or obj is None:
        return obj
    if hasattr(obj, 'Tcw'):  # Camera-like: keep extrinsics (+ intrinsics) as numbers
        cam = {'Tcw': to_json_safe(obj.Tcw, max_numel, max_keys)}
        if hasattr(obj, 'K'):
            cam['K'] = to_json_safe(obj.K, max_numel, max_keys)
        return cam
    return repr(obj)


def dump_batch(batch, max_depth=4, max_items=8):
    """Recursive structured dump of a list-like batch. Tensors / arrays render
    with lovely reprs (this module monkey-patches them on import).
    String lists are treated as leaves (not recursed into)."""
    try:
        lines = [f'batch: type={type(batch).__name__}, len={len(batch)}']
        for i, item in enumerate(batch):
            if i >= max_items:
                lines.append(f'  ... ({len(batch) - max_items} more elements truncated)')
                break
            lines.append(f'  [{i}]:')
            _dump_item_into(item, lines, indent='    ', depth=0,
                            max_depth=max_depth, max_items=max_items)
        return '\n'.join(lines)
    except Exception as e:
        return f'<batch dump failed: {type(e).__name__}: {e}>'


def _dump_item_into(item, lines, indent, depth, max_depth, max_items):
    if depth > max_depth:
        lines.append(f'{indent}... (truncated at depth {max_depth})')
        return
    if isinstance(item, dict):
        for j, (k, v) in enumerate(item.items()):
            if j >= max_items:
                lines.append(f'{indent}... ({len(item) - max_items} more keys truncated)')
                break
            if isinstance(v, dict) or (
                isinstance(v, (list, tuple)) and not _is_string_leaf(v)
            ):
                lines.append(f'{indent}{k!r}:')
                _dump_item_into(v, lines, indent + '  ', depth + 1,
                                max_depth, max_items)
            else:
                lines.append(f'{indent}{k!r}: {safe_repr(v)}')
    elif isinstance(item, (list, tuple)) and not _is_string_leaf(item):
        for j, v in enumerate(item):
            if j >= max_items:
                lines.append(f'{indent}... ({len(item) - max_items} more elements truncated)')
                break
            if isinstance(v, dict) or (
                isinstance(v, (list, tuple)) and not _is_string_leaf(v)
            ):
                lines.append(f'{indent}[{j}]:')
                _dump_item_into(v, lines, indent + '  ', depth + 1,
                                max_depth, max_items)
            else:
                lines.append(f'{indent}[{j}]: {safe_repr(v)}')
    else:
        lines.append(f'{indent}{safe_repr(item)}')