"""
Formula graph utilities and structural inspection helpers.
"""

from __future__ import annotations

import re
from typing import Any, Dict, List, Optional

from g_gremlin.a1 import col_index_to_letter, parse_a1_range

# Matches references like A1, $A$1, Sheet!B2, or 'Sheet Name'!C3
_REF_RE = re.compile(
    r"(?P<sheet>'[^']+'|[A-Za-z0-9_]+)?!?(?P<cell>\$?[A-Za-z]{1,3}\$?\d+)"
)


def parse_formula_references(formula: str) -> List[str]:
    """
    Extract cell references from a formula string.
    Handles A1, $A$1, Sheet!A1, and 'Sheet Name'!A1 patterns.
    Returns unique references in order of appearance.
    """
    refs: List[str] = []
    seen: set[str] = set()
    for match in _REF_RE.finditer(formula or ""):
        sheet_raw = match.group("sheet")
        cell_raw = match.group("cell") or ""
        cell = cell_raw.replace("$", "")
        if not cell:
            continue
        if sheet_raw:
            sheet = sheet_raw.strip("'")
            ref = f"{sheet}!{cell}"
        else:
            ref = cell
        if ref not in seen:
            refs.append(ref)
            seen.add(ref)
    return refs


def classify_formula(formula: str) -> Dict[str, Any]:
    """
    Identify the patterns present in a formula.
    Returns a lightweight classification dictionary.
    """
    text = formula or ""
    upper = text.upper()
    refs = parse_formula_references(text)
    sheets_referenced = sorted(
        {ref.split("!", 1)[0] for ref in refs if "!" in ref}
    )

    uses_offset = "OFFSET(" in upper
    has_iferror = "IFERROR(" in upper
    has_cross_sheet = len(sheets_referenced) > 0

    detected: List[str] = []
    if uses_offset:
        detected.append("offset")
    if has_iferror:
        detected.append("iferror")
    if has_cross_sheet:
        detected.append("cross_sheet")

    if len(detected) > 1:
        formula_type = "compound"
    elif detected:
        formula_type = detected[0]
    else:
        formula_type = "simple"

    offset_details = _extract_offset_details(text) if uses_offset else None

    return {
        "type": formula_type,
        "references_sheets": sheets_referenced,
        "uses_offset": uses_offset,
        "offset_details": offset_details,
    }


def build_dependency_graph(
    cells: Dict[str, Dict[str, Any]], default_sheet: Optional[str] = None
) -> Dict[str, Dict[str, List[str]]]:
    """
    Build a reverse index of dependencies for a set of cells.
    Returns a mapping of cell -> {dependents: [...], references: [...]}.
    """
    graph: Dict[str, Dict[str, List[str]]] = {}
    for cell in cells:
        graph[cell] = {"dependents": [], "references": []}

    for cell, info in cells.items():
        formula = info.get("formula")
        if not (isinstance(formula, str) and formula.startswith("=")):
            continue
        raw_refs = parse_formula_references(formula)
        refs = [_normalize_reference(ref, default_sheet) for ref in raw_refs]
        graph[cell]["references"] = refs
        for ref in refs:
            if ref not in graph:
                graph[ref] = {"dependents": [], "references": []}
            if cell not in graph[ref]["dependents"]:
                graph[ref]["dependents"].append(cell)
    return graph


