"""
LARS Mermaid Visualizer

Generates execution flow diagrams from Echo history showing:
- Cascade structure with nested cells
- Cell-to-cell handoffs
- Takes (Tree of Thought) with parallel attempts and winner selection
- Reforge (iterative refinement) with sequential steps
- Sub-cascades as nested containers
- Wards as validation checkpoints
"""
import json
import re
import os
import sys
import glob
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, List, Optional, Tuple
from dataclasses import dataclass, field
from .echo import Echo
from .config import get_config
from .console_style import S, styled_print


@dataclass
class ExecutionNode:
    """Represents a node in the execution tree."""
    id: str
    node_type: str  # cascade, cell, turn, tool, takes, reforge, etc.
    name: str
    role: str = ""
    content: str = ""
    parent_id: Optional[str] = None
    children: List['ExecutionNode'] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    # Take/Reforge specific
    take_index: Optional[int] = None
    is_winner: bool = False
    reforge_step: Optional[int] = None


def sanitize_label(content: Any, max_length: int = 50) -> str:
    """Sanitize content for Mermaid labels.

    Mermaid stateDiagram labels have strict parsing rules.
    Special characters that cause parse errors must be removed or escaped.

    CRITICAL: Colons (:) in label content break parsing because
    Mermaid uses `state_id : description` syntax.
    """
    if isinstance(content, (list, dict)):
        try:
            if isinstance(content, list):
                return f"[{len(content)} items]"
            return f"{{{len(content)} keys}}"
        except:
            pass

    s = str(content)

    # Remove/escape special Mermaid characters that break parsing
    s = s.replace('`', '')       # Backticks break parsing (```)
    s = s.replace('"', "'")      # Double quotes
    s = s.replace('\n', ' ')     # Newlines
    s = s.replace('#', '')       # Hash (comment start)
    s = s.replace('<', '')       # Angle brackets (HTML-like)
    s = s.replace('>', '')
    s = s.replace('{', '(')      # Curly braces (state block syntax)
    s = s.replace('}', ')')
    s = s.replace('[', '(')      # Square brackets (can conflict with links)
    s = s.replace(']', ')')
    s = s.replace(';', ',')      # Semicolons (statement separator)
    s = s.replace('--', '-')     # Double dash (arrow syntax)
    s = s.replace('->', '-')     # Arrow syntax
    s = s.replace(':', '-')      # Colons break state descriptions!

    # Collapse whitespace
    s = ' '.join(s.split())
    # Truncate
    if len(s) > max_length:
        s = s[:max_length] + "..."
    return s


def safe_id(trace_id: str) -> str:
    """Create a safe Mermaid node ID from trace ID."""
    return "n_" + trace_id.replace("-", "")[:12]


def _validate_mermaid_syntax(content: str) -> Tuple[bool, str]:
    """
    Validate Mermaid syntax using CLI.
    Falls back to basic regex checks if CLI unavailable.

    Returns:
        (is_valid, error_message)
    """
    # Try CLI validation first
    try:
        temp_file = Path("/tmp/lars_mermaid_check.mmd")
        temp_file.write_text(content)

        result = subprocess.run(
            ["mmdc", "-i", str(temp_file), "-o", "/tmp/lars_mermaid_out.svg", "--quiet"],
            capture_output=True,
            text=True,
            timeout=3
        )

        if result.returncode == 0:
            return True, ""
        else:
            return False, result.stderr or "Mermaid CLI validation failed"

    except FileNotFoundError:
        # CLI not installed - do basic sanity checks
        return _basic_mermaid_checks(content)
    except subprocess.TimeoutExpired:
        return False, "Mermaid validation timeout (possible infinite loop in syntax)"
    except Exception as e:
        # Don't fail hard - just warn
        print(f"[Warning] Mermaid validation skipped: {e}")
        return True, ""  # Assume valid if can't check


def _basic_mermaid_checks(content: str) -> Tuple[bool, str]:
    """Basic regex-based syntax checks when CLI unavailable.

    NOTE: We intentionally do NOT validate bracket balance because:
    1. Mermaid label content can contain arbitrary text (JSON, code, etc.)
    2. Labels like `label : T1 [{'tool': '...}]` legitimately have unbalanced brackets
    3. Structural syntax (nested states) is generated by code, not humans
    4. Real validation happens via mmdc CLI when available
    """
    lines = content.strip().split('\n')

    if not lines:
        return False, "Empty diagram"

    # Check for common syntax errors
    first_line = lines[0].strip()
    valid_starts = ['graph', 'flowchart', 'sequenceDiagram', 'stateDiagram', 'classDiagram',
                    'erDiagram', 'journey', 'gantt', 'pie', 'gitGraph', 'mindmap']

    if not any(first_line.startswith(start) for start in valid_starts):
        return False, f"Invalid diagram type: {first_line}"

    # Only validate STRUCTURAL curly braces (nested states), not label content
    # Pattern: "state name {" opens, "}" alone closes
    import re

    structural_opens = 0
    structural_closes = 0

    for line in lines:
        stripped = line.strip()
        # Skip comment lines
        if stripped.startswith('%%'):
            continue

        # Count structural state blocks: "state foo {" or "state foo {"
        if re.match(r'^state\s+\S+.*\{\s*$', stripped):
            structural_opens += 1
        # Count closing braces that are alone (end of nested state)
        elif stripped == '}':
            structural_closes += 1

    if structural_opens != structural_closes:
        return False, f"Unbalanced nested states: {structural_opens} opens vs {structural_closes} closes"

    return True, ""


def _log_invalid_mermaid(content: str, error: str, context: Optional[Dict], output_path: str):
    """Log invalid Mermaid for later debugging"""
    config = get_config()

    log_dir = Path(config.graph_dir) / "mermaid_failures"
    log_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    filename = Path(output_path).stem
    log_file = log_dir / f"{filename}_{timestamp}.json"

    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "original_path": output_path,
        "error": error,
        "mermaid_content": content,
        "context": context or {},
        "content_stats": {
            "line_count": len(content.split('\n')),
            "char_count": len(content),
            "has_takes": "fork" in content.lower(),
            "has_reforge": "reforge" in content.lower(),
        }
    }

    log_file.write_text(json.dumps(log_entry, indent=2))
    #print(f"📝 Invalid Mermaid logged: {log_file}")


def validate_and_write_mermaid(
    mermaid_content: str,
    output_path: str,
    source_context: Optional[Dict] = None
) -> Tuple[bool, str]:
    """
    Validate Mermaid content before writing to file.
    If invalid, log to failures directory and write anyway (with warning comment).

    Returns:
        (success, path_written)
    """
    is_valid, error = _validate_mermaid_syntax(mermaid_content)

    if not is_valid:
        # Log the failure for review
        _log_invalid_mermaid(mermaid_content, error, source_context, output_path)

        # Write with warning comment so it's clear it's invalid
        safe_content = f"""%%{{init: {{'theme':'base'}}}}%%
%% WARNING: This diagram failed validation
%% Error: {error}
%% Generated: {datetime.now().isoformat()}

{mermaid_content}
"""
        with open(output_path, "w") as f:
            f.write(safe_content)

        #print(f"⚠️  Invalid Mermaid diagram written to {output_path} (see mermaid_failures/ for details)")
        return False, output_path

    # Valid - write normally
    with open(output_path, "w") as f:
        f.write(mermaid_content)

    # Try to render to console (silently fails if not supported)
    _maybe_render_mermaid_to_console(mermaid_content)

    return True, output_path


def _maybe_render_mermaid_to_console(mermaid_content: str) -> None:
    """
    Attempt to render Mermaid diagram to console using mermaid-ascii.
    Silently fails if mermaid-ascii is not available or rendering fails.
    """
    # Check env var to disable
    if not os.environ.get("LARS_SHOW_CLI_MERMAID", "true").lower() in ("true", "1"):
        return

    # Check if stdout is a TTY
    if not sys.stdout.isatty():
        return

    try:
        # Try to run mermaid-ascii
        result = subprocess.run(
            ["mermaid-ascii"],
            input=mermaid_content,
            capture_output=True,
            text=True,
            timeout=5  # Don't hang
        )

        if result.returncode == 0 and result.stdout:
            # Successfully rendered
            styled_print(f"\n{S.CHART} Mermaid Diagram:")
            print(result.stdout)
            print("─" * 60 + "\n")
    except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
        # Silently fail - mermaid-ascii not installed, timed out, or other error
        pass


def load_sub_cascade_mermaid(sub_session_id: str, graph_dir: Optional[str] = None) -> Optional[str]:
    """
    Load the mermaid diagram for a sub-cascade session and extract the meaningful content.

    Returns the cell flow portion without style definitions, ready to be embedded.
    Returns None if the file doesn't exist or can't be parsed.

    Args:
        sub_session_id: The session ID of the sub-cascade
        graph_dir: Optional override for graph directory (for testing)
    """
    try:
        # Get graph directory
        if graph_dir is None:
            config = get_config()
            graph_dir = config.graph_dir

        graph_path = os.path.join(graph_dir, f"{sub_session_id}.mmd")

        if not os.path.exists(graph_path):
            return None

        with open(graph_path, 'r') as f:
            content = f.read()

        # Extract everything after the style definitions
        # Strategy: Skip lines until we find actual node/edge definitions
        lines = content.split('\n')
        meaningful_lines = []
        found_content = False

        for line in lines:
            stripped = line.strip()

            # Skip header and style definitions
            if stripped.startswith('graph ') or \
               stripped.startswith('stateDiagram') or \
               stripped.startswith('classDef ') or \
               stripped.startswith('%%') or \
               (not stripped and not found_content):
                continue

            # We've found real content
            found_content = True
            meaningful_lines.append(line)

        # Filter out trailing empty lines
        while meaningful_lines and not meaningful_lines[-1].strip():
            meaningful_lines.pop()

        if not meaningful_lines:
            return None

        return '\n'.join(meaningful_lines)

    except Exception as e:
        # Silently fail - we'll fall back to simple rendering
        return None


def extract_metadata(entry: Dict) -> Dict:
    """Extract and parse metadata from history entry."""
    meta = entry.get("metadata", {})
    if isinstance(meta, str):
        try:
            meta = json.loads(meta)
        except:
            meta = {}
    return {
        'take_index': meta.get('take_index'),
        'is_winner': meta.get('is_winner'),
        'reforge_step': meta.get('reforge_step'),
        'cell_name': meta.get('cell_name'),
        'cascade_id': meta.get('cascade_id'),
        'factor': meta.get('factor'),
        'has_takes': meta.get('has_takes'),
        'has_wards': meta.get('has_wards'),
        'has_sub_cascades': meta.get('has_sub_cascades'),
        'handoffs': meta.get('handoffs', []),
        'winner_index': meta.get('winner_index'),
        'depth': meta.get('depth', 0),
        # Ward-specific
        'ward_type': meta.get('ward_type'),
        'validator': meta.get('validator'),
        'mode': meta.get('mode'),
        'valid': meta.get('valid'),
        'reason': meta.get('reason'),
        # Quartermaster-specific
        'selected_skills': meta.get('selected_skills', []),
        'reasoning': meta.get('reasoning'),
        # Reforge-specific
        'reforge_step': meta.get('reforge_step'),
        'total_steps': meta.get('total_steps'),
        'factor_per_step': meta.get('factor_per_step'),
        'attempt_index': meta.get('attempt_index'),
        'has_mutation': meta.get('has_mutation'),
        # Retry/validation-specific
        'attempt': meta.get('attempt'),
        'max_attempts': meta.get('max_attempts'),
        # Mutation-specific (for takes)
        'mutation_applied': meta.get('mutation_applied'),
        'mutation_type': meta.get('mutation_type'),
        'mutation_template': meta.get('mutation_template'),
        # Tool-specific
        'tool_name': meta.get('tool_name'),
        # Turn-specific
        'turn_number': meta.get('turn_number'),
        'max_turns': meta.get('max_turns'),
    }


def extract_routing_choices(lineage: List[Dict]) -> Dict[str, str]:
    """
    Extract which cell dynamically routed to which target.

    Parses lineage entries like "Dynamically routed to: target_cell"

    Returns:
        Dict mapping source_cell -> target_cell for dynamic routing decisions
    """
    routing_choices = {}
    for item in lineage:
        cell = item.get("cell")
        output = item.get("output", "")
        if isinstance(output, str) and output.startswith("Dynamically routed to: "):
            target = output.replace("Dynamically routed to: ", "")
            routing_choices[cell] = target
    return routing_choices


def extract_ward_retries(history: List[Dict]) -> Dict[str, Dict]:
    """
    Extract ward retry information per cell.

    Identifies cells where retry-mode wards failed and triggered re-execution.

    Returns:
        Dict: cell_name -> {
            'has_retry': bool,
            'retry_count': int,
            'validators': [list of validator names that failed]
        }
    """
    retry_info = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type in ("pre_ward", "post_ward"):
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            mode = meta.get("mode")
            valid = meta.get("valid")
            validator = meta.get("validator", "validator")

            # Track retry wards that failed
            if mode == "retry" and valid is False:
                if cell_name not in retry_info:
                    retry_info[cell_name] = {
                        'has_retry': True,
                        'retry_count': 0,
                        'validators': []
                    }
                retry_info[cell_name]['retry_count'] += 1
                if validator not in retry_info[cell_name]['validators']:
                    retry_info[cell_name]['validators'].append(validator)

        elif node_type == "validation_retry":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            if cell_name not in retry_info:
                retry_info[cell_name] = {
                    'has_retry': True,
                    'retry_count': 1,
                    'validators': []
                }
            else:
                retry_info[cell_name]['retry_count'] += 1

    return retry_info


def extract_validation_retries(history: List[Dict]) -> Dict[str, Dict]:
    """
    Extract validation retry (loop_until) information per cell.

    When max_attempts > 1 and validation fails, the cell re-executes.
    This tracks those retry loops separately from ward retries.

    Returns:
        Dict: cell_name -> {
            'retry_count': int,  # Number of retry attempts (0 = passed first try)
            'max_attempts': int,  # Total attempts allowed
            'reasons': [str],    # Validation failure reasons
            'passed': bool       # Whether it eventually passed
        }
    """
    validation_retries = {}

    for entry in history:
        node_type = entry.get("node_type", "")

        # Track validation_retry nodes (injected when attempt > 0)
        if node_type == "validation_retry":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            content = entry.get("content", "")

            if cell_name not in validation_retries:
                validation_retries[cell_name] = {
                    'retry_count': 0,
                    'max_attempts': meta.get('max_attempts', 1),
                    'reasons': [],
                    'passed': False  # Will update if we see success
                }

            validation_retries[cell_name]['retry_count'] += 1

            # Extract reason from content
            if "rejected" in content.lower() or "failed" in content.lower():
                reason = content[:100] if len(content) > 100 else content
                validation_retries[cell_name]['reasons'].append(reason)

        # Track schema_validation success to mark as passed
        elif node_type == "schema_validation":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            content = entry.get("content", "")

            if cell_name in validation_retries and "passed" in content.lower():
                validation_retries[cell_name]['passed'] = True

        # Also look for loop_until validator results
        elif node_type == "loop_until_validation":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            valid = meta.get("valid")

            if cell_name not in validation_retries:
                validation_retries[cell_name] = {
                    'retry_count': 0,
                    'max_attempts': meta.get('max_attempts', 1),
                    'reasons': [],
                    'passed': False
                }

            if valid:
                validation_retries[cell_name]['passed'] = True
            else:
                validation_retries[cell_name]['retry_count'] += 1
                reason = meta.get('reason', 'Validation failed')
                validation_retries[cell_name]['reasons'].append(reason)

    return validation_retries


def extract_errors(history: List[Dict]) -> Dict[str, List[Dict]]:
    """
    Extract error nodes per cell.

    Returns:
        Dict: cell_name -> [
            {'error_type': str, 'message': str, 'trace_id': str}
        ]
    """
    errors = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type in ("error", "validation_error", "schema_validation_failed"):
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            content = entry.get("content", "")

            error_info = {
                'error_type': node_type,
                'message': content[:100] if len(content) > 100 else content,
                'trace_id': entry.get("trace_id", "")
            }

            if cell_name not in errors:
                errors[cell_name] = []
            errors[cell_name].append(error_info)

    return errors


def extract_quartermaster_selections(history: List[Dict]) -> Dict[str, Dict]:
    """
    Extract Quartermaster (manifest) tool selections per cell.

    Returns:
        Dict: cell_name -> {
            'selected_tools': [list of tool names],
            'reasoning': str
        }
    """
    selections = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type == "quartermaster_result":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            selected_skills = meta.get("selected_skills", [])
            reasoning = meta.get("reasoning", "")

            selections[cell_name] = {
                'selected_tools': selected_skills,
                'reasoning': reasoning[:100] if len(reasoning) > 100 else reasoning
            }

    return selections


def extract_turns(history: List[Dict]) -> Dict[str, List[Dict]]:
    """
    Extract turn information per cell for detailed visualization.

    Returns:
        Dict: cell_name -> [
            {'turn_number': int, 'has_tool_calls': bool, 'content_preview': str}
        ]
    """
    turns = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type == "turn":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            turn_number = meta.get("turn_number", 1)

            if cell_name not in turns:
                turns[cell_name] = []

            turns[cell_name].append({
                'turn_number': turn_number,
                'trace_id': entry.get("trace_id", "")
            })

        # Also track agent responses within turns
        elif node_type == "agent":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            content = entry.get("content", "")
            tool_calls = entry.get("tool_calls")

            # Find the corresponding turn and add info
            if cell_name in turns and turns[cell_name]:
                last_turn = turns[cell_name][-1]
                last_turn['has_tool_calls'] = bool(tool_calls)
                last_turn['content_preview'] = content[:50] if content else ""

    return turns


