"""Display formatting: AST/variable/expression → human-readable strings.

Converts AST DataFrames (queried from the model) into human-readable
expression strings for ``Problem.display()``. The formatting functions
are pure — they take DataFrames and return strings, with no model or
engine dependencies.
"""

from __future__ import annotations

from collections import defaultdict
from typing import Any

from pandas import DataFrame

from .wire_format import (
    _CODE_TO_NAME,
    _EXPR_CODES,
    _EXPR_TYPE_LABELS,
    _HIGHER_ORDER_CODES,
    _INFIX,
    _INFIX_AND_CMP,
)

_MAX_FMT_DEPTH = 400


def format_ast_to_strings(
    expr_df: DataFrame,
    var_df: DataFrame,
    operators_df: DataFrame,
    unordered_df: DataFrame,
    ordered_data_df: DataFrame,
    ordered_hash_df: DataFrame,
) -> None:
    """Add 'expression' column with human-readable strings to *expr_df*.

    Builds lookup tables from query result DataFrames, then recursively
    formats each expression tree from its root Hash node.

    Parameters
    ----------
    expr_df : DataFrame
        Expression DataFrame with 'root' column (mutated in place).
    var_df : DataFrame
        Variable DataFrame with 'variable' and 'name' columns.
    operators_df : DataFrame
        Operator DataFrame with 'node' and 'op' columns.
    unordered_df : DataFrame
        Unordered args DataFrame with 'node' and 'arg' columns.
    ordered_data_df : DataFrame
        Ordered data args with 'node', 'index', 'value' columns.
    ordered_hash_df : DataFrame
        Ordered hash args with 'node', 'index', 'arg' columns.

    Raises
    ------
    RuntimeError
        If a root expression node has no operator entry, indicating the
        engine returned incomplete AST data.
    """
    if expr_df.empty:
        return

    # Build lookups from query results.
    variables: dict[Any, str] = dict(zip(var_df["variable"], var_df["name"])) if not var_df.empty else {}
    operators: dict[Any, int] = dict(zip(operators_df["node"], operators_df["op"])) if not operators_df.empty else {}

    # Build args dict: node -> [child values in order].
    args: dict[Any, list[Any]] = defaultdict(list)
    for node, child in unordered_df.itertuples(index=False):
        args[node].append(child)
    ordered: dict[Any, list[tuple[Any, Any]]] = defaultdict(list)
    for df in (ordered_data_df, ordered_hash_df):
        for node, idx, child in df.itertuples(index=False):
            ordered[node].append((idx, child))
    for node, indexed_children in ordered.items():
        indexed_children.sort(key=lambda x: x[0])
        args[node].extend(child for _, child in indexed_children)

    def fmt(node: Any, _depth: int = 0, _is_root: bool = False) -> str:
        if _depth > _MAX_FMT_DEPTH:
            return "<expression too deep>"
        if node in variables:
            return variables[node]
        op_code = operators.get(node)
        if op_code is None:
            if _is_root:
                raise RuntimeError(
                    "display() received incomplete expression data from the engine. "
                    "This can happen under heavy concurrent load. "
                    "Retrying may resolve the issue."
                )
            return str(node)
        op = _CODE_TO_NAME.get(op_code)
        if op is None:
            return f"<unknown_op:{op_code}>"
        children = args.get(node, [])
        is_infix = op in _INFIX
        parts: list[str] = []
        for child in children:
            s = fmt(child, _depth + 1)
            if is_infix:
                child_op = operators.get(child)
                if child_op is not None and _CODE_TO_NAME.get(child_op, "") in _INFIX_AND_CMP:
                    s = f"({s})"
            parts.append(s)
        # Sort for stable display. SOS1/SOS2 excluded (index order is meaningful).
        if op in _HIGHER_ORDER_CODES and op not in ("special_ordered_set_type_1", "special_ordered_set_type_2"):
            parts.sort()
        if op in _INFIX_AND_CMP:
            return f"{parts[0]} {op} {parts[1]}"
        return f"{op}({', '.join(parts)})"

    expr_df["expression"] = [fmt(root, _is_root=True) for root in expr_df["root"]]


def format_variable_table(var_df: Any, var_types: Any) -> list[str]:
    """Format variable DataFrame into display lines.

    Parameters
    ----------
    var_df : DataFrame
        Variable DataFrame with columns: name, type, lower, upper.
    var_types : Series
        Series of integer type codes.

    Returns
    -------
    list[str]
        Display lines (empty if no variables).
    """
    if var_df.empty:
        return []
    var_df["type"] = var_types.map(_CODE_TO_NAME)
    cols = ["name", "type"] + [c for c in ("lower", "upper") if var_df[c].notna().any()]
    return [
        "\nVariables:",
        var_df.sort_values(by="name")[cols].to_string(index=False),  # type: ignore[call-overload]
    ]


def format_expression_table(
    expr_df: Any,
    expr_types: Any,
    var_df: Any,
    ast_dfs: tuple[Any, Any, Any, Any],
) -> list[str]:
    """Format expression DataFrame into display lines grouped by type.

    Parameters
    ----------
    expr_df : DataFrame
        Expression DataFrame with type, name, root columns.
    expr_types : Series
        Series of integer type codes.
    var_df : DataFrame
        Variable DataFrame (needed for variable name lookups).
    ast_dfs : tuple
        Tuple of (operators_df, ordered_hash_df, ordered_data_df, unordered_df).

    Returns
    -------
    list[str]
        Display lines (empty if no expressions).
    """
    operators_df, ordered_hash_df, ordered_data_df, unordered_df = ast_dfs
    format_ast_to_strings(expr_df, var_df, operators_df, unordered_df, ordered_data_df, ordered_hash_df)
    has_names = "name" in expr_df.columns and expr_df["name"].notna().any()

    lines: list[str] = []
    for type_key, label in _EXPR_TYPE_LABELS:
        type_code = _EXPR_CODES[type_key]
        subset = expr_df[expr_types == type_code]
        if subset.empty:
            continue
        lines.append(f"\n{label}:")
        if has_names:
            lines.append(subset.sort_values(by="name")[["name", "expression"]].to_string(index=False))  # type: ignore[call-overload]
        else:
            lines.append("\n".join(sorted(subset["expression"])))
    return lines