def build_workbook_graph(
    cells: Dict[str, Dict[str, Any]],
) -> Dict[str, Dict[str, Any]]:
    """
    Build a workbook-level dependency view across tabs.
    Returns a dict with both cell-level and sheet-level graphs.
    """
    cell_graph = build_dependency_graph(cells, default_sheet=None)
    sheet_edges: Dict[str, set[str]] = {}
    sheets_seen: set[str] = set()

    for cell, node in cell_graph.items():
        if "!" in cell:
            sheet = cell.split("!", 1)[0]
        else:
            sheet = ""
        sheets_seen.add(sheet)
        for ref in node.get("references", []):
            ref_sheet = ref.split("!", 1)[0] if "!" in ref else sheet
            if sheet and ref_sheet and ref_sheet != sheet:
                sheet_edges.setdefault(sheet, set()).add(ref_sheet)
                sheets_seen.add(ref_sheet)

    sheet_graph: Dict[str, Dict[str, List[str]]] = {}
    for sheet in sorted(s for s in sheets_seen if s):
        feeds = sorted(sheet_edges.get(sheet, set()))
        fed_by = sorted(
            [source for source, targets in sheet_edges.items() if sheet in targets]
        )
        sheet_graph[sheet] = {"feeds": feeds, "fed_by": fed_by}

    return {"cell_graph": cell_graph, "sheet_graph": sheet_graph}


def classify_cells(
    cells: Dict[str, Dict[str, Any]], graph: Dict[str, Dict[str, List[str]]]
) -> Dict[str, str]:
    """
    Classify each cell as driver/intermediary/output/static based on formula status and dependents.
    """
    result: Dict[str, str] = {}
    for cell, info in cells.items():
        node = graph.get(cell, {"dependents": []})
        has_dependents = bool(node.get("dependents"))
        formula = info.get("formula")
        is_formula = isinstance(formula, str) and formula.startswith("=")
        if has_dependents and not is_formula:
            result[cell] = "driver"
        elif has_dependents and is_formula:
            result[cell] = "intermediary"
        elif (not has_dependents) and is_formula:
            result[cell] = "output"
        else:
            result[cell] = "static"
    return result


def compute_structural_profile(
    spreadsheet_id: str,
    a1_range: str,
    values_grid: List[List[Any]],
    formulas_grid: List[List[Any]],
    formatted_grid: Optional[List[List[Any]]] = None,
) -> Dict[str, Any]:
    """
    Build a prompt-ready structural profile focused on formula mechanics and dependencies.
    """
    parsed = parse_a1_range(a1_range)
    sheet_name = parsed.sheet_name
    row_count = max(len(values_grid), len(formulas_grid), len(formatted_grid or []))
    num_cols = _max_columns(values_grid, formulas_grid, formatted_grid)
    total_cells = row_count * num_cols if row_count and num_cols else len(values_grid) * num_cols

    header_row = (formatted_grid or values_grid)[0] if (formatted_grid or values_grid) else []
    cells: Dict[str, Dict[str, Any]] = {}
    for row_idx in range(row_count):
        for col_idx in range(num_cols):
            value = _safe_grid_get(values_grid, row_idx, col_idx)
            formula_cell = _safe_grid_get(formulas_grid, row_idx, col_idx)
            display_value = _safe_grid_get(formatted_grid, row_idx, col_idx)

            has_formula = isinstance(formula_cell, str) and formula_cell.startswith("=")
            has_value = _has_content(value)
            has_display = _has_content(display_value)
            if not (has_formula or has_value or has_display):
                continue

            short_a1 = _short_a1(parsed, row_idx, col_idx)
            full_a1 = _full_address(short_a1, sheet_name)

            cells[full_a1] = {
                "full_a1": full_a1,
                "short_a1": short_a1,
                "value": value,
                "formula": formula_cell if has_formula else None,
                "display": display_value,
                "row_idx": row_idx,
                "col_idx": col_idx,
                "label": _label_for_cell(
                    row_idx, col_idx, formatted_grid, values_grid, header_row
                ),
            }

    graph_input = {
        cell: {"value": meta.get("value"), "formula": meta.get("formula")} for cell, meta in cells.items()
    }
    graph = build_dependency_graph(graph_input, default_sheet=sheet_name)
    classifications = classify_cells(graph_input, graph)

    drivers = _collect_drivers(cells, graph, classifications, sheet_name)
    outputs = _collect_outputs(cells, classifications)
    intermediaries = [cell for cell, kind in classifications.items() if kind == "intermediary"]
    static_count = sum(1 for kind in classifications.values() if kind == "static")

    patterns = _summarize_formula_patterns(cells, sheet_name)
    protected_ranges = _protected_ranges([meta["full_a1"] for meta in cells.values() if meta.get("formula")])
    tab_summary = _build_tab_summary(
        sheet_name=sheet_name,
        drivers=drivers,
        outputs=outputs,
        intermediaries_count=len(intermediaries),
        offset_count=patterns.get("offset_lag", {}).get("count", 0),
        cross_sheet_count=patterns.get("cross_sheet", {}).get("count", 0),
        cross_sheet_targets=patterns.get("cross_sheet", {}).get("sheets_referenced", []),
    )
    sections = _detect_sections(formatted_grid or values_grid, parsed, num_cols)

    structural_profile: Dict[str, Any] = {
        "spreadsheet_id": spreadsheet_id,
        "sheet_name": sheet_name,
        "range": a1_range,
        "tab_summary": tab_summary,
        "summary": {
            "total_cells": total_cells,
            "drivers": len(drivers),
            "intermediaries": len(intermediaries),
            "outputs": len(outputs),
            "static_labels": static_count,
        },
        "drivers": drivers,
        "formula_patterns": patterns,
        "protected_cells": {
            "note": "DO NOT OVERWRITE. These contain formulas.",
            "count": len(protected_ranges["cells"]),
            "ranges": protected_ranges["ranges"],
        },
        "outputs": outputs,
    }

    if sections:
        structural_profile["model_sections"] = sections
    if row_count * num_cols > 5000:
        structural_profile["warnings"] = [
            "Large sheet detected; structural scan may be approximate for very large grids."
        ]

    return structural_profile