def extract_tool_calls(history: List[Dict]) -> Dict[str, List[Dict]]:
    """
    Extract detailed tool call information per cell.

    Returns:
        Dict: cell_name -> [
            {'tool_name': str, 'trace_id': str, 'turn_number': int, 'success': bool}
        ]
    """
    tool_calls = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type == "tool_result":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            tool_name = meta.get("tool_name")
            trace_id = entry.get("trace_id", "")
            turn_number = meta.get("turn_number", 1)

            # Try to extract tool name from content if not in metadata
            if not tool_name:
                content = entry.get("content", "")
                if "Tool Result (" in content:
                    match = re.search(r"Tool Result \((\w+)\)", content)
                    if match:
                        tool_name = match.group(1)

            if tool_name:
                if cell_name not in tool_calls:
                    tool_calls[cell_name] = []

                # Check if this looks like an error result
                content = entry.get("content", "")
                success = "error" not in content.lower()[:100]

                tool_calls[cell_name].append({
                    'tool_name': tool_name,
                    'trace_id': trace_id,
                    'turn_number': turn_number,
                    'success': success
                })

    return tool_calls


def extract_blocked_cells(history: List[Dict]) -> Dict[str, Dict]:
    """
    Extract cells that were blocked by blocking-mode wards.

    Returns:
        Dict: cell_name -> {
            'blocked': True,
            'validator': str,
            'reason': str
        }
    """
    blocked = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type in ("pre_ward", "post_ward", "ward_block"):
            meta = extract_metadata(entry)
            mode = meta.get("mode")
            valid = meta.get("valid")
            cell_name = meta.get("cell_name", "unknown")

            if mode == "blocking" and valid is False:
                blocked[cell_name] = {
                    'blocked': True,
                    'validator': meta.get("validator", "validator"),
                    'reason': meta.get("reason", "Validation failed"),
                    'ward_type': meta.get("ward_type", "post")
                }

    return blocked


def extract_state_changes(history: List[Dict]) -> Dict[str, List[str]]:
    """
    Extract set_state calls to show state flow between cells.

    Returns:
        Dict: cell_name -> [list of state keys set]
    """
    state_changes = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type == "tool_result":
            content = entry.get("content", "")
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")

            # Check if this is a set_state result
            # Format: "State updated: {key} = {value}"
            if isinstance(content, str) and ("state updated" in content.lower() or "set_state" in content.lower()):
                # Try multiple patterns
                match = re.search(r"[Ss]tate\s+updated:\s*(\w+)\s*=", content)
                if not match:
                    match = re.search(r"[Ss]tate\s*['\"](\w+)['\"]", content)
                if match:
                    key = match.group(1)
                    if cell_name not in state_changes:
                        state_changes[cell_name] = []
                    if key not in state_changes[cell_name]:
                        state_changes[cell_name].append(key)

        # Also check tool_calls for set_state
        tool_calls = entry.get("tool_calls")
        if tool_calls:
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            for tc in tool_calls:
                if isinstance(tc, dict):
                    func = tc.get("function", {})
                    if isinstance(func, dict) and func.get("name") == "set_state":
                        args = func.get("arguments", "{}")
                        try:
                            if isinstance(args, str):
                                args = json.loads(args)
                            key = args.get("key", args.get("name"))
                            if key:
                                if cell_name not in state_changes:
                                    state_changes[cell_name] = []
                                if key not in state_changes[cell_name]:
                                    state_changes[cell_name].append(key)
                        except:
                            pass

    return state_changes


def extract_take_mutations(history: List[Dict]) -> Dict[str, Dict[int, Dict]]:
    """
    Extract mutation information for take attempts.

    Returns:
        Dict: cell_name -> {take_index -> {'mutation_applied': str, 'mutation_type': str, 'is_winner': bool}}
    """
    mutations = {}

    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type == "take_attempt":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            take_index = meta.get("take_index", 0)
            mutation_applied = meta.get("mutation_applied")
            is_winner = meta.get("is_winner", False)

            if cell_name not in mutations:
                mutations[cell_name] = {}

            mutations[cell_name][take_index] = {
                'mutation_applied': mutation_applied,
                'is_winner': is_winner
            }

    return mutations


def get_live_session_state(session_id: str) -> dict:
    """
    Get the live running state for a session.

    Returns:
        Dict with status, current_cell, cell_progress, or None if not found.
    """
    try:
        from .state import get_session_state
        return get_session_state(session_id)
    except Exception:
        return None


def format_cell_progress_indicator(cell_progress: dict) -> str:
    """
    Format a compact indicator string showing current position within a cell.

    Uses cell_progress from state.json to show exactly where execution is:
    - Stage: pre_ward, main, post_ward
    - Turn: T1/3 (turn 1 of 3)
    - Attempt: A2/5 (attempt 2 of 5 for validation)
    - Take: S2/5⚖ (take 2 of 5, currently evaluating)
    - Reforge: R1/3 (reforge step 1 of 3)
    - Ward: 🛡️grammar_check (current ward being run)
    - Tool: 🔧run_code (current tool being called)

    Returns compact indicator like: "T2/3 🔧run_code" or "S3/5 ⚖️eval"
    """
    if not cell_progress:
        return ""

    parts = []

    # Stage indicator (only if not 'main')
    stage = cell_progress.get("stage", "main")
    if stage == "pre_ward":
        parts.append("⏵pre")
    elif stage == "post_ward":
        parts.append("⏵post")

    # Turn info
    turn_info = cell_progress.get("turn", {})
    current_turn = turn_info.get("current", 0)
    max_turns = turn_info.get("max", 1)
    if current_turn > 0:
        parts.append(f"T{current_turn}/{max_turns}")

    # Attempt info (validation retries)
    attempt_info = cell_progress.get("attempt", {})
    current_attempt = attempt_info.get("current", 0)
    max_attempts = attempt_info.get("max", 1)
    if current_attempt > 0 and max_attempts > 1:
        parts.append(f"A{current_attempt}/{max_attempts}")

    # Take info
    take_info = cell_progress.get("take")
    if take_info:
        take_idx = take_info.get("index")
        take_factor = take_info.get("factor", 1)
        take_stage = take_info.get("stage", "executing")

        if take_idx is not None:
            stage_icon = "⚖️" if take_stage == "evaluating" else "🔱"
            parts.append(f"{stage_icon}S{take_idx + 1}/{take_factor}")

    # Reforge info
    reforge_info = cell_progress.get("reforge")
    if reforge_info:
        reforge_step = reforge_info.get("step")
        total_steps = reforge_info.get("total_steps", 1)
        if reforge_step is not None:
            parts.append(f"🔨R{reforge_step}/{total_steps}")

    # Ward info
    ward_info = cell_progress.get("ward")
    if ward_info and ward_info.get("name"):
        ward_name = ward_info.get("name", "")
        ward_type = ward_info.get("type", "post")
        ward_idx = ward_info.get("index", 1)
        total_wards = ward_info.get("total", 1)
        type_icon = "🛡️" if ward_type == "pre" else "🔄"
        # Truncate ward name
        short_name = ward_name[:10] + ".." if len(ward_name) > 12 else ward_name
        parts.append(f"{type_icon}{short_name}({ward_idx}/{total_wards})")

    # Tool info
    tool_info = cell_progress.get("tool", {})
    current_tool = tool_info.get("current")
    if current_tool:
        # Truncate tool name
        short_tool = current_tool[:10] + ".." if len(current_tool) > 12 else current_tool
        parts.append(f"🔧{short_tool}")

    # Timing info (optional - elapsed time)
    timing = cell_progress.get("timing", {})
    cell_elapsed = timing.get("cell_elapsed_ms", 0)
    if cell_elapsed > 1000:  # Only show if > 1 second
        seconds = cell_elapsed // 1000
        parts.append(f"⏱{seconds}s")

    return " ".join(parts) if parts else ""


def get_running_internal_node_id(cell_progress: dict, pid: str) -> Optional[str]:
    """
    Determine which internal node ID within a composite cell is currently executing.

    Returns the Mermaid node ID that should be highlighted, or None if not determinable.
    """
    if not cell_progress:
        return None

    stage = cell_progress.get("stage", "main")

    # Pre-ward stage
    if stage == "pre_ward":
        ward_info = cell_progress.get("ward")
        if ward_info:
            ward_idx = ward_info.get("index", 1) - 1  # 0-indexed
            return f"{pid}_pre{ward_idx}"

    # Post-ward stage
    elif stage == "post_ward":
        ward_info = cell_progress.get("ward")
        if ward_info:
            ward_idx = ward_info.get("index", 1) - 1  # 0-indexed
            return f"{pid}_post{ward_idx}"

    # Main stage - could be turn, take, or reforge
    elif stage == "main":
        take_info = cell_progress.get("take")
        if take_info:
            take_idx = take_info.get("index")
            take_stage = take_info.get("stage")

            if take_stage == "evaluating":
                return f"{pid}_eval"
            elif take_idx is not None:
                return f"{pid}_a{take_idx}"

        reforge_info = cell_progress.get("reforge")
        if reforge_info:
            reforge_step = reforge_info.get("step")
            if reforge_step is not None:
                return f"{pid}_rf{reforge_step}"

        # Turn-based progress
        turn_info = cell_progress.get("turn", {})
        current_turn = turn_info.get("current", 0)
        if current_turn > 0:
            return f"{pid}_t{current_turn - 1}"

    return None