def _safe_grid_get(grid: Optional[List[List[Any]]], row: int, col: int) -> Any:
    if grid is None:
        return None
    if row >= len(grid):
        return None
    row_vals = grid[row]
    if col >= len(row_vals):
        return None
    return row_vals[col]


def _has_content(value: Any) -> bool:
    if value is None:
        return False
    if isinstance(value, str) and value.strip() == "":
        return False
    return True


def _short_a1(parsed, row_idx: int, col_idx: int) -> str:
    row_num = parsed.start_row + row_idx
    col_num = parsed.start_col + col_idx
    col_letters = col_index_to_letter(col_num)
    return f"{col_letters}{row_num}"


def _full_address(addr: str, sheet_name: str) -> str:
    if not sheet_name:
        return addr
    if "!" in addr:
        return addr
    return f"{sheet_name}!{addr}"


def _normalize_reference(ref: str, default_sheet: Optional[str]) -> str:
    if "!" in ref or not default_sheet:
        return ref
    return f"{default_sheet}!{ref}"


def _label_for_cell(
    row_idx: int,
    col_idx: int,
    formatted_grid: Optional[List[List[Any]]],
    values_grid: List[List[Any]],
    header_row: List[Any],
) -> Optional[str]:
    """
    Heuristic: prefer the first column on the same row, then the cell's own display,
    then the header for the column.
    """
    if formatted_grid and row_idx < len(formatted_grid):
        row_vals = formatted_grid[row_idx]
        if row_vals:
            if col_idx != 0 and len(row_vals) > 0:
                candidate = row_vals[0]
                if isinstance(candidate, str) and candidate.strip():
                    return candidate.strip()
            if col_idx < len(row_vals):
                candidate = row_vals[col_idx]
                if isinstance(candidate, str) and candidate.strip():
                    return candidate.strip()

    if row_idx < len(values_grid):
        row_vals = values_grid[row_idx]
        if row_vals:
            if col_idx != 0 and len(row_vals) > 0:
                candidate = row_vals[0]
                if isinstance(candidate, str) and candidate.strip():
                    return str(candidate).strip()
            if col_idx < len(row_vals):
                candidate = row_vals[col_idx]
                if isinstance(candidate, str) and candidate.strip():
                    return str(candidate).strip()

    if col_idx < len(header_row):
        candidate = header_row[col_idx]
        if isinstance(candidate, str) and candidate.strip():
            return candidate.strip()
    return None


def _collect_drivers(
    cells: Dict[str, Dict[str, Any]],
    graph: Dict[str, Dict[str, List[str]]],
    classifications: Dict[str, str],
    sheet_name: str,
) -> List[Dict[str, Any]]:
    drivers: List[Dict[str, Any]] = []
    for cell, kind in classifications.items():
        if kind != "driver":
            continue
        meta = cells[cell]
        dependents = graph.get(cell, {}).get("dependents", [])
        drivers.append(
            {
                "cell": meta["full_a1"],
                "label": meta.get("label"),
                "value": meta.get("value"),
                "referenced_by_count": len(dependents),
                "referenced_by_sample": [
                    _full_address(dep, sheet_name) for dep in dependents[:5]
                ],
            }
        )
    return drivers


def _collect_outputs(
    cells: Dict[str, Dict[str, Any]], classifications: Dict[str, str]
) -> List[Dict[str, Any]]:
    outputs: List[Dict[str, Any]] = []
    for cell, kind in classifications.items():
        if kind != "output":
            continue
        meta = cells[cell]
        outputs.append(
            {
                "cell": meta["full_a1"],
                "formula": meta.get("formula"),
                "label": meta.get("label"),
            }
        )
    return outputs


def _summarize_formula_patterns(
    cells: Dict[str, Dict[str, Any]], sheet_name: str
) -> Dict[str, Any]:
    cross_sheet_count = 0
    sheets_referenced: set[str] = set()
    offset_count = 0
    offset_sample: Optional[str] = None
    iferror_count = 0
    iferror_sample: Optional[str] = None
    lag_drivers: List[Dict[str, Any]] = []

    for meta in cells.values():
        formula = meta.get("formula")
        if not formula:
            continue
        details = classify_formula(formula)
        if details["references_sheets"]:
            cross_sheet_count += 1
            sheets_referenced.update(details["references_sheets"])
        if details.get("uses_offset"):
            offset_count += 1
            if offset_sample is None:
                offset_sample = formula
            lag_info = details.get("offset_details") or {}
            lag_ref = lag_info.get("lag_ref")
            if lag_ref:
                lag_value = None
                lag_key = _full_address(lag_ref, sheet_name)
                lag_meta = cells.get(lag_key)
                if lag_meta:
                    lag_value = lag_meta.get("value")
                lag_drivers.append(
                    {
                        "cell": lag_key,
                        "value": lag_value,
                        "direction": lag_info.get("direction"),
                    }
                )
        if "IFERROR(" in (formula or "").upper():
            iferror_count += 1
            if iferror_sample is None:
                iferror_sample = formula

    return {
        "cross_sheet": {
            "count": cross_sheet_count,
            "sheets_referenced": sorted(sheets_referenced),
        },
        "offset_lag": {
            "count": offset_count,
            "warning": "CRITICAL: These create time-lag relationships. Do not remove.",
            "lag_drivers": lag_drivers,
            "sample": offset_sample,
        },
        "iferror": {"count": iferror_count, "sample": iferror_sample},
    }


def _protected_ranges(addresses: List[str]) -> Dict[str, Any]:
    if not addresses:
        return {"cells": [], "ranges": []}
    parsed_cells = [_parse_cell_address(addr) for addr in addresses if _parse_cell_address(addr)]
    if not parsed_cells:
        return {"cells": addresses, "ranges": []}
    rows = [item["row"] for item in parsed_cells]
    cols = [item["col"] for item in parsed_cells]
    min_row, max_row = min(rows), max(rows)
    min_col, max_col = min(cols), max(cols)
    range_str = f"{col_index_to_letter(min_col)}{min_row}:{col_index_to_letter(max_col)}{max_row}"
    return {"cells": addresses, "ranges": [range_str]}