def flatten_history(history: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
    """
    Flatten history entries, extracting nested sub_echo histories.

    When echo.merge() is called, sub-cascade history is stored as:
    {"sub_echo": "session_id", "history": [...]}

    Returns:
        (flattened_entries, sub_echo_entries)
        - flattened_entries: All regular history entries
        - sub_echo_entries: List of {"sub_echo": session_id, "history": [...]} entries
    """
    flattened = []
    sub_echoes = []
    for entry in history:
        if "sub_echo" in entry and "history" in entry:
            # This is a merged sub-cascade - keep track of it
            sub_echoes.append(entry)
            # Also recursively flatten its history
            sub_history = entry.get("history", [])
            nested_flat, nested_subs = flatten_history(sub_history)
            flattened.extend(nested_flat)
            sub_echoes.extend(nested_subs)
        else:
            flattened.append(entry)
    return flattened, sub_echoes


def build_execution_tree(echo: Echo) -> Tuple[List[ExecutionNode], Dict[str, ExecutionNode]]:
    """
    Build an execution tree from Echo history.
    Returns (root_nodes, all_nodes_map)
    """
    # Flatten history to include nested sub_echo entries
    history, _ = flatten_history(echo.history)
    nodes_map: Dict[str, ExecutionNode] = {}
    root_nodes: List[ExecutionNode] = []

    # First pass: create all nodes
    for entry in history:
        trace_id = entry.get("trace_id")
        if not trace_id or trace_id in nodes_map:
            continue

        # Extract metadata
        meta = extract_metadata(entry)

        node = ExecutionNode(
            id=trace_id,
            node_type=entry.get("node_type", "msg"),
            name=entry.get("content", "")[:30] if entry.get("node_type") in ("cascade", "cell") else "",
            role=entry.get("role", ""),
            content=entry.get("content", ""),
            parent_id=entry.get("parent_id"),
            metadata=meta,
            take_index=meta.get('take_index'),
            is_winner=meta.get('is_winner', False),
            reforge_step=meta.get('reforge_step')
        )
        nodes_map[trace_id] = node

    # Second pass: build parent-child relationships
    for node in nodes_map.values():
        if node.parent_id and node.parent_id in nodes_map:
            nodes_map[node.parent_id].children.append(node)
        else:
            root_nodes.append(node)

    return root_nodes, nodes_map


def collect_cells(nodes_map: Dict[str, ExecutionNode]) -> List[ExecutionNode]:
    """Collect all cell nodes in order of appearance."""
    cells = []
    seen = set()
    for node in nodes_map.values():
        if node.node_type == "cell" and node.id not in seen:
            cells.append(node)
            seen.add(node.id)
    return cells


def collect_takes(nodes_map: Dict[str, ExecutionNode], cell_id: str) -> Dict[int, List[ExecutionNode]]:
    """Collect take attempts for a cell, grouped by take_index."""
    takes: Dict[int, List[ExecutionNode]] = {}

    for node in nodes_map.values():
        if node.parent_id == cell_id or (node.parent_id and nodes_map.get(node.parent_id, ExecutionNode("", "", "")).parent_id == cell_id):
            if node.take_index is not None:
                if node.take_index not in takes:
                    takes[node.take_index] = []
                takes[node.take_index].append(node)

    return takes


def collect_reforge_steps(nodes_map: Dict[str, ExecutionNode], cell_id: str) -> Dict[int, List[ExecutionNode]]:
    """Collect reforge steps for a cell, grouped by reforge_step."""
    reforge_steps: Dict[int, List[ExecutionNode]] = {}

    for node in nodes_map.values():
        if node.reforge_step is not None:
            if node.reforge_step not in reforge_steps:
                reforge_steps[node.reforge_step] = []
            reforge_steps[node.reforge_step].append(node)

    return reforge_steps


def export_execution_graph_json(echo: Echo, output_path: str) -> str:
    """
    Export execution graph as structured JSON for UI consumption.

    Provides easy-to-query structure with trace_ids for DB lookups.
    Includes nodes, edges, and metadata without needing to parse mermaid.

    Returns:
        Path to written JSON file
    """
    root_nodes, nodes_map = build_execution_tree(echo)
    history, sub_echoes = flatten_history(echo.history)

    # Build nodes list
    nodes = []
    edges = []

    for trace_id, node in nodes_map.items():
        # Create node entry
        node_entry = {
            "trace_id": trace_id,
            "node_type": node.node_type,
            "role": node.role,
            "parent_id": node.parent_id,
            "depth": node.metadata.get("depth", 0),

            # Metadata
            "cell_name": node.metadata.get("cell_name"),
            "cascade_id": node.metadata.get("cascade_id"),

            # Takes/Reforge
            "take_index": node.take_index,
            "is_winner": node.is_winner,
            "reforge_step": node.reforge_step,

            # Content preview (truncated for JSON size)
            "content_preview": str(node.content)[:100] if node.content else None,

            # Additional metadata
            "metadata": {
                k: v for k, v in node.metadata.items()
                if k not in ("cell_name", "cascade_id", "depth")
            }
        }
        nodes.append(node_entry)

        # Create edge if has parent
        if node.parent_id:
            edges.append({
                "source": node.parent_id,
                "target": trace_id,
                "edge_type": "parent_child"
            })

    # Collect cells in order
    cells = []
    for item in echo.lineage:
        cells.append({
            "cell": item.get("cell"),  # Note: lineage still uses "cell" key internally
            "trace_id": item.get("trace_id"),
            "output_preview": str(item.get("output", ""))[:100]
        })

    # Build cell connections from history
    cell_nodes = [n for n in nodes if n["node_type"] == "cell"]
    for i in range(len(cell_nodes) - 1):
        edges.append({
            "source": cell_nodes[i]["trace_id"],
            "target": cell_nodes[i+1]["trace_id"],
            "edge_type": "cell_sequence"
        })

    # Collect takes info
    takes_groups = {}
    for node in nodes:
        if node["take_index"] is not None:
            cell = node.get("cell_name", "unknown")
            if cell not in takes_groups:
                takes_groups[cell] = []
            takes_groups[cell].append({
                "trace_id": node["trace_id"],
                "take_index": node["take_index"],
                "is_winner": node["is_winner"],
                "reforge_step": node.get("reforge_step")
            })

    # Build sub-cascade mapping for mermaid node IDs
    # Maps mermaid node IDs (like "call_child_a0") to sub-cascade session IDs
    sub_cascade_mapping = {}

    # Scan for existing sub-cascade mermaid files using predictable naming pattern
    # Pattern: {parent_session_id}_sub_{index}.json
    output_dir = os.path.dirname(output_path) if output_path else "graphs"

    # Look for all sub-cascade JSON files
    pattern = os.path.join(output_dir, f"{echo.session_id}_sub_*.json")
    sub_files = glob.glob(pattern)

    for sub_file in sub_files:
        # Extract sub-cascade session ID from filename
        basename = os.path.basename(sub_file)
        sub_session_id = basename.replace(".json", "")

        # Extract take index from session ID
        if "_sub_" in sub_session_id:
            try:
                idx_str = sub_session_id.split("_sub_")[-1]
                take_index = int(idx_str)
            except (ValueError, IndexError):
                continue

            # Try to load the JSON to get cascade_id
            cascade_id = None
            try:
                with open(sub_file, 'r') as f:
                    sub_data = json.load(f)
                    # Find cascade node in nodes
                    for node in sub_data.get("nodes", []):
                        if node.get("node_type") == "cascade":
                            cascade_id = node.get("cascade_id")
                            break
            except Exception:
                pass  # Gracefully handle missing/corrupt files

            # Build mermaid node ID by checking which cell has this take
            # For now, we need to infer the cell name from the mermaid structure
            # Since we're generating this mapping AFTER mermaid is generated,
            # we can scan the actual mermaid content for the node IDs

            # This mapping is session_id-based, UI can match by pattern
            # Store with a generic key that UI can search
            sub_cascade_mapping[f"take_{take_index}"] = {
                "session_id": sub_session_id,
                "cascade_id": cascade_id,
                "take_index": take_index
            }

    # Build output structure
    graph = {
        "session_id": echo.session_id,
        "generated_at": None,  # Could add timestamp

        "nodes": nodes,
        "edges": edges,

        "cells": cells,
        "takes": takes_groups,

        # NEW: Sub-cascade mapping for UI
        "sub_cascade_mapping": sub_cascade_mapping,

        "summary": {
            "total_nodes": len(nodes),
            "total_edges": len(edges),
            "total_cells": len(cells),
            "has_takes": len(takes_groups) > 0,
            "has_sub_cascades": len(sub_echoes) > 0
        }
    }

    # Write JSON
    with open(output_path, "w") as f:
        json.dump(graph, f, indent=2, default=str)

    return output_path


def export_react_flow_graph(echo: Echo, output_path: str) -> str:
    """
    Export execution graph in React Flow format for direct UI use.

    React Flow format: https://reactflow.dev/
    Ready to drop into React Flow component with custom node types.

    Returns:
        Path to written JSON file
    """
    root_nodes, nodes_map = build_execution_tree(echo)
    history, sub_echoes = flatten_history(echo.history)

    rf_nodes = []
    rf_edges = []

    # Position layout (simple left-to-right based on depth and order)
    node_positions = {}
    y_offset = 0
    x_spacing = 250
    y_spacing = 100

    # Group nodes by depth for layout
    depth_groups = {}
    for trace_id, node in nodes_map.items():
        depth = node.metadata.get("depth", 0)
        if depth not in depth_groups:
            depth_groups[depth] = []
        depth_groups[depth].append((trace_id, node))

    # Assign positions
    for depth in sorted(depth_groups.keys()):
        nodes_at_depth = depth_groups[depth]
        for i, (trace_id, node) in enumerate(nodes_at_depth):
            node_positions[trace_id] = {
                "x": depth * x_spacing,
                "y": i * y_spacing
            }

    # Build React Flow nodes
    for trace_id, node in nodes_map.items():
        pos = node_positions.get(trace_id, {"x": 0, "y": y_offset})
        y_offset += y_spacing

        # Determine node type for custom rendering
        if node.node_type == "cell":
            rf_type = "cellNode"
        elif node.node_type == "cascade":
            rf_type = "cascadeNode"
        elif node.node_type in ("take_attempt", "takes"):
            rf_type = "takeNode"
        elif node.node_type in ("reforge_step", "reforge_attempt"):
            rf_type = "reforgeNode"
        elif node.node_type == "tool_result":
            rf_type = "toolNode"
        else:
            rf_type = "default"

        rf_node = {
            "id": trace_id,
            "type": rf_type,
            "position": pos,
            "data": {
                "label": node.content[:50] if node.content else node.node_type,
                "trace_id": trace_id,
                "node_type": node.node_type,
                "role": node.role,
                "cell_name": node.metadata.get("cell_name"),
                "cascade_id": node.metadata.get("cascade_id"),
                "take_index": node.take_index,
                "is_winner": node.is_winner,
                "reforge_step": node.reforge_step,
                "metadata": node.metadata
            }
        }

        # Add parent node reference for grouping
        if node.parent_id:
            rf_node["parentNode"] = node.parent_id
            rf_node["extent"] = "parent"  # Constrain to parent bounds

        rf_nodes.append(rf_node)

    # Build React Flow edges
    for trace_id, node in nodes_map.items():
        if node.parent_id:
            # Determine edge style based on relationship
            edge_style = {}
            animated = False
            edge_type = "default"

            if node.is_winner:
                edge_style = {"stroke": "#00ff00", "strokeWidth": 3}
                animated = True
                edge_type = "winner"
            elif node.take_index is not None:
                edge_style = {"stroke": "#fab005", "strokeDasharray": "5 5"}
                edge_type = "take"
            elif node.node_type == "cell":
                edge_style = {"stroke": "#1c7ed6", "strokeWidth": 2}
                edge_type = "cell"

            rf_edge = {
                "id": f"e_{node.parent_id}_{trace_id}",
                "source": node.parent_id,
                "target": trace_id,
                "type": edge_type,
                "animated": animated,
                "style": edge_style,
                "data": {
                    "edge_type": "parent_child"
                }
            }

            rf_edges.append(rf_edge)

    # Add cell sequence edges
    cells = [item.get("trace_id") for item in echo.lineage if item.get("trace_id")]
    for i in range(len(cells) - 1):
        rf_edges.append({
            "id": f"seq_{cells[i]}_{cells[i+1]}",
            "source": cells[i],
            "target": cells[i+1],
            "type": "cell_sequence",
            "animated": True,
            "style": {"stroke": "#1c7ed6", "strokeWidth": 2},
            "data": {"edge_type": "cell_sequence"}
        })

    # Output React Flow format
    react_flow_data = {
        "nodes": rf_nodes,
        "edges": rf_edges,
        "meta": {
            "session_id": echo.session_id,
            "total_nodes": len(rf_nodes),
            "total_edges": len(rf_edges)
        }
    }

    with open(output_path, "w") as f:
        json.dump(react_flow_data, f, indent=2, default=str)

    return output_path


def generate_mermaid_string(echo: Echo) -> str:
    """ NOT USED ANYMORE, WE USE STATE DIAGRAMS NOW!!! THIS ONLY FOR LEGACY REFERENCE!!!
    Generate a Mermaid flowchart string from Echo history.

    The diagram shows:
    - Cascade as the outer container
    - Cells as nodes connected by handoffs
    - Takes as parallel branches with winner highlighting
    - Reforge as sequential refinement steps
    - Sub-cascades as nested groups

    Returns the mermaid diagram as a string without writing to file.
    """
    root_nodes, nodes_map = build_execution_tree(echo)

    # Flatten history to include nested sub_echo entries
    history, sub_echoes = flatten_history(echo.history)

    lines = ["graph TD"]

    # Style definitions - Midnight Fjord Dark Theme
    # Background: #0B1219 (Midnight Fjord), #16202A (Abyssal Slate)
    # Accents: #2DD4BF (Glacial Ice), #D9A553 (Compass Brass)
    # Text: #F0F4F8 (Frosted White), #9AA5B1 (Mist Gray)
    lines.extend([
        "    %% Node Styles - Midnight Fjord Dark Theme",
        "    classDef cascade fill:#16202A,stroke:#2C3B4B,stroke-width:2px,color:#F0F4F8;",
        "    classDef cell fill:#16202A,stroke:#2DD4BF,stroke-width:2px,color:#F0F4F8;",
        "    classDef cell_active fill:#1a2a3a,stroke:#2DD4BF,stroke-width:3px,color:#F0F4F8;",
        "    classDef system fill:#16202A,stroke:#60a5fa,color:#F0F4F8;",
        "    classDef user fill:#16202A,stroke:#D9A553,color:#F0F4F8;",
        "    classDef tool fill:#16202A,stroke:#f472b6,color:#F0F4F8;",
        "    classDef takes_group fill:#16202A,stroke:#D9A553,stroke-width:2px,color:#F0F4F8;",
        "    classDef attempt fill:#16202A,stroke:#D9A553,stroke-dasharray:3 3,color:#F0F4F8;",
        "    classDef winner fill:#16202A,stroke:#F0F4F8,stroke-width:4px,color:#F0F4F8;",
        "    classDef loser fill:#16202A,stroke:#9AA5B1,stroke-dasharray:5 5,color:#9AA5B1;",
        "    classDef reforge_group fill:#16202A,stroke:#D9A553,stroke-width:2px,color:#F0F4F8;",
        "    classDef reforge_step fill:#16202A,stroke:#D9A553,color:#F0F4F8;",
        "    classDef evaluator fill:#16202A,stroke:#a78bfa,stroke-width:2px,color:#F0F4F8;",
        "    classDef agent fill:#16202A,stroke:#2DD4BF,color:#F0F4F8;",
        "    classDef sub_cascade fill:#16202A,stroke:#a78bfa,stroke-width:2px,color:#F0F4F8;",
        "    classDef ward_pre fill:#16202A,stroke:#2DD4BF,color:#F0F4F8;",
        "    classDef ward_post fill:#16202A,stroke:#CF6679,color:#F0F4F8;",
        "    classDef ward_fail fill:#16202A,stroke:#CF6679,stroke-width:3px,color:#CF6679;",
        "",
    ])

    # Collect sub-cascade trace IDs to filter out their cells from top-level rendering
    sub_cascade_trace_ids = set()
    for sub_echo in sub_echoes:
        sub_history = sub_echo.get("history", [])
        for entry in sub_history:
            if entry.get("node_type") == "cascade":
                sub_cascade_trace_ids.add(entry.get("trace_id"))

    # Collect structural entries from history with their metadata
    cascade_entry = None
    cell_entries = []
    takes_entries = []
    take_attempts = []
    evaluator_entries = []
    ward_entries = []
    quartermaster_entries = []

    for entry in history:
        node_type = entry.get("node_type")
        meta = extract_metadata(entry)

        if node_type == "cascade":
            # Skip sub-cascade entries - they'll be rendered inline
            if entry.get("trace_id") in sub_cascade_trace_ids:
                continue
            cascade_entry = entry
        elif node_type == "cell":
            # Skip cells that belong to sub-cascades
            if entry.get("parent_id") in sub_cascade_trace_ids:
                continue
            cell_entries.append(entry)
        elif node_type == "takes":
            takes_entries.append(entry)
        elif node_type == "take_attempt":
            take_attempts.append(entry)
        elif node_type in ("evaluator", "evaluation"):
            evaluator_entries.append(entry)
        elif node_type in ("pre_ward", "post_ward"):
            ward_entries.append(entry)
        elif node_type == "quartermaster_result":
            quartermaster_entries.append(entry)
        elif node_type in ("validation", "schema_validation"):
            # Collect validation entries (loop_until and output_schema)
            # We'll render these inline in the cell
            pass  # Handled via cell messages below
        elif node_type == "validation_retry":
            # Retry messages are already captured as user messages
            pass
        elif node_type == "cascade_takes":
            # Cascade-level takes start marker
            pass  # Collected separately below
        elif node_type == "cascade_take_attempt":
            # Individual cascade take attempt
            pass  # Collected separately below
        elif node_type == "cascade_evaluator":
            # Cascade takes evaluator
            pass  # Collected separately below
        elif node_type == "cascade_takes_result":
            # Cascade takes result/winner
            pass  # Collected separately below
        elif node_type in ("reforge_step", "reforge_attempt", "reforge_evaluator", "reforge_winner"):
            # Reforge entries - collected separately by cell
            pass

    # Extract cell order from lineage
    cell_order = [item.get("cell") for item in echo.lineage]

    # Sort cells by lineage order if available
    if cell_order:
        def cell_sort_key(entry):
            content = entry.get("content", "")
            name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
            try:
                return cell_order.index(name)
            except ValueError:
                return 999
        cell_entries.sort(key=cell_sort_key)

    # Note: Removed cascade container wrapper for cleaner visualization
    # Diagram renders cells directly without outer border box

    # Group take attempts by cell
    takes_by_cell: Dict[str, List[Dict]] = {}
    for sa in take_attempts:
        meta = extract_metadata(sa)
        cell_name = meta.get("cell_name", "unknown")
        if cell_name not in takes_by_cell:
            takes_by_cell[cell_name] = []
        takes_by_cell[cell_name].append(sa)

    # Group evaluator entries by cell
    evaluators_by_cell: Dict[str, Dict] = {}
    for ev in evaluator_entries:
        meta = extract_metadata(ev)
        cell_name = meta.get("cell_name", "unknown")
        evaluators_by_cell[cell_name] = ev

    # Group quartermaster entries by cell
    quartermaster_by_cell: Dict[str, Dict] = {}
    for qm in quartermaster_entries:
        meta = extract_metadata(qm)
        cell_name = meta.get("cell_name", "unknown")
        quartermaster_by_cell[cell_name] = qm

    # Group reforge entries by cell and step
    # Structure: {cell_name: {step: {steps: [], attempts: [], evaluator: None, winner: None}}}
    reforge_by_cell: Dict[str, Dict[int, Dict]] = {}
    for entry in history:
        node_type = entry.get("node_type")
        if node_type in ("reforge_step", "reforge_attempt", "reforge_evaluator", "reforge_winner"):
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            step = meta.get("reforge_step", 1)

            if cell_name not in reforge_by_cell:
                reforge_by_cell[cell_name] = {}
            if step not in reforge_by_cell[cell_name]:
                reforge_by_cell[cell_name][step] = {
                    "step_entry": None,
                    "attempts": [],
                    "evaluator": None,
                    "winner": None
                }

            if node_type == "reforge_step":
                reforge_by_cell[cell_name][step]["step_entry"] = entry
            elif node_type == "reforge_attempt":
                reforge_by_cell[cell_name][step]["attempts"].append(entry)
            elif node_type == "reforge_evaluator":
                reforge_by_cell[cell_name][step]["evaluator"] = entry
            elif node_type == "reforge_winner":
                reforge_by_cell[cell_name][step]["winner"] = entry

    # Group wards by cell (using parent_id to trace to cell)
    wards_by_cell: Dict[str, Dict[str, List[Dict]]] = {}  # cell_name -> {pre: [...], post: [...]}
    cell_trace_to_name = {pe.get("trace_id"): pe.get("content", "").replace("Cell: ", "") for pe in cell_entries}

    for ward in ward_entries:
        ward_meta = extract_metadata(ward)
        parent_id = ward.get("parent_id")
        cell_name = cell_trace_to_name.get(parent_id, "unknown")

        if cell_name not in wards_by_cell:
            wards_by_cell[cell_name] = {"pre": [], "post": []}

        ward_type = ward_meta.get("ward_type", "")
        if ward_type == "pre":
            wards_by_cell[cell_name]["pre"].append(ward)
        elif ward_type == "post":
            wards_by_cell[cell_name]["post"].append(ward)

    # Determine cell completion status
    # A cell is complete if it appears in lineage
    completed_cells = {item.get("cell") for item in echo.lineage}

    # Check which cells have any activity (messages logged)
    active_cells = set()
    for entry in history:
        meta = extract_metadata(entry)
        cell_name = meta.get("cell_name")
        if cell_name:
            active_cells.add(cell_name)

    def get_cell_status(cell_name: str) -> str:
        """Get status icon for a cell."""
        if cell_name in completed_cells:
            return "+"  # Completed
        elif cell_name in active_cells:
            return ">"  # In progress (has activity but not in lineage yet)
        else:
            return "-"  # Pending

    # Collect cascade-level takes entries
    cascade_takes_start = None
    cascade_take_attempts = []
    cascade_evaluator = None
    cascade_takes_result = None

    for entry in history:
        node_type = entry.get("node_type")
        if node_type == "cascade_takes":
            cascade_takes_start = entry
        elif node_type == "cascade_take_attempt":
            cascade_take_attempts.append(entry)
        elif node_type == "cascade_evaluator":
            cascade_evaluator = entry
        elif node_type == "cascade_takes_result":
            cascade_takes_result = entry

    # Render cascade-level takes if present (appears before cells)
    cascade_takes_node_id = None
    if cascade_takes_start and cascade_take_attempts:
        cs_meta = extract_metadata(cascade_takes_start)
        factor = cs_meta.get("factor", len(cascade_take_attempts))
        cs_id = "n_cascade_takes"

        lines.append(f'        subgraph {cs_id}["Cascade Takes ({factor} executions)"]')
        lines.append("        direction TB")

        # Render attempts as boxes with sub-session links
        lines.append(f'            subgraph {cs_id}_attempts["Parallel Cascade Executions"]')
        lines.append("            direction LR")

        attempt_ids = []
        winner_index = None
        if cascade_takes_result:
            result_meta = extract_metadata(cascade_takes_result)
            winner_index = result_meta.get("winner_index")

        # Group cells by take_index for cascade-level takes
        cells_by_take: Dict[int, List[Dict]] = {}
        for cell_entry in cell_entries:
            cell_meta = extract_metadata(cell_entry)
            take_idx = cell_meta.get("take_index")
            if take_idx is not None:
                if take_idx not in cells_by_take:
                    cells_by_take[take_idx] = []
                cells_by_take[take_idx].append(cell_entry)

        for attempt in sorted(cascade_take_attempts, key=lambda a: extract_metadata(a).get("take_index", 0)):
            a_meta = extract_metadata(attempt)
            idx = a_meta.get("take_index", 0)
            sub_session_id = a_meta.get("sub_session_id", f"take_{idx}")
            is_winner = (winner_index is not None and idx == winner_index)

            attempt_id = f"{cs_id}_a{idx}"
            attempt_ids.append(attempt_id)

            # Try to load the complete mermaid diagram from the sub-session
            sub_mermaid = load_sub_cascade_mermaid(sub_session_id)

            if sub_mermaid:
                # Embed the complete sub-cascade diagram
                winner_mark = " ✓" if is_winner else ""
                lines.append(f'                subgraph {attempt_id}["Attempt #{idx+1}{winner_mark}"]')
                lines.append("                direction TB")

                # Add the sub-cascade content with proper indentation
                for line in sub_mermaid.split('\n'):
                    if line.strip():
                        # Add extra indentation for embedding
                        lines.append(f"                    {line}")

                lines.append("                end")
                if is_winner:
                    lines.append(f"                class {attempt_id} winner")
                else:
                    lines.append(f"                class {attempt_id} loser")
            else:
                # Fallback: Try to reconstruct from cells in parent history
                take_cells = cells_by_take.get(idx, [])

                if take_cells:
                    # Render as subgraph containing the cells
                    winner_mark = " ✓" if is_winner else ""
                    lines.append(f'                subgraph {attempt_id}["Attempt #{idx+1}{winner_mark}"]')
                    lines.append("                direction TB")

                    # Render each cell in this take
                    for j, cell_entry in enumerate(take_cells):
                        cell_content = cell_entry.get("content", "")
                        cell_name = cell_content.replace("Cell: ", "") if cell_content.startswith("Cell: ") else cell_content
                        cell_id = f"{attempt_id}_p{j}"

                        # Simple cell node
                        cell_label = sanitize_label(cell_name, 25)
                        lines.append(f'                    {cell_id}["{cell_label}"]')
                        lines.append(f"                    class {cell_id} cell")

                        # Connect cells sequentially within the take
                        if j > 0:
                            prev_cell_id = f"{attempt_id}_p{j-1}"
                            lines.append(f"                    {prev_cell_id} --> {cell_id}")

                    lines.append("                end")
                    if is_winner:
                        lines.append(f"                class {attempt_id} winner")
                    else:
                        lines.append(f"                class {attempt_id} loser")
                else:
                    # Last fallback: simple box
                    content_preview = sanitize_label(attempt.get("content", ""), 20)
                    if is_winner:
                        label = f"#{idx+1} ✓: {content_preview}" if content_preview else f"#{idx+1} ✓"
                        lines.append(f'                {attempt_id}["{label}"]')
                        lines.append(f"                class {attempt_id} winner")
                    else:
                        label = f"#{idx+1}: {content_preview}" if content_preview else f"#{idx+1}"
                        lines.append(f'                {attempt_id}["{label}"]')
                        lines.append(f"                class {attempt_id} loser")

        lines.append("            end")

        # Add evaluator node
        eval_id = f"{cs_id}_eval"
        if cascade_evaluator:
            eval_preview = sanitize_label(cascade_evaluator.get("content", ""), 35)
            lines.append(f'            {eval_id}{{"Eval: {eval_preview}"}}')
        else:
            lines.append(f'            {eval_id}{{"Evaluate"}}')
        lines.append(f"            class {eval_id} evaluator")

        # Connect attempts to evaluator
        for aid in attempt_ids:
            lines.append(f"            {aid} --> {eval_id}")

        # Add winner output
        if winner_index is not None:
            winner_out = f"{cs_id}_winner"
            lines.append(f'            {winner_out}(["* Cascade #{winner_index+1} selected"])')
            lines.append(f"            class {winner_out} winner")
            lines.append(f"            {eval_id} ==> {winner_out}")
            cascade_takes_node_id = winner_out
        else:
            cascade_takes_node_id = eval_id

        lines.append("        end")
        lines.append(f"        class {cs_id} takes_group")

    # Render cells and their connections
    cell_ids = []
    cell_id_map = {}

    for i, cell_entry in enumerate(cell_entries):
        meta = extract_metadata(cell_entry)

        # Skip cells that belong to cascade-level takes (they're already rendered inside take boxes)
        if meta.get("take_index") is not None:
            continue

        content = cell_entry.get("content", "")
        cell_name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
        cell_id = safe_id(cell_entry.get("trace_id", f"cell_{i}"))
        cell_ids.append(cell_id)
        cell_id_map[cell_name] = cell_id

        # Get cell status
        status_icon = get_cell_status(cell_name)

        has_takes = meta.get("has_takes", False) or cell_name in takes_by_cell
        has_wards = meta.get("has_wards", False) or cell_name in wards_by_cell
        cell_wards = wards_by_cell.get(cell_name, {"pre": [], "post": []})

        if has_takes and cell_name in takes_by_cell:
            # Render takes group
            lines.append(f'        subgraph {cell_id}["{status_icon} {sanitize_label(cell_name, 30)}"]')
            lines.append("        direction TB")

            # Get take attempts for this cell
            attempts = takes_by_cell[cell_name]
            winner_index = None

            # Build attempts info with content
            attempts_info = {}
            for sa in attempts:
                sa_meta = extract_metadata(sa)
                idx = sa_meta.get("take_index", 0)
                is_winner = sa_meta.get("is_winner", False)
                content = sa.get("content", "")
                attempts_info[idx] = {"is_winner": is_winner, "content": content}
                if is_winner:
                    winner_index = idx

            # Get evaluator content if available
            eval_entry = evaluators_by_cell.get(cell_name, {})
            eval_content = eval_entry.get("content", "") if eval_entry else ""

            # Try to load sub-cascade mermaid diagram (for cells with sub_cascades config)
            # Sub-cascades get rendered to {session_id}_sub.mmd
            sub_cascade_mermaid = load_sub_cascade_mermaid(f"{echo.session_id}_sub")

            # Also group sub-cascade cells by take_index as fallback
            # (happens when cell has both takes + sub_cascades)
            cells_by_take_idx: Dict[int, List[Dict]] = {}
            for pe in cell_entries:
                pe_meta = extract_metadata(pe)
                pe_take_idx = pe_meta.get("take_index")
                # Check if this cell belongs to a take of the current cell we're rendering
                if pe_take_idx is not None and pe_meta.get("cell_name") != cell_name:
                    # This is a sub-cascade cell inside a take
                    if pe_take_idx not in cells_by_take_idx:
                        cells_by_take_idx[pe_take_idx] = []
                    cells_by_take_idx[pe_take_idx].append(pe)

            # Render parallel attempts
            if attempts_info:
                lines.append(f'            subgraph {cell_id}_attempts["Attempts"]')
                lines.append("            direction LR")

                attempt_ids = []
                for idx in sorted(attempts_info.keys()):
                    info = attempts_info[idx]
                    attempt_id = f"{cell_id}_a{idx}"
                    attempt_ids.append(attempt_id)

                    # Priority 1: Try to embed the complete sub-cascade diagram
                    # Priority 2: Reconstruct from cells in parent history
                    # Priority 3: Simple box

                    if sub_cascade_mermaid:
                        # Embed the complete sub-cascade diagram
                        winner_mark = " ✓" if info["is_winner"] else ""
                        lines.append(f'                subgraph {attempt_id}["Attempt #{idx+1}{winner_mark}"]')
                        lines.append("                direction TB")

                        # Add the sub-cascade content with proper indentation
                        for line in sub_cascade_mermaid.split('\n'):
                            if line.strip():
                                # Add extra indentation for embedding
                                lines.append(f"                    {line}")

                        lines.append("                end")
                        if info["is_winner"]:
                            lines.append(f"                class {attempt_id} winner")
                        else:
                            lines.append(f"                class {attempt_id} loser")
                    else:
                        # Fallback: Check if this take has sub-cascade cells to render
                        take_cells = cells_by_take_idx.get(idx, [])

                        if take_cells:
                            # Render as subgraph containing the sub-cascade cells
                            winner_mark = " ✓" if info["is_winner"] else ""
                            lines.append(f'                subgraph {attempt_id}["Attempt #{idx+1}{winner_mark}"]')
                            lines.append("                direction TB")

                            # Render each sub-cascade cell in this take
                            for j, sub_cell_entry in enumerate(take_cells):
                                sub_cell_content = sub_cell_entry.get("content", "")
                                sub_cell_name = sub_cell_content.replace("Cell: ", "") if sub_cell_content.startswith("Cell: ") else sub_cell_content
                                sub_cell_id = f"{attempt_id}_p{j}"

                                # Simple cell node
                                sub_cell_label = sanitize_label(sub_cell_name, 25)
                                lines.append(f'                    {sub_cell_id}["{sub_cell_label}"]')
                                lines.append(f"                    class {sub_cell_id} cell")

                                # Connect cells sequentially within the take
                                if j > 0:
                                    prev_cell_id = f"{attempt_id}_p{j-1}"
                                    lines.append(f"                    {prev_cell_id} --> {sub_cell_id}")

                            lines.append("                end")
                            if info["is_winner"]:
                                lines.append(f"                class {attempt_id} winner")
                            else:
                                lines.append(f"                class {attempt_id} loser")
                        else:
                            # Last fallback: simple box if no sub-cascade data
                            content_preview = sanitize_label(info["content"], 25) if info["content"] else ""
                            if info["is_winner"]:
                                label = f"#{idx+1} ✓"
                                if content_preview:
                                    label = f"#{idx+1} ✓: {content_preview}"
                                lines.append(f'                {attempt_id}["{label}"]')
                                lines.append(f"                class {attempt_id} winner")
                            else:
                                label = f"#{idx+1}"
                                if content_preview:
                                    label = f"#{idx+1}: {content_preview}"
                                lines.append(f'                {attempt_id}["{label}"]')
                                lines.append(f"                class {attempt_id} loser")

                lines.append("            end")

                # Add evaluator node with content preview
                eval_id = f"{cell_id}_eval"
                if eval_content:
                    eval_preview = sanitize_label(eval_content, 35)
                    lines.append(f'            {eval_id}{{"Eval: {eval_preview}"}}')
                else:
                    lines.append(f'            {eval_id}{{"Evaluate"}}')
                lines.append(f"            class {eval_id} evaluator")

                # Connect attempts to evaluator
                for aid in attempt_ids:
                    lines.append(f"            {aid} --> {eval_id}")

                # Add winner output
                last_node_id = eval_id
                if winner_index is not None:
                    winner_out = f"{cell_id}_winner"
                    winner_content = attempts_info.get(winner_index, {}).get("content", "")
                    if winner_content:
                        winner_preview = sanitize_label(winner_content, 30)
                        lines.append(f'            {winner_out}(["* #{winner_index+1}: {winner_preview}"])')
                    else:
                        lines.append(f'            {winner_out}(["* Winner #{winner_index+1}"])')
                    lines.append(f"            class {winner_out} winner")
                    lines.append(f"            {eval_id} ==> {winner_out}")
                    last_node_id = winner_out

                # Render reforge steps if present
                cell_reforge = reforge_by_cell.get(cell_name, {})
                if cell_reforge:
                    for step_num in sorted(cell_reforge.keys()):
                        step_data = cell_reforge[step_num]
                        rf_step_entry = step_data.get("step_entry")
                        rf_attempts = step_data.get("attempts", [])
                        rf_evaluator = step_data.get("evaluator")
                        rf_winner = step_data.get("winner")

                        # Render if we have any reforge data (step entry, attempts, winner, or evaluator)
                        if rf_step_entry or rf_attempts or rf_evaluator or rf_winner:
                            # Create reforge step subgraph
                            rf_step_id = f"{cell_id}_rf{step_num}"
                            lines.append(f'            subgraph {rf_step_id}["Reforge Step {step_num}"]')
                            lines.append("            direction TB")

                            rf_attempt_ids = []
                            rf_winner_index = None
                            if rf_winner:
                                rf_winner_meta = extract_metadata(rf_winner)
                                rf_winner_index = rf_winner_meta.get("winner_index")

                            # Render refinement attempts if we have them
                            if rf_attempts:
                                lines.append(f'                subgraph {rf_step_id}_attempts["Refinements"]')
                                lines.append("                direction LR")

                                for rf_att in sorted(rf_attempts, key=lambda a: extract_metadata(a).get("attempt_index", 0)):
                                    rf_att_meta = extract_metadata(rf_att)
                                    rf_idx = rf_att_meta.get("attempt_index", 0)
                                    rf_is_winner = (rf_winner_index is not None and rf_idx == rf_winner_index)

                                    rf_att_id = f"{rf_step_id}_a{rf_idx}"
                                    rf_attempt_ids.append(rf_att_id)

                                    rf_content = rf_att.get("content", "")
                                    rf_preview = sanitize_label(rf_content, 20) if rf_content else ""

                                    if rf_is_winner:
                                        rf_label = f"R{rf_idx+1} ✓: {rf_preview}" if rf_preview else f"R{rf_idx+1} ✓"
                                        lines.append(f'                    {rf_att_id}["{rf_label}"]')
                                        lines.append(f"                    class {rf_att_id} winner")
                                    else:
                                        rf_label = f"R{rf_idx+1}: {rf_preview}" if rf_preview else f"R{rf_idx+1}"
                                        lines.append(f'                    {rf_att_id}["{rf_label}"]')
                                        lines.append(f"                    class {rf_att_id} loser")

                                lines.append("                end")

                            # Add reforge evaluator
                            rf_eval_id = f"{rf_step_id}_eval"
                            if rf_evaluator:
                                rf_eval_content = rf_evaluator.get("content", "")
                                rf_eval_preview = sanitize_label(rf_eval_content, 30)
                                lines.append(f'                {rf_eval_id}{{"Eval: {rf_eval_preview}"}}')
                            else:
                                lines.append(f'                {rf_eval_id}{{"Evaluate"}}')
                            lines.append(f"                class {rf_eval_id} evaluator")

                            # Connect refinement attempts to evaluator
                            for rf_aid in rf_attempt_ids:
                                lines.append(f"                {rf_aid} --> {rf_eval_id}")

                            # Add step winner
                            if rf_winner_index is not None:
                                rf_step_winner = f"{rf_step_id}_winner"
                                lines.append(f'                {rf_step_winner}(["* R{rf_winner_index+1}"])')
                                lines.append(f"                class {rf_step_winner} winner")
                                lines.append(f"                {rf_eval_id} ==> {rf_step_winner}")

                            lines.append("            end")
                            lines.append(f"            class {rf_step_id} reforge_group")

                            # Connect from previous step/winner to this reforge step
                            lines.append(f"            {last_node_id} --> {rf_step_id}")
                            last_node_id = rf_step_winner if rf_winner_index is not None else rf_eval_id

            lines.append("        end")
            lines.append(f"        class {cell_id} takes_group")

        elif cell_wards["pre"] or cell_wards["post"]:
            # Cell with wards - render as subgraph with checkpoints
            lines.append(f'        subgraph {cell_id}["{status_icon} {sanitize_label(cell_name, 30)}"]')
            lines.append("        direction TB")

            internal_nodes = []

            # Pre-wards
            for j, ward in enumerate(cell_wards["pre"]):
                ward_meta = extract_metadata(ward)
                validator = ward_meta.get("validator", "validator")
                valid = ward_meta.get("valid", True)
                mode = ward_meta.get("mode", "blocking")

                ward_id = f"{cell_id}_pre{j}"
                internal_nodes.append(ward_id)

                mode_label = "[B]" if mode == "blocking" else ("[R]" if mode == "retry" else "[A]")
                status_mark = "+" if valid else "x"
                label = f"{mode_label} {validator} {status_mark}"

                lines.append(f'            {ward_id}(["{sanitize_label(label, 30)}"])')
                style = "ward_pre" if valid else "ward_fail"
                lines.append(f"            class {ward_id} {style}")

            # Main cell execution node
            exec_id = f"{cell_id}_exec"
            internal_nodes.append(exec_id)
            lines.append(f'            {exec_id}["{sanitize_label(cell_name, 25)}"]')
            lines.append(f"            class {exec_id} cell")

            # Post-wards
            for j, ward in enumerate(cell_wards["post"]):
                ward_meta = extract_metadata(ward)
                validator = ward_meta.get("validator", "validator")
                valid = ward_meta.get("valid", True)
                mode = ward_meta.get("mode", "blocking")

                ward_id = f"{cell_id}_post{j}"
                internal_nodes.append(ward_id)

                mode_label = "[B]" if mode == "blocking" else ("[R]" if mode == "retry" else "[A]")
                status_mark = "+" if valid else "x"
                label = f"{mode_label} {validator} {status_mark}"

                lines.append(f'            {ward_id}(["{sanitize_label(label, 30)}"])')
                style = "ward_post" if valid else "ward_fail"
                lines.append(f"            class {ward_id} {style}")

            # Connect internal nodes sequentially
            for k in range(len(internal_nodes) - 1):
                lines.append(f"            {internal_nodes[k]} --> {internal_nodes[k+1]}")

            lines.append("        end")
            lines.append(f"        class {cell_id} cell")

        else:
            # Cell with messages - render as subgraph containing message nodes
            cell_trace_id = cell_entry.get("trace_id")

            # Collect all turn traces under this cell
            turn_traces = set()
            for entry in history:
                if entry.get("parent_id") == cell_trace_id and entry.get("node_type") == "turn":
                    turn_traces.add(entry.get("trace_id"))

            # Check for sub-cascades under this cell
            # Sub-cascades are stored as sub_echo entries after merge
            sub_cascade_entries = []
            cell_meta = extract_metadata(cell_entry)
            if cell_meta.get('has_sub_cascades'):
                # Look for sub_echo entries that contain cascade entries for this cell
                for sub_echo in sub_echoes:
                    sub_history = sub_echo.get("history", [])
                    # Find the cascade entry in the sub_echo's history
                    for entry in sub_history:
                        if entry.get("node_type") == "cascade":
                            sub_cascade_entries.append({
                                "entry": entry,
                                "history": sub_history
                            })
                            break  # Only one cascade entry per sub_echo

            # Collect tool traces under turn traces
            tool_traces = set()
            for entry in history:
                if entry.get("parent_id") in turn_traces and entry.get("node_type") == "tool":
                    tool_traces.add(entry.get("trace_id"))

            # Collect messages belonging to this cell
            cell_messages = []
            for entry in history:
                entry_parent = entry.get("parent_id")
                entry_type = entry.get("node_type", "")
                entry_content = entry.get("content", "")

                # Direct children of cell (system, user, injection, validation, validation_retry)
                if entry_parent == cell_trace_id:
                    if entry_type in ("system", "user", "injection", "validation", "schema_validation", "validation_retry"):
                        cell_messages.append(entry)
                # Children of turn traces (turn_output, turn_input, follow_up)
                elif entry_parent in turn_traces:
                    if entry_type == "turn_output":
                        # Skip empty turn_output (tool call without text response)
                        if entry_content and entry_content.strip():
                            cell_messages.append(entry)
                    elif entry_type in ("turn_input", "follow_up"):
                        cell_messages.append(entry)
                # Tool results under tool traces
                elif entry_parent in tool_traces:
                    if entry_type == "tool_result":
                        cell_messages.append(entry)
                # Grandchildren (tool results under tool traces under turns) - fallback
                elif entry_parent and entry_parent in nodes_map:
                    grandparent_id = nodes_map[entry_parent].parent_id
                    if grandparent_id in turn_traces:
                        if entry_type in ("tool_result", "injection"):
                            cell_messages.append(entry)

            # Check for quartermaster result for this cell
            qm_entry = quartermaster_by_cell.get(cell_name)

            # If we have messages, sub-cascades, or quartermaster, render as subgraph
            if cell_messages or sub_cascade_entries or qm_entry:
                lines.append(f'        subgraph {cell_id}["{status_icon} {sanitize_label(cell_name, 32)}"]')
                lines.append("        direction TB")

                all_node_ids = []

                # Render quartermaster decision node first (if present)
                if qm_entry:
                    qm_id = f"{cell_id}_qm"
                    all_node_ids.append(qm_id)
                    qm_content = qm_entry.get("content", "")
                    qm_meta = extract_metadata(qm_entry)
                    selected_skills = qm_meta.get("selected_skills", [])
                    # Format as comma-separated list of tools
                    if selected_skills:
                        tools_preview = ", ".join(selected_skills[:3])
                        if len(selected_skills) > 3:
                            tools_preview += f" +{len(selected_skills) - 3}"
                        qm_label = f"Tackle: {tools_preview}"
                    else:
                        qm_label = "No tools"
                    lines.append(f'            {qm_id}{{"{sanitize_label(qm_label, 35)}"}}')
                    lines.append(f"            class {qm_id} evaluator")

                # Render cell messages
                for j, msg in enumerate(cell_messages):
                    msg_id = f"{cell_id}_m{j}"
                    all_node_ids.append(msg_id)

                    msg_type = msg.get("node_type", "msg")
                    msg_role = msg.get("role", "")
                    msg_content = msg.get("content", "")

                    # Determine icon and style based on type
                    if msg_type == "system" or msg_role == "system":
                        icon = "SYS"
                        style = "system"
                        preview = sanitize_label(msg_content, 30)
                        label = f"{icon}: {preview}" if preview else icon
                    elif msg_type in ("user", "turn_input") or msg_role == "user":
                        icon = "USE"
                        style = "user"
                        preview = sanitize_label(msg_content, 30)
                        label = f"{icon}: {preview}" if preview else icon
                    elif msg_type in ("turn_output", "follow_up") or msg_role == "assistant":
                        icon = "OUT"
                        style = "agent"
                        preview = sanitize_label(msg_content, 35)
                        label = f"{icon}: {preview}"
                    elif msg_type == "tool_result":
                        icon = "TOOL"
                        style = "tool"
                        preview = sanitize_label(msg_content, 25)
                        label = f"{icon}: {preview}" if preview else "Tool"
                    elif msg_type == "injection":
                        icon = "INJ"
                        style = "user"
                        preview = sanitize_label(msg_content, 25)
                        label = f"{icon}: {preview}" if preview else "Inject"
                    elif msg_type in ("validation", "schema_validation"):
                        # Validation result - render as diamond decision node
                        msg_meta = extract_metadata(msg)
                        is_valid = msg_meta.get("valid", False)
                        validator = msg_meta.get("validator", "schema")
                        status_mark = "+" if is_valid else "x"
                        style = "winner" if is_valid else "ward_fail"
                        preview = sanitize_label(msg_content, 30)
                        label = f"Valid {status_mark}: {preview}"
                        # Use diamond shape for validation
                        lines.append(f'            {msg_id}{{"{label}"}}')
                        lines.append(f"            class {msg_id} {style}")
                        continue  # Skip normal rendering below (msg_id already in all_node_ids)
                    elif msg_type == "validation_retry":
                        icon = "RETRY"
                        style = "user"
                        preview = sanitize_label(msg_content, 30)
                        label = f"{icon}: {preview}" if preview else "Retry"
                    else:
                        icon = "•"
                        style = "cell"
                        label = f"{icon} {msg_type}"

                    lines.append(f'            {msg_id}["{label}"]')
                    lines.append(f"            class {msg_id} {style}")

                # Render sub-cascades as nested subgraphs
                for sc_idx, sc_data in enumerate(sub_cascade_entries):
                    sc_entry = sc_data.get("entry", {})
                    sc_history = sc_data.get("history", [])
                    sc_trace_id = sc_entry.get("trace_id")
                    sc_content = sc_entry.get("content", "")
                    sc_name = sc_content.replace("Cascade: ", "") if sc_content.startswith("Cascade: ") else sc_content
                    sc_id = f"{cell_id}_sc{sc_idx}"
                    all_node_ids.append(sc_id)

                    # Find cells under this sub-cascade from its history
                    sc_cell_entries = []
                    for entry in sc_history:
                        if entry.get("parent_id") == sc_trace_id and entry.get("node_type") == "cell":
                            sc_cell_entries.append(entry)

                    lines.append(f'            subgraph {sc_id}["Sub: {sanitize_label(sc_name, 25)}"]')
                    lines.append("            direction TB")

                    sc_cell_ids = []
                    for sp_idx, sp_entry in enumerate(sc_cell_entries):
                        sp_trace_id = sp_entry.get("trace_id")
                        sp_content = sp_entry.get("content", "")
                        sp_name = sp_content.replace("Cell: ", "") if sp_content.startswith("Cell: ") else sp_content
                        sp_id = f"{sc_id}_p{sp_idx}"
                        sc_cell_ids.append(sp_id)

                        # Find turn traces for this sub-cell (from sub-cascade's history)
                        sp_turn_traces = set()
                        for entry in sc_history:
                            if entry.get("parent_id") == sp_trace_id and entry.get("node_type") == "turn":
                                sp_turn_traces.add(entry.get("trace_id"))

                        # Collect messages for this sub-cell (from sub-cascade's history)
                        sp_messages = []
                        for entry in sc_history:
                            entry_parent = entry.get("parent_id")
                            entry_type = entry.get("node_type", "")

                            if entry_parent == sp_trace_id:
                                if entry_type in ("system", "user", "injection"):
                                    sp_messages.append(entry)
                            elif entry_parent in sp_turn_traces:
                                if entry_type in ("turn_output", "turn_input", "tool_result", "follow_up"):
                                    sp_messages.append(entry)

                        if sp_messages:
                            lines.append(f'                subgraph {sp_id}["{sanitize_label(sp_name, 20)}"]')
                            lines.append("                direction TB")

                            sp_msg_ids = []
                            for m_idx, sp_msg in enumerate(sp_messages):
                                sp_msg_id = f"{sp_id}_m{m_idx}"
                                sp_msg_ids.append(sp_msg_id)

                                msg_type = sp_msg.get("node_type", "msg")
                                msg_role = sp_msg.get("role", "")
                                msg_content = sp_msg.get("content", "")

                                if msg_type == "system" or msg_role == "system":
                                    icon = "SYS"
                                    style = "system"
                                    preview = sanitize_label(msg_content, 20)
                                    label = f"{icon}: {preview}" if preview else icon
                                elif msg_type in ("user", "turn_input") or msg_role == "user":
                                    icon = "USE"
                                    style = "user"
                                    preview = sanitize_label(msg_content, 20)
                                    label = f"{icon}: {preview}" if preview else icon
                                elif msg_type in ("turn_output", "follow_up") or msg_role == "assistant":
                                    icon = "OUT"
                                    style = "agent"
                                    preview = sanitize_label(msg_content, 25)
                                    label = f"{icon}: {preview}"
                                else:
                                    icon = "-"
                                    style = "cell"
                                    label = f"{icon} {msg_type}"

                                lines.append(f'                    {sp_msg_id}["{label}"]')
                                lines.append(f"                    class {sp_msg_id} {style}")

                            # Connect sub-cell messages
                            for k in range(len(sp_msg_ids) - 1):
                                lines.append(f"                    {sp_msg_ids[k]} --> {sp_msg_ids[k+1]}")

                            lines.append("                end")
                            lines.append(f"                class {sp_id} cell")
                        else:
                            lines.append(f'                {sp_id}["{sanitize_label(sp_name, 20)}"]')
                            lines.append(f"                class {sp_id} cell")

                    # Connect sub-cascade cells
                    for k in range(len(sc_cell_ids) - 1):
                        lines.append(f"                {sc_cell_ids[k]} --> {sc_cell_ids[k+1]}")

                    lines.append("            end")
                    lines.append(f"            class {sc_id} sub_cascade")

                # Connect all nodes (messages + sub-cascades) sequentially
                for k in range(len(all_node_ids) - 1):
                    lines.append(f"            {all_node_ids[k]} --> {all_node_ids[k+1]}")

                lines.append("        end")
                lines.append(f"        class {cell_id} cell")
            else:
                # No messages found, just render cell name with status
                lines.append(f'        {cell_id}["{status_icon} {sanitize_label(cell_name, 35)}"]')
                lines.append(f"        class {cell_id} cell")

    # If we have cascade takes, connect the winner to the first cell
    if cascade_takes_node_id and cell_ids:
        lines.append(f"        {cascade_takes_node_id} ==> {cell_ids[0]}")

    # Connect cells using handoffs from metadata, or in order
    # Only connect cells that were actually rendered (not filtered out)
    for i in range(len(cell_ids)):
        # Find the corresponding cell_entry for this cell_id
        # (skip cells that were filtered out due to take_index)
        rendered_cell_idx = 0
        cell_entry = None
        for pe in cell_entries:
            meta = extract_metadata(pe)
            if meta.get("take_index") is not None:
                continue  # Skip filtered cells
            if rendered_cell_idx == i:
                cell_entry = pe
                break
            rendered_cell_idx += 1

        if not cell_entry:
            continue

        meta = extract_metadata(cell_entry)
        handoffs = meta.get("handoffs", [])
        current_id = cell_ids[i]

        if handoffs:
            # Use explicit handoffs
            for target in handoffs:
                if target in cell_id_map:
                    lines.append(f"        {current_id} --> {cell_id_map[target]}")
        elif i + 1 < len(cell_ids):
            # Default: connect to next cell
            lines.append(f"        {current_id} --> {cell_ids[i+1]}")

    # Note: Cascade container removed for cleaner visualization without outer border
    # The structural diagram focuses on flow; use logs/lineage for content details.

    # Return the mermaid diagram as a string
    return "\n".join(lines)


def generate_state_diagram_string(echo: Echo) -> str:
    """
    Generate a Mermaid state diagram string from Echo history.

    State diagrams provide a compact, semantic representation with encapsulated complexity:
    - Cells are composite states that contain their internal complexity
    - Fork/join pseudo-states show parallel execution (takes)
    - Sub-cascades appear as nested composite states within parent cells
    - Wards chain as sequential validation steps
    - Reforge chains sequentially after takes
    - Dynamic routing shown with taken (✓) vs available (○) paths
    - Blocked cells terminate at error states (⛔)
    - Live running state shows currently executing cell (▶)
    - Mutation strategies shown on take attempts

    Visual language:
    - ▶ running (currently executing cell) + gold glow border
    - ✓ completed, ○ pending, ⛔ blocked
    - 🔱 takes (parallel attempts)
    - 🔨 reforge (iterative refinement)
    - 📦 sub-cascade (nested workflow)
    - 🛡️ blocking ward, 🔄 retry ward, ℹ️ advisory ward
    - ⚖️ evaluator, ★ winner
    - 🔧 tool usage
    - 🔄N retry count (N retries occurred)
    - 📝key1,key2 state keys set (via set_state)
    - ⏸️ HITL checkpoint (paused waiting for human input)
    - ▶️ Resumed (checkpoint resumed with human response)

    CSS Classes (applied via Mermaid classDef):
    - running: thick gold border (4px), gold fill - for currently executing cell
    - blocked: red border, red fill - for blocked cells
    - checkpoint_paused: yellow border - for cells waiting for HITL input
    - checkpoint_resumed: blue border - for checkpoints that have been resumed

    Take labels:
    - [baseline] = first attempt, no mutation
    - [mutation...] = mutation strategy applied (truncated)
    - ✓ = winner

    Routing transitions:
    - "✓ route" = taken path (agent chose this)
    - "○ available" = available but not taken
    - "⛔ validator" = blocked by validator

    Returns the mermaid state diagram as a string.
    """
    history, sub_echoes = flatten_history(echo.history)

    # Helper to create safe state IDs
    def sid(name: str) -> str:
        # Replace problematic characters and ensure valid ID
        s = name.replace("-", "_").replace(" ", "_").replace(".", "_")
        s = s.replace("(", "").replace(")", "").replace(":", "").replace("/", "_")
        return s

    # Helper to prefix all state IDs in an embedded sub-cascade diagram
    def prefix_state_ids(mermaid_content: str, prefix: str) -> str:
        """
        Rewrites all state IDs in a mermaid diagram by adding a prefix.
        This prevents ID collisions when embedding multiple sub-cascades.

        Examples:
            "state foo" -> "state prefix_foo"
            'state "label" as foo' -> 'state "label" as prefix_foo'
            "foo --> bar" -> "prefix_foo --> prefix_bar"
            "state foo {" -> "state prefix_foo {"
        """
        # First pass: collect all state IDs that are defined in the diagram
        state_ids = set()
        for line in mermaid_content.split('\n'):
            stripped = line.strip()

            # Find state IDs from "as foo" pattern
            match = re.search(r'\bas\s+([a-zA-Z_][a-zA-Z0-9_]*)\b', stripped)
            if match:
                state_ids.add(match.group(1))

            # Find state IDs from "state foo" pattern (not followed by quotes)
            match = re.search(r'\bstate\s+([a-zA-Z_][a-zA-Z0-9_]*)\b', stripped)
            if match and not stripped.startswith('state "'):
                state_ids.add(match.group(1))

        # Second pass: replace all occurrences of these state IDs
        lines = []
        for line in mermaid_content.split('\n'):
            stripped = line.strip()

            # Skip empty lines, directives, and classDef definitions
            if not stripped or stripped.startswith('direction') or stripped.startswith('classDef') or stripped.startswith('%%'):
                lines.append(line)
                continue

            # Replace each state ID with its prefixed version
            for state_id in state_ids:
                # Use word boundaries to avoid partial matches
                line = re.sub(r'\b' + state_id + r'\b', f'{prefix}_{state_id}', line)

            # Handle [*] specially - don't prefix the entry/exit nodes
            line = line.replace(f'{prefix}_[*]', '[*]')

            lines.append(line)

        return '\n'.join(lines)

    # Helper to get status icon
    def status_icon(cell_name: str, completed_cells: set) -> str:
        if is_running and cell_name == running_cell:
            return "▶"  # Currently running
        elif cell_name in completed_cells:
            return "✓"  # Completed
        else:
            return "○"  # Pending

    # Collect data structures
    completed_cells = {item.get("cell") for item in echo.lineage}

    # Extract routing choices (which cell dynamically routed to which target)
    routing_choices = extract_routing_choices(echo.lineage)

    # Extract ward retry information
    ward_retries = extract_ward_retries(history)

    # Extract cells blocked by blocking wards
    blocked_cells = extract_blocked_cells(history)

    # Extract state changes (set_state calls per cell)
    state_changes = extract_state_changes(history)

    # Extract take mutations
    take_mutations = extract_take_mutations(history)

    # Extract validation retries (loop_until / max_attempts > 1)
    validation_retries = extract_validation_retries(history)

    # Extract errors per cell
    errors_by_cell = extract_errors(history)

    # Extract Quartermaster tool selections
    qm_selections = extract_quartermaster_selections(history)

    # Extract detailed turn info per cell
    turns_detail = extract_turns(history)

    # Extract detailed tool call info per cell
    tool_calls_detail = extract_tool_calls(history)

    # Get live session state for running indicator
    live_state = get_live_session_state(echo.session_id)
    running_cell = None
    is_running = False
    cell_progress = None
    progress_indicator = ""
    running_internal_node = None

    if live_state:
        is_running = live_state.get("status") == "running"
        running_cell = live_state.get("current_cell")
        cell_progress = live_state.get("cell_progress")

        # Generate compact progress indicator from cell_progress
        if cell_progress and is_running:
            progress_indicator = format_cell_progress_indicator(cell_progress)

    # Build lookup maps
    sub_cascade_trace_ids = set()
    sub_cascades_by_cell: Dict[str, List[Dict]] = {}  # cell_trace_id -> [sub_echo_data]

    for sub_echo in sub_echoes:
        sub_history = sub_echo.get("history", [])
        for entry in sub_history:
            if entry.get("node_type") == "cascade":
                sub_cascade_trace_ids.add(entry.get("trace_id"))

    # Collect all entries by type
    cell_entries = []
    takes_by_cell: Dict[str, List[Dict]] = {}
    reforge_by_cell: Dict[str, Dict[int, Dict]] = {}
    wards_by_cell: Dict[str, Dict[str, List[Dict]]] = {}
    tools_by_cell: Dict[str, List[str]] = {}
    turns_by_cell: Dict[str, int] = {}
    handoffs_by_cell: Dict[str, List[str]] = {}  # cell_name -> list of available handoff targets
    cascade_take_attempts = []
    cascade_takes_result = None
    checkpoint_entries: List[Dict] = []  # checkpoint_created entries
    checkpoint_resume_entries: List[Dict] = []  # checkpoint_resume entries

    for entry in history:
        node_type = entry.get("node_type", "")
        meta = extract_metadata(entry)

        if node_type == "cell":
            # Only collect top-level cells (not sub-cascade cells)
            if entry.get("parent_id") not in sub_cascade_trace_ids:
                cell_entries.append(entry)
                # Also capture handoffs for this cell
                content = entry.get("content", "")
                cell_name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
                handoffs = meta.get("handoffs", [])
                if handoffs:
                    handoffs_by_cell[cell_name] = handoffs

        elif node_type == "take_attempt":
            cell_name = meta.get("cell_name", "unknown")
            if cell_name not in takes_by_cell:
                takes_by_cell[cell_name] = []
            takes_by_cell[cell_name].append(entry)

        elif node_type in ("reforge_step", "reforge_attempt", "reforge_evaluator", "reforge_winner"):
            cell_name = meta.get("cell_name", "unknown")
            step = meta.get("reforge_step", 1)
            if cell_name not in reforge_by_cell:
                reforge_by_cell[cell_name] = {}
            if step not in reforge_by_cell[cell_name]:
                reforge_by_cell[cell_name][step] = {"attempts": [], "winner": None}

            if node_type == "reforge_attempt":
                reforge_by_cell[cell_name][step]["attempts"].append(entry)
            elif node_type == "reforge_winner":
                reforge_by_cell[cell_name][step]["winner"] = entry

        elif node_type in ("pre_ward", "post_ward"):
            cell_name = meta.get("cell_name", "unknown")
            if cell_name not in wards_by_cell:
                wards_by_cell[cell_name] = {"pre": [], "post": []}
            ward_type = meta.get("ward_type", "")
            if ward_type in ("pre", "post"):
                wards_by_cell[cell_name][ward_type].append(entry)

        elif node_type == "tool_result":
            cell_name = meta.get("cell_name", "unknown")
            if cell_name not in tools_by_cell:
                tools_by_cell[cell_name] = []
            # Extract tool name from metadata (preferred) or content
            tool_name = meta.get("tool_name")
            if not tool_name:
                # Try to extract from content format "Tool Result (tool_name):"
                content = entry.get("content", "")
                if "Tool Result (" in content:
                    match = re.search(r"Tool Result \((\w+)\)", content)
                    if match:
                        tool_name = match.group(1)
            if tool_name and tool_name not in tools_by_cell[cell_name]:
                tools_by_cell[cell_name].append(tool_name)

        elif node_type == "turn":
            cell_name = meta.get("cell_name", "unknown")
            turns_by_cell[cell_name] = turns_by_cell.get(cell_name, 0) + 1

        elif node_type == "cascade_take_attempt":
            cascade_take_attempts.append(entry)

        elif node_type == "cascade_takes_result":
            cascade_takes_result = entry

        elif node_type == "checkpoint_created":
            checkpoint_entries.append(entry)

        elif node_type == "checkpoint_resume":
            checkpoint_resume_entries.append(entry)

    # Map cell trace_id to cell data for ward lookup
    cell_trace_to_name = {}
    cell_trace_to_entry = {}
    for pe in cell_entries:
        content = pe.get("content", "")
        name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
        cell_trace_to_name[pe.get("trace_id")] = name
        cell_trace_to_entry[pe.get("trace_id")] = pe

    # Re-process wards with correct cell names (using parent_id lookup)
    wards_by_cell = {}
    for entry in history:
        node_type = entry.get("node_type", "")
        if node_type in ("pre_ward", "post_ward"):
            parent_id = entry.get("parent_id")
            cell_name = cell_trace_to_name.get(parent_id, "unknown")
            meta = extract_metadata(entry)
            if cell_name not in wards_by_cell:
                wards_by_cell[cell_name] = {"pre": [], "post": []}
            ward_type = meta.get("ward_type", "")
            if ward_type in ("pre", "post"):
                wards_by_cell[cell_name][ward_type].append(entry)

    # Map sub-cascades to their parent cells
    sub_cascades_by_cell_name: Dict[str, List[Dict]] = {}
    for sub_echo in sub_echoes:
        sub_history = sub_echo.get("history", [])
        # Find cascade entry and its cells
        cascade_entry = None
        sub_cells = []
        for entry in sub_history:
            if entry.get("node_type") == "cascade":
                cascade_entry = entry
            elif entry.get("node_type") == "cell":
                sub_cells.append(entry)

        if cascade_entry:
            # Find which parent cell this belongs to
            # Look for has_sub_cascades metadata in cell entries
            for pe in cell_entries:
                pe_meta = extract_metadata(pe)
                if pe_meta.get("has_sub_cascades"):
                    cell_name = cell_trace_to_name.get(pe.get("trace_id"), "unknown")
                    if cell_name not in sub_cascades_by_cell_name:
                        sub_cascades_by_cell_name[cell_name] = []
                    sub_cascades_by_cell_name[cell_name].append({
                        "cascade": cascade_entry,
                        "cells": sub_cells,
                        "sub_echo": sub_echo.get("sub_echo", "sub")
                    })

    # Sort cells by lineage order
    cell_order = [item.get("cell") for item in echo.lineage]
    if cell_order:
        def cell_sort_key(entry):
            content = entry.get("content", "")
            name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
            try:
                return cell_order.index(name)
            except ValueError:
                return 999
        cell_entries.sort(key=cell_sort_key)

    # =========================================================================
    # BUILD THE STATE DIAGRAM
    # =========================================================================

    lines = ["stateDiagram-v2"]
    lines.append("    direction LR")
    lines.append("")

    # Add style definitions for running states (thick border, glow effect via stroke)
    # Mermaid state diagrams support classDef for styling
    lines.append("    %% Style for running states - thick gold border for visibility")
    lines.append("    classDef running fill:#1f1a0f,stroke:#D9A553,stroke-width:4px,color:#D9A553")
    lines.append("    classDef blocked fill:#2a1a1a,stroke:#ff4444,stroke-width:3px,color:#ff4444")
    lines.append("    classDef failed fill:#2a1a1a,stroke:#ff6b6b,stroke-width:2px,color:#ff6b6b")
    lines.append("    classDef sub_cascade fill:#1a1a2a,stroke:#a78bfa,stroke-width:3px,color:#a78bfa")
    lines.append("    classDef checkpoint_paused fill:#2a2a1a,stroke:#fbbf24,stroke-width:3px,color:#fbbf24")
    lines.append("    classDef checkpoint_resumed fill:#1a2a2a,stroke:#60a5fa,stroke-width:3px,color:#60a5fa")
    lines.append("")

    # Track which cell IDs need styling classes applied
    running_cell_ids = []
    failed_cell_ids = []

    # Cascade-level takes (if present)
    first_state = None
    if cascade_take_attempts:
        cs_meta = extract_metadata(cascade_takes_result) if cascade_takes_result else {}
        winner_index = cs_meta.get("winner_index")
        factor = len(cascade_take_attempts)

        lines.append("    state cascade_takes {")
        lines.append("        cs_label : 🔱 Cascade Takes")
        lines.append("        [*] --> cs_fork")
        lines.append("        state cs_fork <<fork>>")

        for attempt in sorted(cascade_take_attempts, key=lambda a: extract_metadata(a).get("take_index", 0)):
            a_meta = extract_metadata(attempt)
            idx = a_meta.get("take_index", 0)
            is_winner = (winner_index is not None and idx == winner_index)
            marker = " ✓" if is_winner else ""
            sub_session_id = a_meta.get("sub_session_id", f"take_{idx}")

            # Try to load the complete mermaid diagram for this cascade take
            sub_mermaid = load_sub_cascade_mermaid(sub_session_id)

            if sub_mermaid:
                # Embed the complete cascade execution as a composite state
                cs_attempt_id = f"cs_a{idx}"
                lines.append(f"        state {cs_attempt_id} {{")
                lines.append(f"            {cs_attempt_id}_label : Attempt #{idx+1}{marker}")
                lines.append("")

                # Prefix all state IDs in sub-cascade to prevent collisions
                prefixed_mermaid = prefix_state_ids(sub_mermaid, f"cs{idx}")

                # Embed with proper indentation (12 spaces for nesting)
                for line in prefixed_mermaid.split('\n'):
                    if line.strip():
                        lines.append(f"            {line}")

                lines.append(f"        }}")
                lines.append(f"        class {cs_attempt_id} sub_cascade")
                lines.append(f"        cs_fork --> {cs_attempt_id}")
            else:
                # Fallback to simple node
                lines.append(f"        cs_fork --> cs_a{idx}")
                lines.append(f"        cs_a{idx} : #{idx+1}{marker}")

        lines.append("        state cs_join <<join>>")
        for i in range(factor):
            attempt_id = f"cs_a{i}"
            lines.append(f"        {attempt_id} --> cs_join")

        lines.append("        cs_join --> cs_eval")
        lines.append("        cs_eval : ⚖️ Evaluate")

        if winner_index is not None:
            lines.append("        cs_eval --> cs_winner")
            lines.append(f"        cs_winner : ★ #{winner_index+1}")
            lines.append("        cs_winner --> [*]")
        else:
            lines.append("        cs_eval --> [*]")

        lines.append("    }")
        lines.append("")
        first_state = "cascade_takes"

    # =========================================================================
    # RENDER EACH CELL
    # =========================================================================

    cell_ids = []

    for cell_entry in cell_entries:
        content = cell_entry.get("content", "")
        cell_name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
        pid = sid(cell_name)
        cell_ids.append(pid)
        cell_trace_id = cell_entry.get("trace_id")

        meta = extract_metadata(cell_entry)
        status = status_icon(cell_name, completed_cells)

        # Determine what complexity this cell has
        has_takes = cell_name in takes_by_cell
        has_reforge = cell_name in reforge_by_cell
        has_wards = cell_name in wards_by_cell and (wards_by_cell[cell_name]["pre"] or wards_by_cell[cell_name]["post"])
        has_sub_cascades = cell_name in sub_cascades_by_cell_name
        turn_count = turns_by_cell.get(cell_name, 0)
        has_retries = cell_name in ward_retries
        is_blocked = cell_name in blocked_cells
        has_validation_retries = cell_name in validation_retries
        has_errors = cell_name in errors_by_cell
        has_qm_selection = cell_name in qm_selections
        detailed_turns = turns_detail.get(cell_name, [])
        has_multi_turns = len(detailed_turns) > 1

        # Get tool calls detail for this cell
        cell_tool_calls = tool_calls_detail.get(cell_name, [])
        has_many_tool_calls = len(cell_tool_calls) >= 3  # Show composite if 3+ tool calls

        # Get tools and state changes for this cell
        cell_tools = tools_by_cell.get(cell_name, [])
        cell_state_changes = state_changes.get(cell_name, [])

        # Build annotation string for simple cells
        annotations = []
        # Show tool names (up to 3) instead of just count
        if cell_tools:
            if len(cell_tools) <= 3:
                tool_str = ",".join(cell_tools)
            else:
                tool_str = ",".join(cell_tools[:2]) + f"+{len(cell_tools)-2}"
            annotations.append(f"🔧{tool_str}")
        if turn_count > 1:
            annotations.append(f"↻{turn_count}")
        # Show ward retry count if retries occurred
        if has_retries:
            retry_count = ward_retries[cell_name].get('retry_count', 0)
            if retry_count > 0:
                annotations.append(f"🔄{retry_count}")
        # Show validation retry loop (loop_until / max_attempts) - different icon
        if has_validation_retries:
            vr = validation_retries[cell_name]
            retry_count = vr.get('retry_count', 0)
            passed = vr.get('passed', False)
            if retry_count > 0:
                # Show attempt/max format: ⟳2/3 means 2 retries out of 3 max
                max_attempts = vr.get('max_attempts', retry_count + 1)
                loop_icon = "✓" if passed else "✗"
                annotations.append(f"⟳{retry_count+1}/{max_attempts}{loop_icon}")
        # Show blocked indicator if cell was blocked
        if is_blocked:
            annotations.append("⛔")
        # Show error indicator if errors occurred
        if has_errors:
            error_count = len(errors_by_cell[cell_name])
            annotations.append(f"❌{error_count}")
        # Show Quartermaster (manifest) indicator
        if has_qm_selection:
            qm_tools = qm_selections[cell_name].get('selected_tools', [])
            annotations.append(f"🧭{len(qm_tools)}")
        # Show state changes (keys set)
        if cell_state_changes:
            state_keys = ",".join(cell_state_changes[:2])  # Limit to 2 keys
            if len(cell_state_changes) > 2:
                state_keys += f"+{len(cell_state_changes)-2}"
            annotations.append(f"📝{state_keys}")
        annotation_str = " " + " ".join(annotations) if annotations else ""

        # Override status if blocked or has errors
        if is_blocked:
            status = "⛔"
        elif has_errors:
            status = "✗"  # Failed but not blocked

        # Track running, blocked, and failed cells for styling
        is_cell_running = is_running and cell_name == running_cell
        cell_running_internal_node = None

        if is_cell_running:
            running_cell_ids.append(pid)
            # Calculate which internal node is currently executing
            if cell_progress:
                cell_running_internal_node = get_running_internal_node_id(cell_progress, pid)

        # Track cells with errors for red styling (use outer-scope list)
        if has_errors and not is_blocked:
            failed_cell_ids.append(pid)

        # Decide if cell needs composite state
        # Note: errors alone don't force composite - show as simple failed cell
        # Only render composite if there's meaningful internal structure to show
        needs_composite = (has_takes or has_reforge or has_wards or has_sub_cascades
                          or has_multi_turns or has_qm_selection or has_many_tool_calls)

        if not needs_composite:
            # Simple cell - single state
            # Add progress indicator if this cell is currently running
            progress_str = ""
            if is_cell_running and progress_indicator:
                progress_str = f" [{progress_indicator}]"
            lines.append(f'    state "{status} {sanitize_label(cell_name, 28)}{annotation_str}{progress_str}" as {pid}')

        else:
            # Composite state
            lines.append(f"    state {pid} {{")

            # Cell label with status
            label_parts = [status, sanitize_label(cell_name, 25)]
            if has_takes:
                factor = len(takes_by_cell[cell_name])
                label_parts.append(f"🔱{factor}")
            if has_reforge:
                steps = len(reforge_by_cell[cell_name])
                label_parts.append(f"🔨{steps}")
            if has_sub_cascades:
                count = len(sub_cascades_by_cell_name[cell_name])
                label_parts.append(f"📦{count}")
            # Show tool names for composite cells
            if cell_tools:
                if len(cell_tools) <= 2:
                    tool_str = ",".join(cell_tools)
                else:
                    tool_str = ",".join(cell_tools[:2]) + f"+{len(cell_tools)-2}"
                label_parts.append(f"🔧{tool_str}")
            # Add ward retry annotation for composite cells
            if has_retries:
                retry_count = ward_retries[cell_name].get('retry_count', 0)
                if retry_count > 0:
                    label_parts.append(f"🔄{retry_count}")
            # Add validation retry (loop_until) annotation for composite cells
            if has_validation_retries:
                vr = validation_retries[cell_name]
                retry_count = vr.get('retry_count', 0)
                passed = vr.get('passed', False)
                if retry_count > 0:
                    max_attempts = vr.get('max_attempts', retry_count + 1)
                    loop_icon = "✓" if passed else "✗"
                    label_parts.append(f"⟳{retry_count+1}/{max_attempts}{loop_icon}")
            # Add error indicator for composite cells
            if has_errors:
                error_count = len(errors_by_cell[cell_name])
                label_parts.append(f"❌{error_count}")
            # Add Quartermaster indicator for composite cells
            if has_qm_selection:
                qm_tool_count = len(qm_selections[cell_name].get('selected_tools', []))
                label_parts.append(f"🧭{qm_tool_count}")
            # Add state changes for composite cells
            if cell_state_changes:
                state_keys = ",".join(cell_state_changes[:2])
                if len(cell_state_changes) > 2:
                    state_keys += f"+{len(cell_state_changes)-2}"
                label_parts.append(f"📝{state_keys}")

            # Add progress indicator for running composite cells
            if is_cell_running and progress_indicator:
                label_parts.append(f"[{progress_indicator}]")

            lines.append(f"        {pid}_label : {' '.join(label_parts)}")
            lines.append("")

            # Track last node for chaining
            last_node = None
            first_internal = None

            # Track internal nodes that need running style
            running_internal_nodes = []
            if cell_running_internal_node:
                running_internal_nodes.append(cell_running_internal_node)

            # QUARTERMASTER SELECTION (show which tools were auto-selected)
            if has_qm_selection:
                qm_data = qm_selections[cell_name]
                qm_tools = qm_data.get('selected_tools', [])
                if qm_tools:
                    qm_id = f"{pid}_qm"
                    # Show up to 3 tool names
                    if len(qm_tools) <= 3:
                        tools_display = ", ".join(qm_tools)
                    else:
                        tools_display = ", ".join(qm_tools[:3]) + f" +{len(qm_tools)-3}"
                    lines.append(f"        {qm_id} : 🧭 manifest: {sanitize_label(tools_display, 30)}")
                    if last_node:
                        lines.append(f"        {last_node} --> {qm_id}")
                    else:
                        first_internal = qm_id
                    last_node = qm_id

            # PRE-WARDS
            if has_wards and wards_by_cell[cell_name]["pre"]:
                for j, ward in enumerate(wards_by_cell[cell_name]["pre"]):
                    ward_meta = extract_metadata(ward)
                    validator = ward_meta.get("validator", "check")
                    valid = ward_meta.get("valid", True)
                    mode = ward_meta.get("mode", "blocking")

                    mode_icon = "🛡️" if mode == "blocking" else ("🔄" if mode == "retry" else "ℹ️")
                    status_mark = "✓" if valid else "✗"

                    ward_id = f"{pid}_pre{j}"
                    lines.append(f"        {ward_id} : {mode_icon} {sanitize_label(validator, 12)} {status_mark}")

                    if last_node:
                        lines.append(f"        {last_node} --> {ward_id}")
                    else:
                        first_internal = ward_id
                    last_node = ward_id

            # TAKES
            if has_takes:
                attempts = takes_by_cell[cell_name]
                attempts_info = {}
                winner_index = None

                for sa in attempts:
                    sa_meta = extract_metadata(sa)
                    idx = sa_meta.get("take_index", 0)
                    is_winner = sa_meta.get("is_winner", False)
                    attempts_info[idx] = {"is_winner": is_winner}
                    if is_winner:
                        winner_index = idx

                # Fork
                fork_id = f"{pid}_fork"
                lines.append(f"        state {fork_id} <<fork>>")

                if last_node:
                    lines.append(f"        {last_node} --> {fork_id}")
                else:
                    first_internal = fork_id

                # Attempts - include mutation info if available
                cell_mutations = take_mutations.get(cell_name, {})
                for idx in sorted(attempts_info.keys()):
                    info = attempts_info[idx]
                    marker = " ✓" if info["is_winner"] else ""

                    # Check for mutation info
                    mutation_info = cell_mutations.get(idx, {})
                    mutation_applied = mutation_info.get('mutation_applied')

                    # Determine mutation label (shortened)
                    mutation_label = ""
                    if mutation_applied:
                        # Shorten the mutation description
                        if len(mutation_applied) > 15:
                            mutation_label = f" [{mutation_applied[:12]}...]"
                        else:
                            mutation_label = f" [{mutation_applied}]"
                    elif idx == 0:
                        mutation_label = " [baseline]"

                    # Always try to load sub-cascade diagram for this take attempt
                    # For cell-level takes with sub-cascades, the session ID pattern is:
                    # {parent_session_id}_sub_{take_index}
                    sub_session_id = f"{echo.session_id}_sub_{idx}"
                    sub_mermaid = load_sub_cascade_mermaid(sub_session_id)

                    if sub_mermaid:
                        # Render as composite state with embedded sub-cascade diagram
                        attempt_id = f"{pid}_a{idx}"
                        lines.append(f"        state {attempt_id} {{")
                        lines.append(f"            {attempt_id}_label : 📦 Attempt #{idx+1}{mutation_label}{marker}")
                        lines.append("")

                        # Prefix all state IDs in sub-cascade to prevent collisions
                        prefixed_mermaid = prefix_state_ids(sub_mermaid, f"sub{idx}")

                        # Embed with proper indentation (12 spaces for nesting)
                        for line in prefixed_mermaid.split('\n'):
                            if line.strip():
                                lines.append(f"            {line}")

                        lines.append(f"        }}")
                        lines.append(f"        class {attempt_id} sub_cascade")
                        lines.append(f"        {fork_id} --> {attempt_id}")
                    else:
                        # Fallback to simple node if sub-cascade diagram not found
                        lines.append(f"        {fork_id} --> {pid}_a{idx}")
                        lines.append(f"        {pid}_a{idx} : #{idx+1}{mutation_label}{marker}")

                # Join
                join_id = f"{pid}_join"
                lines.append(f"        state {join_id} <<join>>")
                for idx in sorted(attempts_info.keys()):
                    lines.append(f"        {pid}_a{idx} --> {join_id}")

                # Evaluator
                eval_id = f"{pid}_eval"
                lines.append(f"        {join_id} --> {eval_id}")
                lines.append(f"        {eval_id} : ⚖️ Evaluate")

                # Winner
                if winner_index is not None:
                    winner_id = f"{pid}_winner"
                    lines.append(f"        {eval_id} --> {winner_id}")
                    lines.append(f"        {winner_id} : ★ #{winner_index+1}")
                    last_node = winner_id
                else:
                    last_node = eval_id

                # REFORGE (chains after takes)
                if has_reforge:
                    for step_num in sorted(reforge_by_cell[cell_name].keys()):
                        step_data = reforge_by_cell[cell_name][step_num]
                        rf_attempts = step_data.get("attempts", [])
                        rf_winner = step_data.get("winner")

                        if rf_attempts:
                            rf_id = f"{pid}_rf{step_num}"
                            rf_factor = len(rf_attempts)
                            rf_winner_meta = extract_metadata(rf_winner) if rf_winner else {}
                            rf_winner_index = rf_winner_meta.get("winner_index")

                            # Nested composite for reforge step
                            lines.append(f"        state {rf_id} {{")
                            lines.append(f"            {rf_id}_label : 🔨 Reforge {step_num}")
                            lines.append(f"            [*] --> {rf_id}_fork")
                            lines.append(f"            state {rf_id}_fork <<fork>>")

                            for rf_att in sorted(rf_attempts, key=lambda a: extract_metadata(a).get("attempt_index", 0)):
                                rf_att_meta = extract_metadata(rf_att)
                                rf_idx = rf_att_meta.get("attempt_index", 0)
                                rf_is_winner = (rf_winner_index is not None and rf_idx == rf_winner_index)
                                rf_marker = " ✓" if rf_is_winner else ""
                                lines.append(f"            {rf_id}_fork --> {rf_id}_a{rf_idx}")
                                lines.append(f"            {rf_id}_a{rf_idx} : R{rf_idx+1}{rf_marker}")

                            lines.append(f"            state {rf_id}_join <<join>>")
                            for j in range(rf_factor):
                                lines.append(f"            {rf_id}_a{j} --> {rf_id}_join")

                            if rf_winner_index is not None:
                                lines.append(f"            {rf_id}_join --> {rf_id}_winner")
                                lines.append(f"            {rf_id}_winner : ★ R{rf_winner_index+1}")
                                lines.append(f"            {rf_id}_winner --> [*]")
                            else:
                                lines.append(f"            {rf_id}_join --> [*]")

                            lines.append(f"        }}")
                            lines.append(f"        {last_node} --> {rf_id}")
                            last_node = rf_id

            # SUB-CASCADES
            if has_sub_cascades:
                for sc_idx, sc_data in enumerate(sub_cascades_by_cell_name[cell_name]):
                    sc_cascade = sc_data["cascade"]
                    sc_cells = sc_data["cells"]
                    sc_name = sc_data.get("sub_echo", f"sub_{sc_idx}")

                    sc_content = sc_cascade.get("content", "")
                    sc_cascade_name = sc_content.replace("Cascade: ", "") if sc_content.startswith("Cascade: ") else sc_content
                    sc_id = f"{pid}_sc{sc_idx}"

                    # Try to load the complete mermaid diagram for this sub-cascade
                    sub_cascade_mermaid = load_sub_cascade_mermaid(sc_name)

                    # Nested composite for sub-cascade
                    lines.append(f"        state {sc_id} {{")
                    lines.append(f"            {sc_id}_label : 📦 Sub-Cascade: {sanitize_label(sc_cascade_name, 20)}")
                    lines.append("")

                    if sub_cascade_mermaid:
                        # Prefix all state IDs in sub-cascade to prevent collisions
                        prefixed_mermaid = prefix_state_ids(sub_cascade_mermaid, f"sc{sc_idx}")

                        # Embed the complete sub-cascade state diagram
                        # Add extra indentation for proper nesting (12 spaces = 3 levels of 4-space indents)
                        for line in prefixed_mermaid.split('\n'):
                            if line.strip():
                                # Add extra indentation to nest within parent composite state
                                lines.append(f"            {line}")
                    elif sc_cells:
                        # Fallback: Render sub-cascade cells as simple nodes
                        sc_cell_ids = []
                        for sp_idx, sp_entry in enumerate(sc_cells):
                            sp_content = sp_entry.get("content", "")
                            sp_name = sp_content.replace("Cell: ", "") if sp_content.startswith("Cell: ") else sp_content
                            sp_id = f"{sc_id}_p{sp_idx}"
                            sc_cell_ids.append(sp_id)

                            # Check if sub-cell is complete (in parent lineage or has output)
                            sp_status = "✓" if sp_name in completed_cells else "○"
                            lines.append(f"            {sp_id} : {sp_status} {sanitize_label(sp_name, 18)}")

                        # Connect sub-cells
                        if sc_cell_ids:
                            lines.append(f"            [*] --> {sc_cell_ids[0]}")
                            for k in range(len(sc_cell_ids) - 1):
                                lines.append(f"            {sc_cell_ids[k]} --> {sc_cell_ids[k+1]}")
                            lines.append(f"            {sc_cell_ids[-1]} --> [*]")
                    else:
                        # Empty sub-cascade placeholder
                        lines.append(f"            {sc_id}_empty : (no cells)")
                        lines.append(f"            [*] --> {sc_id}_empty")
                        lines.append(f"            {sc_id}_empty --> [*]")

                    lines.append(f"        }}")
                    lines.append(f"        class {sc_id} sub_cascade")

                    if last_node:
                        lines.append(f"        {last_node} --> {sc_id}")
                    else:
                        first_internal = sc_id
                    last_node = sc_id

            # POST-WARDS
            if has_wards and wards_by_cell[cell_name]["post"]:
                for j, ward in enumerate(wards_by_cell[cell_name]["post"]):
                    ward_meta = extract_metadata(ward)
                    validator = ward_meta.get("validator", "check")
                    valid = ward_meta.get("valid", True)
                    mode = ward_meta.get("mode", "blocking")

                    mode_icon = "🛡️" if mode == "blocking" else ("🔄" if mode == "retry" else "ℹ️")
                    status_mark = "✓" if valid else "✗"

                    ward_id = f"{pid}_post{j}"
                    lines.append(f"        {ward_id} : {mode_icon} {sanitize_label(validator, 12)} {status_mark}")

                    if last_node:
                        lines.append(f"        {last_node} --> {ward_id}")
                    else:
                        first_internal = ward_id
                    last_node = ward_id

            # TURNS (show individual turn progression for multi-turn cells)
            if has_multi_turns and not has_takes:  # Don't show turns if takes handles it
                lines.append("")
                lines.append(f"        %% Turn progression")
                turn_ids = []
                for t_idx, turn_info in enumerate(detailed_turns):
                    t_num = turn_info.get('turn_number', t_idx + 1)
                    has_tools = turn_info.get('has_tool_calls', False)
                    content_preview = turn_info.get('content_preview', '')
                    t_id = f"{pid}_t{t_idx}"
                    turn_ids.append(t_id)

                    # Show turn with tool indicator and content preview
                    tool_mark = "🔧" if has_tools else ""
                    # Add short content preview like takes do with mutations
                    context_label = ""
                    if content_preview:
                        # Truncate and clean for Mermaid
                        preview = sanitize_label(content_preview, 18)
                        if preview:
                            context_label = f" [{preview}]"
                    lines.append(f"        {t_id} : T{t_num}{context_label} {tool_mark}")

                # Connect turns in sequence
                if turn_ids:
                    if last_node:
                        lines.append(f"        {last_node} --> {turn_ids[0]}")
                    else:
                        first_internal = turn_ids[0]
                    for k in range(len(turn_ids) - 1):
                        lines.append(f"        {turn_ids[k]} --> {turn_ids[k+1]}")
                    last_node = turn_ids[-1]

            # TOOL CALLS (show individual tool calls for cells with many tools but not multi-turn)
            if has_many_tool_calls and not has_multi_turns and not has_takes:
                lines.append("")
                lines.append(f"        %% Tool calls")
                tool_ids = []
                for tc_idx, tc_info in enumerate(cell_tool_calls):
                    tool_name = tc_info.get('tool_name', 'tool')
                    success = tc_info.get('success', True)
                    tc_id = f"{pid}_tc{tc_idx}"
                    tool_ids.append(tc_id)

                    # Show tool with success/failure indicator
                    status_mark = "✓" if success else "✗"
                    # Truncate long tool names
                    short_name = tool_name[:12] + ".." if len(tool_name) > 14 else tool_name
                    lines.append(f"        {tc_id} : 🔧 {short_name} {status_mark}")

                # Connect tool calls in sequence
                if tool_ids:
                    if last_node:
                        lines.append(f"        {last_node} --> {tool_ids[0]}")
                    else:
                        first_internal = tool_ids[0]
                    for k in range(len(tool_ids) - 1):
                        lines.append(f"        {tool_ids[k]} --> {tool_ids[k+1]}")
                    last_node = tool_ids[-1]

            # ERRORS (show error nodes if errors occurred)
            if has_errors:
                lines.append("")
                lines.append(f"        %% Errors")
                for e_idx, error_info in enumerate(errors_by_cell[cell_name]):
                    error_type = error_info.get('error_type', 'error')
                    error_msg = sanitize_label(error_info.get('message', 'Error'), 25)
                    e_id = f"{pid}_err{e_idx}"

                    # Different icons for different error types
                    if error_type == "validation_error":
                        err_icon = "⚠️"
                    elif error_type == "schema_validation_failed":
                        err_icon = "📋❌"
                    else:
                        err_icon = "❌"

                    lines.append(f"        {e_id} : {err_icon} {error_msg}")

                    # Connect error nodes into the flow
                    if last_node:
                        # Errors branch off from the last node
                        lines.append(f"        {last_node} --> {e_id}")
                    else:
                        # Error is the first (and possibly only) internal node
                        if first_internal is None:
                            first_internal = e_id
                    # Update last_node so subsequent errors chain together
                    last_node = e_id

            # Connect internal flow
            if first_internal:
                lines.append(f"        [*] --> {first_internal}")
            if last_node:
                lines.append(f"        {last_node} --> [*]")

            # Apply running style to internal nodes that are currently executing
            if running_internal_nodes:
                lines.append("")
                lines.append("        %% Highlight currently executing internal node")
                for internal_node_id in running_internal_nodes:
                    lines.append(f"        class {internal_node_id} running")

            lines.append("    }")

        lines.append("")

    # =========================================================================
    # CONNECT CELLS
    # =========================================================================

    # Build cell_name -> pid mapping
    cell_name_to_pid = {}
    for cell_entry in cell_entries:
        content = cell_entry.get("content", "")
        pname = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
        cell_name_to_pid[pname] = sid(pname)

    if first_state is None and cell_ids:
        first_state = cell_ids[0]
        lines.append(f"    [*] --> {first_state}")
    elif first_state and cell_ids:
        lines.append(f"    {first_state} --> {cell_ids[0]}")

    # Track which cells were actually executed (from lineage)
    executed_cells = [item.get("cell") for item in echo.lineage]

    # Cell-to-cell transitions with routing differentiation
    for i, cell_entry in enumerate(cell_entries):
        content = cell_entry.get("content", "")
        cell_name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
        pid = cell_ids[i]

        # Check if this cell was blocked
        if cell_name in blocked_cells:
            # Add blocked terminal state
            blocked_id = f"{pid}_blocked"
            validator = blocked_cells[cell_name].get('validator', 'ward')
            reason_short = sanitize_label(blocked_cells[cell_name].get('reason', 'blocked')[:20], 20)
            lines.append(f'    state "{reason_short}" as {blocked_id}')
            lines.append(f"    {pid} --> {blocked_id} : ⛔ {validator}")
            # No further transitions from blocked cell
            continue

        # Get handoffs for this cell
        cell_handoffs = handoffs_by_cell.get(cell_name, [])
        routing_target = routing_choices.get(cell_name)

        if cell_handoffs and len(cell_handoffs) > 1:
            # Multiple handoff options - show routing choices
            for target in cell_handoffs:
                target_pid = cell_name_to_pid.get(target)
                if target_pid:
                    if routing_target == target:
                        # This was the taken path - bold arrow with checkmark
                        lines.append(f"    {pid} --> {target_pid} : ✓ route")
                    else:
                        # Available but not taken - note syntax for dashed (Mermaid state diagrams don't support dashed, so we use note)
                        # Use different notation to indicate not-taken
                        lines.append(f"    {pid} --> {target_pid} : ○ available")
        elif cell_handoffs and len(cell_handoffs) == 1:
            # Single handoff - show taken path
            target = cell_handoffs[0]
            target_pid = cell_name_to_pid.get(target)
            if target_pid:
                lines.append(f"    {pid} --> {target_pid}")
        elif i + 1 < len(cell_ids):
            # Default sequential - connect to next cell
            lines.append(f"    {cell_ids[i]} --> {cell_ids[i+1]}")

    # =========================================================================
    # SELF-LOOP ARROWS FOR VALIDATION RETRIES
    # =========================================================================
    # Add self-loop arrows for cells that had validation retries (loop_until / max_attempts)
    lines.append("")
    lines.append("    %% Validation retry loops")
    for cell_entry in cell_entries:
        content = cell_entry.get("content", "")
        pname = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
        pid = cell_name_to_pid.get(pname)

        if pname in validation_retries and pid:
            vr = validation_retries[pname]
            retry_count = vr.get('retry_count', 0)
            if retry_count > 0:
                max_attempts = vr.get('max_attempts', retry_count + 1)
                passed = vr.get('passed', False)
                loop_status = "✓" if passed else "✗"
                # Self-loop showing retry behavior
                lines.append(f"    {pid} --> {pid} : ⟳ retry {retry_count}x {loop_status}")

    # =========================================================================
    # CHECKPOINT VISUALIZATION (Human-in-the-Loop)
    # =========================================================================
    # Show checkpoints and their resume connections
    if checkpoint_entries or checkpoint_resume_entries:
        lines.append("")
        lines.append("    %% Checkpoint nodes (Human-in-the-Loop)")

        # Map cell trace IDs to checkpoint info
        checkpoints_by_cell: Dict[str, List[Dict]] = {}
        for cp_entry in checkpoint_entries:
            cp_meta = extract_metadata(cp_entry)
            # Find the cell this checkpoint belongs to by parent_id
            parent_id = cp_entry.get("parent_id")
            cell_name = cell_trace_to_name.get(parent_id, "unknown")
            if cell_name not in checkpoints_by_cell:
                checkpoints_by_cell[cell_name] = []
            checkpoints_by_cell[cell_name].append(cp_entry)

        # Track checkpoint IDs for styling
        checkpoint_node_ids = []
        resume_node_ids = []

        # Render checkpoint nodes for cells that have them
        for cell_name, cp_list in checkpoints_by_cell.items():
            pid = cell_name_to_pid.get(cell_name)
            if not pid:
                continue

            for cp_idx, cp_entry in enumerate(cp_list):
                cp_meta = extract_metadata(cp_entry)
                checkpoint_id = cp_meta.get("checkpoint_id", f"cp_{cp_idx}")
                checkpoint_type = cp_meta.get("checkpoint_type", "cell_input")
                cp_node_id = f"{pid}_cp{cp_idx}"
                checkpoint_node_ids.append(cp_node_id)

                # Show checkpoint type indicator
                if checkpoint_type == "take_eval":
                    cp_label = "⏸️ HITL Take Eval"
                else:
                    cp_label = "⏸️ HITL Input"

                lines.append(f'    state "{cp_label}" as {cp_node_id}')
                lines.append(f"    {pid} --> {cp_node_id} : waiting")

        # Render resume connections
        for resume_entry in checkpoint_resume_entries:
            r_meta = extract_metadata(resume_entry)
            checkpoint_id = r_meta.get("checkpoint_id", "unknown")
            original_cell_trace = r_meta.get("original_trace_id")

            # Find which cell this connects back to
            original_cell_name = cell_trace_to_name.get(original_cell_trace, None)
            if original_cell_name:
                original_pid = cell_name_to_pid.get(original_cell_name)
                if original_pid:
                    # Find the corresponding checkpoint node
                    resume_node_id = f"{original_pid}_resume"
                    resume_node_ids.append(resume_node_id)
                    lines.append(f'    state "▶️ Resumed" as {resume_node_id}')

                    # Find checkpoint node for this cell
                    for cp_node_id in checkpoint_node_ids:
                        if cp_node_id.startswith(f"{original_pid}_cp"):
                            lines.append(f"    {cp_node_id} --> {resume_node_id} : human response")
                            break

        # Apply checkpoint styles
        if checkpoint_node_ids:
            lines.append("")
            lines.append("    %% Apply checkpoint styles")
            for cp_id in checkpoint_node_ids:
                lines.append(f"    class {cp_id} checkpoint_paused")
        if resume_node_ids:
            for r_id in resume_node_ids:
                lines.append(f"    class {r_id} checkpoint_resumed")

    # End state - only connect if last cell wasn't blocked
    if cell_ids:
        last_cell_entry = cell_entries[-1] if cell_entries else None
        if last_cell_entry:
            last_content = last_cell_entry.get("content", "")
            last_cell_name = last_content.replace("Cell: ", "") if last_content.startswith("Cell: ") else last_content
            if last_cell_name not in blocked_cells:
                lines.append(f"    {cell_ids[-1]} --> [*]")

    # =========================================================================
    # APPLY STYLING CLASSES
    # =========================================================================

    # Apply running class to running cells (thick green border)
    if running_cell_ids:
        lines.append("")
        lines.append("    %% Apply running style to active cells")
        for pid in running_cell_ids:
            lines.append(f"    class {pid} running")

    # Apply failed class to cells with errors (red border)
    if failed_cell_ids:
        lines.append("")
        lines.append("    %% Apply failed style to cells with errors")
        for pid in failed_cell_ids:
            lines.append(f"    class {pid} failed")

    return "\n".join(lines)


def generate_state_diagram_with_metadata(echo: Echo, include_click_handlers: bool = True) -> Tuple[str, Dict[str, Any]]:
    """
    Generate a Mermaid state diagram with companion metadata for interactive debugging.

    Returns:
        Tuple of (mermaid_string, metadata_dict)

    The metadata dict contains:
    - node_map: Maps Mermaid node IDs to trace_ids and metadata
    - session_id: The session this diagram represents
    - cells: List of cell info with trace_ids
    - click_handlers: If enabled, adds click callbacks to the Mermaid

    This enables:
    - Clicking nodes to navigate to log entries
    - Highlighting nodes dynamically (e.g., for current execution)
    - Querying logs by trace_id for detailed inspection
    """
    history, sub_echoes = flatten_history(echo.history)

    # Build node metadata map
    node_map: Dict[str, Dict[str, Any]] = {}

    # Helper to create safe state IDs (same as in main function)
    def sid(name: str) -> str:
        s = name.replace("-", "_").replace(" ", "_").replace(".", "_")
        s = s.replace("(", "").replace(")", "").replace(":", "").replace("/", "_")
        return s

    # Collect cell info with trace_ids
    cells_info = []
    for entry in history:
        if entry.get("node_type") == "cell":
            content = entry.get("content", "")
            cell_name = content.replace("Cell: ", "") if content.startswith("Cell: ") else content
            trace_id = entry.get("trace_id", "")
            pid = sid(cell_name)

            cell_info = {
                "node_id": pid,
                "cell_name": cell_name,
                "trace_id": trace_id,
                "parent_id": entry.get("parent_id"),
                "node_type": "cell"
            }
            cells_info.append(cell_info)
            node_map[pid] = cell_info

    # Collect turn info
    for entry in history:
        if entry.get("node_type") == "turn":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            turn_number = meta.get("turn_number", 1)
            trace_id = entry.get("trace_id", "")
            pid = sid(cell_name)
            turn_id = f"{pid}_t{turn_number - 1}"

            node_map[turn_id] = {
                "node_id": turn_id,
                "cell_name": cell_name,
                "turn_number": turn_number,
                "trace_id": trace_id,
                "parent_id": entry.get("parent_id"),
                "node_type": "turn"
            }

    # Collect take attempts
    for entry in history:
        if entry.get("node_type") == "take_attempt":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            take_index = meta.get("take_index", 0)
            is_winner = meta.get("is_winner", False)
            trace_id = entry.get("trace_id", "")
            pid = sid(cell_name)
            take_id = f"{pid}_a{take_index}"

            node_map[take_id] = {
                "node_id": take_id,
                "cell_name": cell_name,
                "take_index": take_index,
                "is_winner": is_winner,
                "trace_id": trace_id,
                "parent_id": entry.get("parent_id"),
                "node_type": "take_attempt"
            }

    # Collect ward info
    for entry in history:
        if entry.get("node_type") in ("pre_ward", "post_ward"):
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            validator = meta.get("validator", "check")
            ward_type = "pre" if entry.get("node_type") == "pre_ward" else "post"
            trace_id = entry.get("trace_id", "")
            pid = sid(cell_name)
            # Ward IDs are indexed, find next available
            ward_idx = sum(1 for k in node_map if k.startswith(f"{pid}_{ward_type}"))
            ward_id = f"{pid}_{ward_type}{ward_idx}"

            node_map[ward_id] = {
                "node_id": ward_id,
                "cell_name": cell_name,
                "validator": validator,
                "ward_type": ward_type,
                "valid": meta.get("valid"),
                "mode": meta.get("mode"),
                "trace_id": trace_id,
                "parent_id": entry.get("parent_id"),
                "node_type": entry.get("node_type")
            }

    # Collect error info
    for entry in history:
        if entry.get("node_type") in ("error", "validation_error", "schema_validation_failed"):
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            trace_id = entry.get("trace_id", "")
            pid = sid(cell_name)
            err_idx = sum(1 for k in node_map if k.startswith(f"{pid}_err"))
            err_id = f"{pid}_err{err_idx}"

            node_map[err_id] = {
                "node_id": err_id,
                "cell_name": cell_name,
                "error_type": entry.get("node_type"),
                "message": entry.get("content", "")[:100],
                "trace_id": trace_id,
                "parent_id": entry.get("parent_id"),
                "node_type": entry.get("node_type")
            }

    # Collect quartermaster info
    for entry in history:
        if entry.get("node_type") == "quartermaster_result":
            meta = extract_metadata(entry)
            cell_name = meta.get("cell_name", "unknown")
            trace_id = entry.get("trace_id", "")
            pid = sid(cell_name)
            qm_id = f"{pid}_qm"

            node_map[qm_id] = {
                "node_id": qm_id,
                "cell_name": cell_name,
                "selected_tools": meta.get("selected_skills", []),
                "reasoning": meta.get("reasoning", ""),
                "trace_id": trace_id,
                "parent_id": entry.get("parent_id"),
                "node_type": "quartermaster"
            }

    # Generate the base diagram
    mermaid_str = generate_state_diagram_string(echo)

    # Add click handlers if requested
    if include_click_handlers:
        click_lines = ["\n    %% Click handlers for interactive debugging"]
        for node_id, node_info in node_map.items():
            trace_id = node_info.get("trace_id", "")
            if trace_id:
                # Mermaid click syntax: click nodeId callback or click nodeId "url"
                # Using callback style for maximum flexibility
                click_lines.append(f'    click {node_id} call handleNodeClick("{node_id}", "{trace_id}")')

        mermaid_str += "\n".join(click_lines)

    # Get live session state for additional metadata
    live_state = get_live_session_state(echo.session_id)
    running_info = None
    if live_state and live_state.get("status") == "running":
        running_cell = live_state.get("current_cell")
        cell_progress = live_state.get("cell_progress")
        progress_indicator = format_cell_progress_indicator(cell_progress) if cell_progress else ""

        running_info = {
            "running_cell": running_cell,
            "cell_progress": cell_progress,
            "progress_indicator": progress_indicator,
            "running_internal_node": get_running_internal_node_id(
                cell_progress,
                sid(running_cell) if running_cell else ""
            ) if cell_progress else None
        }

    # Build metadata structure
    metadata = {
        "session_id": echo.session_id,
        "node_map": node_map,
        "cells": cells_info,
        "lineage": echo.lineage,
        "node_count": len(node_map),
        "has_click_handlers": include_click_handlers,
        # Live execution state (null if not running)
        "running_state": running_info
    }

    return mermaid_str, metadata


def generate_state_diagram(echo: Echo, output_path: str) -> str:
    """
    Generate a Mermaid state diagram from Echo history and write to file.

    This is an alternative to generate_mermaid() that produces a more compact,
    semantically meaningful visualization using state diagram syntax.

    Returns:
        Path to written .mmd file
    """
    content = generate_state_diagram_string(echo)

    with open(output_path, "w") as f:
        f.write(content)

    return output_path


def generate_mermaid(echo: Echo, output_path: str) -> str:
    """
    Generate a Mermaid state diagram from Echo history and write to file.

    The diagram shows:
    - Cells as composite states with internal complexity
    - Takes as fork/join parallel branches with winner highlighting
    - Reforge as nested refinement states
    - Sub-cascades as nested composite states
    - Wards as entry/exit validation states

    Also generates companion JSON files with execution graph structure.

    Returns:
        Path to written .mmd file
    """
    # Generate the state diagram (more compact than flowchart)
    mermaid_content = generate_state_diagram_string(echo)

    # Generate companion JSON file
    json_path = output_path.replace(".mmd", ".json")

    try:
        export_execution_graph_json(echo, json_path)
    except Exception as e:
        # Don't fail mermaid generation if JSON fails
        print(f"[Warning] Failed to generate execution graph JSON: {e}")

    # Validate and write with source context for debugging
    context = {
        "session_id": echo.session_id,
        "cell_count": len(echo.lineage),
        "message_count": len(echo.history),
        "has_takes": any("take_index" in str(msg) for msg in echo.history),
    }

    is_valid, path = validate_and_write_mermaid(mermaid_content, output_path, context)

    # if not is_valid:
    #     print(f"⚠️  Generated diagram may not render correctly")

    return path


def generate_mermaid_string_from_config(config: Any) -> str:
    """
    Generate a static Mermaid diagram string from CascadeConfig.
    Shows the intended structure before execution.

    This is the string-returning variant - useful for APIs.

    Args:
        config: CascadeConfig object

    Returns:
        Mermaid diagram as a string
    """
    lines = ["stateDiagram-v2"]

    # Styles - Midnight Fjord Dark Theme
    lines.extend([
        "    %% Static Structure Styles - Midnight Fjord Dark Theme",
        "    classDef cell fill:#16202A,stroke:#2DD4BF,stroke-width:2px,color:#F0F4F8;",
        "    classDef deterministic fill:#16202A,stroke:#6366f1,stroke-width:2px,color:#F0F4F8;",
        "    classDef takes fill:#16202A,stroke:#D9A553,stroke-width:2px,color:#F0F4F8;",
        "    classDef reforge fill:#16202A,stroke:#D9A553,stroke-width:2px,color:#F0F4F8;",
        "    classDef sub_cascade fill:#16202A,stroke:#a78bfa,stroke-width:2px,color:#F0F4F8;",
        "    classDef ward fill:#16202A,stroke:#f97316,stroke-width:1px,color:#F0F4F8;",
        "",
    ])

    lines.append(f"    [*] --> {config.cells[0].name.replace('-', '_').replace(' ', '_')}")
    lines.append("")

    cell_ids = []

    for cell in config.cells:
        cell_id = f"{cell.name.replace('-', '_').replace(' ', '_')}"
        cell_ids.append(cell_id)

        # Determine cell type and decorations
        is_deterministic = cell.is_deterministic()
        has_takes = cell.takes and cell.takes.factor > 1
        has_reforge = has_takes and cell.takes.reforge
        has_sub_cascades = bool(cell.sub_cascades)
        has_wards = cell.wards and (cell.wards.pre or cell.wards.post)

        # Build label with icons
        icons = []
        if has_takes:
            icons.append("🔱ToT")
        if has_reforge:
            icons.append("🔨R")
        if has_sub_cascades:
            icons.append("📦Sub")
        if has_wards:
            icons.append("🛡️W")
        if is_deterministic:
            icons.append(f"⚡{cell.tool.split(':')[0] if ':' in cell.tool else cell.tool}")

        label = f"{cell.name}"
        if icons:
            label += f" [{', '.join(icons)}]"

        # Determine style class
        if has_reforge:
            style = "reforge"
        elif has_takes:
            style = "takes"
        elif has_sub_cascades:
            style = "sub_cascade"
        elif is_deterministic:
            style = "deterministic"
        else:
            style = "cell"

        # Create state with description
        lines.append(f"    state \"{sanitize_label(label, 60)}\" as {cell_id}")
        lines.append(f"    class {cell_id} {style}")

        # Add note with cell details
        details = []
        if is_deterministic:
            details.append(f"Tool: {cell.tool}")
            if cell.routing:
                routes = ", ".join(f"{k}→{v}" for k, v in cell.routing.items())
                details.append(f"Routes: {routes}")
        else:
            if cell.skills:
                if cell.skills == "manifest":
                    details.append("Skills: Auto (Quartermaster)")
                else:
                    tools = ", ".join(cell.skills[:3])
                    if len(cell.skills) > 3:
                        tools += f"... (+{len(cell.skills) - 3})"
                    details.append(f"Tools: {tools}")

        if cell.rules and cell.rules.max_turns:
            details.append(f"Max turns: {cell.rules.max_turns}")

        if has_takes:
            details.append(f"Takes: {cell.takes.factor}x")
            if cell.takes.mode == "aggregate":
                details.append("Mode: Aggregate")

        if details:
            note_text = "\\n".join(details)
            lines.append(f"    note right of {cell_id}")
            for detail in details:
                lines.append(f"        {sanitize_label(detail, 80)}")
            lines.append(f"    end note")

        lines.append("")

    # Connect cells (handoffs and routing)
    for i, cell in enumerate(config.cells):
        current_id = cell_ids[i]

        # Check for explicit handoffs
        if cell.handoffs:
            for handoff in cell.handoffs:
                target = handoff.target if hasattr(handoff, 'target') else handoff
                # Find target cell
                for j, p in enumerate(config.cells):
                    if p.name == target:
                        target_id = cell_ids[j]
                        # Label the edge if there's a description
                        if hasattr(handoff, 'description') and handoff.description:
                            label = sanitize_label(handoff.description, 30)
                            lines.append(f"    {current_id} --> {target_id} : {label}")
                        else:
                            lines.append(f"    {current_id} --> {target_id}")
                        break
        # Check for deterministic routing
        elif cell.routing:
            for route_key, target_name in cell.routing.items():
                # Find target cell
                for j, p in enumerate(config.cells):
                    if p.name == target_name:
                        target_id = cell_ids[j]
                        lines.append(f"    {current_id} --> {target_id} : {route_key}")
                        break
        # Default sequential connection
        elif i + 1 < len(cell_ids):
            lines.append(f"    {current_id} --> {cell_ids[i+1]}")

    # Add final state if last cell has no handoffs
    last_cell = config.cells[-1]
    last_id = cell_ids[-1]
    if not last_cell.handoffs and not last_cell.routing:
        lines.append(f"    {last_id} --> [*]")

    return "\n".join(lines)


def generate_mermaid_from_config(config: Any, output_path: str) -> str:
    """
    Generate a static Mermaid diagram from CascadeConfig and write to file.
    Shows the intended structure before execution.

    This is the file-writing variant - wraps generate_mermaid_string_from_config.

    Args:
        config: CascadeConfig object
        output_path: Path to write the .mmd file

    Returns:
        Path to the written file
    """
    mermaid_content = generate_mermaid_string_from_config(config)

    with open(output_path, "w") as f:
        f.write(mermaid_content)

    return output_path