def _parse_cell_address(addr: str) -> Optional[Dict[str, int]]:
    target = addr.split("!", 1)[-1]
    match = re.match(r"^([A-Za-z]+)(\d+)$", target)
    if not match:
        return None
    col_letters, row_str = match.groups()
    return {"row": int(row_str), "col": _letters_to_index(col_letters)}


def _letters_to_index(letters: str) -> int:
    total = 0
    for ch in letters.upper():
        total = total * 26 + (ord(ch) - ord("A") + 1)
    return total


def _detect_sections(
    grid: List[List[Any]], parsed, num_cols: int
) -> List[Dict[str, Any]]:
    if not grid:
        return []
    sections: List[Dict[str, Any]] = []
    current_label: Optional[str] = None
    start_row_idx: Optional[int] = None
    for row_idx, row in enumerate(grid):
        label = None
        if row and len(row) > 0:
            candidate = row[0]
            if isinstance(candidate, str) and candidate.strip():
                label = candidate.strip()
        if label and label != current_label:
            if current_label is not None and start_row_idx is not None:
                sections.append(
                    {
                        "label": current_label,
                        "start": _full_address(
                            _short_a1(parsed, start_row_idx, 0), parsed.sheet_name
                        ),
                        "end": _full_address(
                            _short_a1(parsed, row_idx - 1, max(0, num_cols - 1)),
                            parsed.sheet_name,
                        ),
                    }
                )
            current_label = label
            start_row_idx = row_idx
        elif label is None and current_label is not None and start_row_idx is not None:
            sections.append(
                {
                    "label": current_label,
                    "start": _full_address(
                        _short_a1(parsed, start_row_idx, 0), parsed.sheet_name
                    ),
                    "end": _full_address(
                        _short_a1(parsed, row_idx - 1, max(0, num_cols - 1)),
                        parsed.sheet_name,
                    ),
                }
            )
            current_label = None
            start_row_idx = None

    if current_label is not None and start_row_idx is not None:
        sections.append(
            {
                "label": current_label,
                "start": _full_address(
                    _short_a1(parsed, start_row_idx, 0), parsed.sheet_name
                ),
                "end": _full_address(
                    _short_a1(parsed, len(grid) - 1, max(0, num_cols - 1)),
                    parsed.sheet_name,
                ),
            }
        )
    return sections


def _max_columns(
    values_grid: List[List[Any]],
    formulas_grid: List[List[Any]],
    formatted_grid: Optional[List[List[Any]]],
) -> int:
    candidates = [
        max((len(r) for r in values_grid), default=0),
        max((len(r) for r in formulas_grid), default=0),
    ]
    if formatted_grid is not None:
        candidates.append(max((len(r) for r in formatted_grid), default=0))
    return max(candidates) if candidates else 0


def _build_tab_summary(
    sheet_name: str,
    drivers: List[Dict[str, Any]],
    outputs: List[Dict[str, Any]],
    intermediaries_count: int,
    offset_count: int,
    cross_sheet_count: int,
    cross_sheet_targets: List[str],
) -> Dict[str, Any]:
    formula_count = intermediaries_count + len(outputs)
    role = "generic"
    if formula_count == 0 and drivers:
        role = "drivers"
    elif offset_count > 0 or cross_sheet_count > 0:
        role = "calc_model"
    elif outputs and formula_count > 0:
        role = "calc_model"
    elif outputs and not drivers and formula_count > 0:
        role = "output"

    key_drivers = sorted(drivers, key=lambda d: d.get("referenced_by_count", 0), reverse=True)[:5]
    key_outputs = outputs[:5]

    return {
        "tab": sheet_name,
        "role": role,
        "key_drivers": key_drivers,
        "key_outputs": key_outputs,
        "feeds_tabs": cross_sheet_targets,
        "notes": [
            "Contains OFFSET formulas" if offset_count > 0 else None,
            "References other tabs" if cross_sheet_count > 0 else None,
        ],
    }


def _extract_offset_details(formula: str) -> Optional[Dict[str, Any]]:
    upper = formula.upper()
    start_idx = upper.find("OFFSET(")
    if start_idx == -1:
        return None
    inner = _extract_inner_args(formula, start_idx + len("OFFSET("))
    if inner is None:
        return None
    args = _split_top_level(inner)
    if len(args) < 3:
        return None
    rows_expr = args[1].strip()
    cols_expr = args[2].strip()

    if _looks_like_cell_ref(rows_expr):
        return {"direction": "row", "lag_ref": _normalize_ref(rows_expr)}
    if _looks_like_cell_ref(cols_expr):
        return {"direction": "column", "lag_ref": _normalize_ref(cols_expr)}
    return None


def _extract_inner_args(text: str, start: int) -> Optional[str]:
    depth = 1
    parts: List[str] = []
    for idx in range(start, len(text)):
        char = text[idx]
        if char == "(":
            depth += 1
        elif char == ")":
            depth -= 1
            if depth == 0:
                return "".join(parts)
        if depth >= 1:
            parts.append(char)
    return None


def _split_top_level(arg_str: str) -> List[str]:
    args: List[str] = []
    current: List[str] = []
    depth = 0
    for char in arg_str:
        if char == "," and depth == 0:
            args.append("".join(current).strip())
            current = []
            continue
        if char == "(":
            depth += 1
        elif char == ")":
            depth = max(0, depth - 1)
        current.append(char)
    if current:
        args.append("".join(current).strip())
    return args


def _looks_like_cell_ref(expr: str) -> bool:
    return bool(re.match(r"^(?:'[^']+'|[A-Za-z0-9_]+)?!?\$?[A-Za-z]{1,3}\$?\d+$", expr.strip()))


def _normalize_ref(expr: str) -> str:
    clean = expr.replace("$", "").strip()
    if "!" not in clean:
        return clean
    sheet, cell = clean.split("!", 1)
    sheet_clean = sheet.strip("'")
    return f"{sheet_clean}!{cell}"


# ─────────────────────────────────────────────────────────────────────────────
# Formula Trace Functions
# ─────────────────────────────────────────────────────────────────────────────


def detect_volatile_refs(formula: str) -> List[str]:
    """
    Detect volatile functions that cannot be statically traced.
    Returns list of volatile function names found (e.g., ['INDIRECT', 'OFFSET']).
    """
    if not formula:
        return []
    # Handle non-string values (e.g., when cell contains a number, not a formula)
    if not isinstance(formula, str):
        return []
    upper = formula.upper()
    volatile: List[str] = []
    if "INDIRECT(" in upper:
        volatile.append("INDIRECT")
    if "OFFSET(" in upper:
        volatile.append("OFFSET")
    return volatile


def trace_upstream(
    graph: Dict[str, Dict[str, List[str]]],
    cell: str,
    max_depth: int = 10,
    cells_metadata: Optional[Dict[str, Dict[str, Any]]] = None,
) -> List[Dict[str, Any]]:
    """
    Trace all cells that feed INTO the given cell (BFS traversal).

    Args:
        graph: Dependency graph from build_dependency_graph()
        cell: Starting cell (e.g., "Sheet!C36")
        max_depth: Maximum traversal depth
        cells_metadata: Optional dict with formula/value info for volatile detection

    Returns:
        List of {cell, depth, formula, is_driver, volatile, parent} ordered by depth
    """
    results: List[Dict[str, Any]] = []
    visited: set[str] = set()
    # Queue: (cell, depth, parent)
    queue: List[tuple[str, int, Optional[str]]] = [(cell, 0, None)]

    while queue:
        current, depth, parent = queue.pop(0)
        if current in visited or depth > max_depth:
            continue
        visited.add(current)

        node = graph.get(current, {})
        refs = node.get("references", [])
        dependents = node.get("dependents", [])

        # Get metadata if available
        meta = (cells_metadata or {}).get(current, {})
        formula = meta.get("formula")
        volatile = detect_volatile_refs(formula) if formula else []

        is_driver = not formula and len(dependents) > 0

        if depth > 0:  # Don't include the starting cell
            results.append({
                "cell": current,
                "depth": depth,
                "formula": formula,
                "is_driver": is_driver,
                "volatile": volatile,
                "direction": "upstream",
                "parent": parent,  # Track parent for edge drawing
            })

        # Queue upstream references
        for ref in refs:
            if ref not in visited:
                queue.append((ref, depth + 1, current))

    return results


def trace_downstream(
    graph: Dict[str, Dict[str, List[str]]],
    cell: str,
    max_depth: int = 10,
    cells_metadata: Optional[Dict[str, Dict[str, Any]]] = None,
) -> List[Dict[str, Any]]:
    """
    Trace all cells that DEPEND on the given cell (BFS traversal).

    Args:
        graph: Dependency graph from build_dependency_graph()
        cell: Starting cell (e.g., "Sheet!B55")
        max_depth: Maximum traversal depth
        cells_metadata: Optional dict with formula/value info

    Returns:
        List of {cell, depth, formula, dependent_count, parent} ordered by depth
    """
    results: List[Dict[str, Any]] = []
    visited: set[str] = set()
    # Queue: (cell, depth, parent)
    queue: List[tuple[str, int, Optional[str]]] = [(cell, 0, None)]

    while queue:
        current, depth, parent = queue.pop(0)
        if current in visited or depth > max_depth:
            continue
        visited.add(current)

        node = graph.get(current, {})
        dependents = node.get("dependents", [])

        # Get metadata if available
        meta = (cells_metadata or {}).get(current, {})
        formula = meta.get("formula")
        volatile = detect_volatile_refs(formula) if formula else []

        if depth > 0:  # Don't include the starting cell
            results.append({
                "cell": current,
                "depth": depth,
                "formula": formula,
                "dependent_count": len(dependents),
                "volatile": volatile,
                "direction": "downstream",
                "parent": parent,  # Track parent for edge drawing
            })

        # Queue downstream dependents
        for dep in dependents:
            if dep not in visited:
                queue.append((dep, depth + 1, current))

    return results


def format_trace_tree(
    results: List[Dict[str, Any]],
    root_cell: str,
    direction: str = "upstream",
) -> str:
    """
    Format trace results as indented tree for CLI output.

    Args:
        results: List from trace_upstream or trace_downstream
        root_cell: The starting cell for the trace
        direction: "upstream" or "downstream"

    Returns:
        Formatted tree string
    """
    lines: List[str] = [root_cell]
    arrow = "<-" if direction == "upstream" else "->"

    # Group by depth for better visualization
    by_depth: Dict[int, List[Dict[str, Any]]] = {}
    for item in results:
        depth = item.get("depth", 1)
        if depth not in by_depth:
            by_depth[depth] = []
        by_depth[depth].append(item)

    for depth in sorted(by_depth.keys()):
        for item in by_depth[depth]:
            indent = "  " * depth
            cell = item.get("cell", "?")

            # Build cell label
            label_parts: List[str] = [cell]

            if item.get("is_driver"):
                label_parts.append("(driver)")
            elif item.get("formula"):
                # Truncate formula for display
                formula = item["formula"]
                if len(formula) > 40:
                    formula = formula[:37] + "..."
                label_parts.append(f'"{formula}"')

            if item.get("volatile"):
                volatiles = item["volatile"]
                label_parts.append(f"[!] {', '.join(volatiles)} (volatile)")

            lines.append(f"{indent}{arrow} {' '.join(label_parts)}")

    return "\n".join(lines)


def format_trace_mermaid(
    results: List[Dict[str, Any]],
    root_cell: str,
    direction: str = "upstream",
) -> str:
    """
    Format trace results as Mermaid.js flowchart syntax.

    Args:
        results: List from trace_upstream or trace_downstream
        root_cell: The starting cell for the trace
        direction: "upstream" or "downstream"

    Returns:
        Mermaid flowchart string
    """
    lines: List[str] = ["```mermaid", "flowchart TD"]

    # Create node ID mapping (sanitize cell names for mermaid)
    def node_id(cell: str) -> str:
        return cell.replace("!", "_").replace(" ", "_").replace("'", "").replace("-", "_")

    def node_label(cell: str, item: Optional[Dict[str, Any]] = None) -> str:
        parts = [cell]
        if item:
            if item.get("is_driver"):
                parts.append("(driver)")
            if item.get("volatile"):
                parts.append(f"[!] {', '.join(item['volatile'])}")
        return "<br/>".join(parts)

    # Add root node
    root_id = node_id(root_cell)
    lines.append(f'    {root_id}["{root_cell}"]')

    # Add nodes
    for item in results:
        cell = item.get("cell", "?")
        cell_id = node_id(cell)

        # Node definition with label
        label = node_label(cell, item)
        lines.append(f'    {cell_id}["{label}"]')

    # Add edges using parent info (draws edges at ALL depths)
    edges: set[str] = set()
    for item in results:
        cell = item.get("cell", "?")
        cell_id = node_id(cell)
        parent = item.get("parent")

        if parent:
            parent_id = node_id(parent)
            # For upstream: cell feeds into parent (cell --> parent)
            # For downstream: parent feeds into cell (parent --> cell)
            if direction == "upstream":
                edge = f"    {cell_id} --> {parent_id}"
            else:
                edge = f"    {parent_id} --> {cell_id}"
            if edge not in edges:
                edges.add(edge)
                lines.append(edge)

    # Add styling for volatile nodes
    volatile_nodes = [node_id(item["cell"]) for item in results if item.get("volatile")]
    if volatile_nodes:
        lines.append(f"    style {' '.join(volatile_nodes)} fill:#ffcccc")

    lines.append("```")
    return "\n".join(lines)


def format_trace_dot(
    results: List[Dict[str, Any]],
    root_cell: str,
    direction: str = "upstream",
) -> str:
    """
    Format trace results as Graphviz DOT format.

    Args:
        results: List from trace_upstream or trace_downstream
        root_cell: The starting cell for the trace
        direction: "upstream" or "downstream"

    Returns:
        DOT format string
    """
    lines: List[str] = ["digraph formula_trace {", "    rankdir=TB;", "    node [shape=box];"]

    # Create node ID mapping
    def node_id(cell: str) -> str:
        return f'"{cell}"'

    # Add root node with special styling
    root_nid = node_id(root_cell)
    lines.append(f"    {root_nid} [style=filled, fillcolor=lightblue];")

    # Add nodes
    for item in results:
        cell = item.get("cell", "?")
        cell_nid = node_id(cell)

        attrs: List[str] = []
        if item.get("is_driver"):
            attrs.append("style=filled")
            attrs.append("fillcolor=lightgreen")
        elif item.get("volatile"):
            attrs.append("style=filled")
            attrs.append("fillcolor=lightyellow")
            attrs.append(f'label="{cell}\\n[!] {", ".join(item["volatile"])}"')

        if attrs:
            lines.append(f"    {cell_nid} [{', '.join(attrs)}];")

    # Add edges using parent info (draws edges at ALL depths)
    edges: set[str] = set()
    for item in results:
        cell = item.get("cell", "?")
        cell_nid = node_id(cell)
        parent = item.get("parent")

        if parent:
            parent_nid = node_id(parent)
            # For upstream: cell feeds into parent (cell -> parent)
            # For downstream: parent feeds into cell (parent -> cell)
            if direction == "upstream":
                edge = f"    {cell_nid} -> {parent_nid};"
            else:
                edge = f"    {parent_nid} -> {cell_nid};"
            if edge not in edges:
                edges.add(edge)
                lines.append(edge)

    lines.append("}")
    return "\n".join(lines)
