# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

from typing import TYPE_CHECKING, Any

from marimo._plugins.ui._impl.dataframes.transforms.types import (
    FilterCondition,
    RangeValue,
    Transform,
    TransformType,
)
from marimo._utils.assert_never import assert_never

if TYPE_CHECKING:
    from collections.abc import Callable


def python_print_transforms(
    df_name: str,
    all_columns: list[str],
    transforms: list[Transform],
    print_transform: Callable[[str, list[str], Transform], str],
) -> str:
    df_next_name = f"{df_name}_next"
    strs: list[str] = []
    for transform in transforms:
        strs.append(
            f"{df_next_name} = {print_transform(df_next_name, all_columns, transform)}"
        )
    return "\n".join([f"{df_next_name} = {df_name}"] + strs)


def python_print_pandas(
    df_name: str, all_columns: list[str], transform: Transform
) -> str:
    def generate_where_clause(df_name: str, where: FilterCondition) -> str:
        column_id, operator, value = (
            where.column_id,
            where.operator,
            where.value,
        )
        col = f"{df_name}[{_as_literal(column_id)}]"

        result: str
        if operator == "==" or operator == "equals":
            result = f"{col} == {_as_literal(value)}"
        elif operator == "!=" or operator == "does_not_equal":
            result = f"{col} != {_as_literal(value)}"
        elif operator in [">", ">=", "<", "<="]:
            result = f"{col} {operator} {_as_literal(value)}"
        elif operator == "contains":
            result = f"{col}.str.contains({_as_literal(value)})"
        elif operator == "regex":
            result = f"{col}.str.contains({_as_literal(value)}, regex=True)"
        elif operator == "starts_with":
            result = f"{col}.str.startswith({_as_literal(value)})"
        elif operator == "ends_with":
            result = f"{col}.str.endswith({_as_literal(value)})"
        elif operator == "in" or operator == "not_in":
            expr = f"{col}.isin({_list_of_strings(value)})"
            result = expr if operator == "in" else f"~{expr}"
        elif operator == "is_null":
            result = f"{col}.isna()"
        elif operator == "is_not_null":
            result = f"{col}.notna()"
        elif operator == "is_true":
            result = f"{col}.eq(True)"
        elif operator == "is_false":
            result = f"{col}.eq(False)"
        elif operator == "between":
            if isinstance(value, RangeValue):
                result = f"{col}.between({_as_literal(value.min)}, {_as_literal(value.max)})"
            else:
                raise ValueError(
                    f"between requires RangeValue, got {type(value)}"
                )
        elif operator == "is_empty":
            result = f"({col} == '')"
        else:
            raise ValueError(f"Unknown operator: {operator}")

        if where.negate:
            result = f"~({result})"

        return result

    if transform.type == TransformType.COLUMN_CONVERSION:
        column_id, data_type, errors = (
            transform.column_id,
            transform.data_type,
            transform.errors,
        )
        return f'{df_name}\n{df_name}[{_as_literal(column_id)}] = {df_name}[{_as_literal(column_id)}].astype("{data_type}", errors="{errors}")'

    elif transform.type == TransformType.RENAME_COLUMN:
        column_id, new_column_id = (
            transform.column_id,
            transform.new_column_id,
        )
        return f"{df_name}.rename(columns={{{_as_literal(column_id)}: {_as_literal(new_column_id)}}})"

    elif transform.type == TransformType.SORT_COLUMN:
        column_id, ascending, na_position = (
            transform.column_id,
            transform.ascending,
            transform.na_position,
        )
        args = _args_list(
            f"by={_as_literal(column_id)}",
            f"ascending={ascending}",
            f"na_position={_as_literal(na_position)}",
        )
        return f"{df_name}.sort_values({args})"

    elif transform.type == TransformType.FILTER_ROWS:
        operation, group = transform.operation, transform.where
        if not group.children:
            return df_name
        where_clauses = [
            generate_where_clause(df_name, condition)
            for condition in group.children
            if isinstance(
                condition, FilterCondition
            )  # TODO: handle nested FilterGroup (OR, negate)
        ]
        if operation == "keep_rows" and len(where_clauses) == 1:
            return f"{df_name}[{where_clauses[0]}]"
        expression = " & ".join(f"({clause})" for clause in where_clauses)
        return (
            f"{df_name}[{expression}]"
            if operation == "keep_rows"
            else f"{df_name}[~({expression})]"
        )

    elif transform.type == TransformType.AGGREGATE:
        column_ids, aggregations = (
            transform.column_ids,
            transform.aggregations,
        )
        if not column_ids:
            return f"{df_name}.agg({_list_of_strings(aggregations)})"
        # Generate code that matches narwhals behavior: columns named like 'column_agg'
        # Use pd.DataFrame to create a single-row dataframe with proper column names
        agg_parts = []
        for agg in aggregations:
            for col in column_ids:
                agg_parts.append(
                    f"{_as_literal(f'{col}_{agg}')}: [{df_name}[{_as_literal(col)}].{agg}()]"
                )
        return f"pd.DataFrame({{{', '.join(agg_parts)}}})"

    elif transform.type == TransformType.GROUP_BY:
        column_ids, aggregation, drop_na = (
            transform.column_ids,
            transform.aggregation,
            transform.drop_na,
        )
        # Use explicit aggregation columns if provided, otherwise all except group-by columns
        # Filter out group-by columns from aggregation columns to match narwhals behavior
        group_by_column_id_set = set(column_ids)
        if transform.aggregation_column_ids:
            aggregation_columns = [
                col
                for col in transform.aggregation_column_ids
                if col not in group_by_column_id_set
            ]
        else:
            aggregation_columns = [
                col for col in all_columns if col not in group_by_column_id_set
            ]
        args = _args_list(_list_of_strings(column_ids), f"dropna={drop_na}")
        group_by = f"{df_name}.groupby({args})"
        # Narwhals adds suffixes to aggregated columns like 'column_count'
        # We need to replicate this behavior by using agg() with explicit column names
        if aggregation == "count":
            agg_func = "count"
        elif aggregation == "sum":
            agg_func = "sum"
        elif aggregation == "mean":
            agg_func = "mean"
        elif aggregation == "median":
            agg_func = "median"
        elif aggregation == "min":
            agg_func = "min"
        elif aggregation == "max":
            agg_func = "max"
        else:
            assert_never(aggregation)

        # If aggregation_columns is empty after filtering, just return unique grouped columns
        # This matches narwhals behavior when agg() is called with empty list
        if not aggregation_columns:
            return f"{df_name}[{_list_of_strings(column_ids)}].drop_duplicates().reset_index(drop=True)"

        # If specific aggregation columns are provided, only aggregate those and rename explicitly.
        agg_dict = ", ".join(
            f"{_as_literal(f'{col}_{aggregation}')} : ({_as_literal(col)}, {_as_literal(agg_func)})"
            for col in aggregation_columns
        )
        return f"{group_by}.agg(**{{{agg_dict}}}).reset_index()"

    elif transform.type == TransformType.SELECT_COLUMNS:
        column_ids = transform.column_ids
        if not column_ids:
            return df_name
        return (
            f"{df_name}[{_as_literal(column_ids[0])}]"
            if len(column_ids) == 1
            else f"{df_name}[{_list_of_strings(column_ids)}]"
        )

    elif transform.type == TransformType.SAMPLE_ROWS:
        n = transform.n
        return f"{df_name}.sample(n={n})"

    elif transform.type == TransformType.SHUFFLE_ROWS:
        return f"{df_name}.sample(frac=1)"

    elif transform.type == TransformType.EXPLODE_COLUMNS:
        column_ids = transform.column_ids
        return f"{df_name}.explode({_list_of_strings(column_ids)})"

    elif transform.type == TransformType.EXPAND_DICT:
        column_id = _as_literal(transform.column_id)
        return (
            f"{df_name}\n"
            f"_dict_index = {df_name}.columns.get_loc({column_id})\n"
            f"_expanded = pd.json_normalize({df_name}.pop({column_id}).map(lambda value: {{}} if value is None or (isinstance(value, float) and value != value) else value), max_level=0).set_axis({df_name}.index, axis=0)\n"
            f"{df_name} = pd.concat([{df_name}.iloc[:, :_dict_index], _expanded, {df_name}.iloc[:, _dict_index:]], axis=1)"
        )

    elif transform.type == TransformType.UNIQUE:
        column_ids = transform.column_ids
        return f"{df_name}.drop_duplicates({_list_of_strings(column_ids)}, keep={_as_literal(transform.keep)})"

    elif transform.type == TransformType.PIVOT:
        if not transform.index_column_ids:
            index_columns = _list_of_strings(
                list(
                    filter(
                        lambda col: (
                            col not in transform.column_ids
                            and col not in transform.value_column_ids
                        ),
                        all_columns,
                    )
                )
            )
        else:
            index_columns = _list_of_strings(transform.index_column_ids)

        if not transform.value_column_ids:
            value_columns = _list_of_strings(
                list(
                    filter(
                        lambda col: (
                            col not in transform.column_ids
                            and col not in transform.index_column_ids
                        ),
                        all_columns,
                    )
                )
            )
        else:
            value_columns = _list_of_strings(transform.value_column_ids)
        column_ids = transform.column_ids
        agg_func = transform.aggregation

        args = _args_list(
            f"index={index_columns}",
            f"columns={_list_of_strings(column_ids)}",
            f"values={value_columns}",
            f"aggfunc={_as_literal(agg_func)}",
            "sort=False",
        )
        pivot_code = f"{df_name}.pivot_table({args}).sort_index(axis=0)"
        # Fill NaN with 0 for count/sum aggregations to match narwhals behavior
        fill_code = (
            f"{df_name} = {df_name}.fillna(0)"
            if agg_func in ("count", "sum")
            else ""
        )
        flatten_columns_code = (
            f"{df_name}.columns = ["
            f"f\"{{'_'.join(map(str, col)).strip()}}_{agg_func}\" "
            f'if isinstance(col, tuple) else f"{{col}}_{agg_func}" '
            f"for col in {df_name}.columns]"
        )
        reset_index_code = f"{df_name} = {df_name}.reset_index()"
        code_parts = [
            pivot_code,
            fill_code,
            flatten_columns_code,
            reset_index_code,
        ]
        return "\n".join(part for part in code_parts if part)

    assert_never(transform.type)


def python_print_polars(
    df_name: str, all_columns: list[str], transform: Transform
) -> str:
    def generate_where_clause_polars(where: FilterCondition) -> str:
        column_id, operator, value = (
            where.column_id,
            where.operator,
            where.value,
        )
        col = f"pl.col({_as_literal(column_id)})"

        result: str
        if operator == "==" or operator == "equals":
            result = f"{col} == {_as_literal(value)}"
        elif operator == "!=" or operator == "does_not_equal":
            result = f"{col} != {_as_literal(value)}"
        elif operator in [">", ">=", "<", "<="]:
            result = f"{col} {operator} {_as_literal(value)}"
        elif operator == "contains":
            result = f"{col}.str.contains({_as_literal(value)})"
        elif operator == "regex":
            result = f"{col}.str.contains({_as_literal(value)}, literal=False)"
        elif operator == "starts_with":
            result = f"{col}.str.starts_with({_as_literal(value)})"
        elif operator == "ends_with":
            result = f"{col}.str.ends_with({_as_literal(value)})"
        elif operator == "in" or operator == "not_in":
            expr = f"{col}.is_in({_list_of_strings(value)})"
            result = expr if operator == "in" else f"~{expr}"
        elif operator == "is_null":
            result = f"{col}.is_null()"
        elif operator == "is_not_null":
            result = f"{col}.is_not_null()"
        elif operator == "is_true":
            result = f"{col} == True"
        elif operator == "is_false":
            result = f"{col} == False"
        elif operator == "between":
            if isinstance(value, RangeValue):
                result = f"{col}.is_between({_as_literal(value.min)}, {_as_literal(value.max)})"
            else:
                raise ValueError(
                    f"between requires RangeValue, got {type(value)}"
                )
        elif operator == "is_empty":
            result = f'({col} == "")'
        else:
            raise ValueError(f"Unknown operator: {operator}")

        if where.negate:
            result = f"~({result})"

        return result

    if transform.type == TransformType.COLUMN_CONVERSION:
        column_id, data_type = transform.column_id, transform.data_type
        try:
            import polars.datatypes as pl_datatypes

            data_type = str(pl_datatypes.numpy_char_code_to_dtype(data_type))
        except Exception:
            pass
        return f"{df_name}.cast({{{_as_literal(column_id)}: pl.{data_type}}}, strict={transform.errors == 'raise'})"

    elif transform.type == TransformType.RENAME_COLUMN:
        column_id, new_column_id = (
            transform.column_id,
            transform.new_column_id,
        )
        # Update column names in place
        all_columns[:] = [
            str(new_column_id) if col == column_id else col
            for col in all_columns
        ]
        return f"{df_name}.rename({{{_as_literal(column_id)}: {_as_literal(new_column_id)}}})"

    elif transform.type == TransformType.SORT_COLUMN:
        column_id, ascending, na_position = (
            transform.column_id,
            transform.ascending,
            transform.na_position,
        )
        return f"{df_name}.sort({_as_literal(column_id)}, descending={not ascending}, nulls_last={na_position == 'last'})"

    elif transform.type == TransformType.FILTER_ROWS:
        operation, group = transform.operation, transform.where
        if not group.children:
            return df_name
        where_clauses = [
            generate_where_clause_polars(condition)
            for condition in group.children
            if isinstance(
                condition, FilterCondition
            )  # TODO: handle nested FilterGroup (OR, negate)
        ]
        if operation == "keep_rows" and len(where_clauses) == 1:
            return f"{df_name}.filter({where_clauses[0]})"
        expression = " & ".join(f"({clause})" for clause in where_clauses)
        return (
            f"{df_name}.filter({expression})"
            if operation == "keep_rows"
            else f"{df_name}.filter(~({expression}))"
        )

    elif transform.type == TransformType.AGGREGATE:
        column_ids, aggregations = transform.column_ids, transform.aggregations
        # Build aggregation expressions that match narwhals behavior
        # Each column/aggregation combination produces a column named "column_agg"
        agg_exprs = []
        for agg_func in aggregations:
            for column_id in column_ids:
                agg_exprs.append(
                    f"pl.col({_as_literal(column_id)}).{agg_func}().alias({_as_literal(f'{column_id}_{agg_func}')})"
                )
        return f"{df_name}.select([{', '.join(agg_exprs)}])"

    elif transform.type == TransformType.GROUP_BY:
        column_ids, aggregation = transform.column_ids, transform.aggregation
        columns = transform.aggregation_column_ids or all_columns
        aggregation_columns = [col for col in columns if col not in column_ids]
        aggs: list[str] = []
        # Use _as_literal to properly escape column names
        for column_id in aggregation_columns:
            col_ref = _as_literal(column_id)
            agg_alias = f"{column_id}_{aggregation}"
            if aggregation == "count":
                aggs.append(
                    f"pl.col({col_ref}).count().alias({_as_literal(agg_alias)})"
                )
            elif aggregation == "sum":
                aggs.append(
                    f"pl.col({col_ref}).sum().alias({_as_literal(agg_alias)})"
                )
            elif aggregation == "mean":
                aggs.append(
                    f"pl.col({col_ref}).mean().alias({_as_literal(agg_alias)})"
                )
            elif aggregation == "median":
                aggs.append(
                    f"pl.col({col_ref}).median().alias({_as_literal(agg_alias)})"
                )
            elif aggregation == "min":
                aggs.append(
                    f"pl.col({col_ref}).min().alias({_as_literal(agg_alias)})"
                )
            elif aggregation == "max":
                aggs.append(
                    f"pl.col({col_ref}).max().alias({_as_literal(agg_alias)})"
                )
        group_cols = [f"pl.col({_as_literal(col)})" for col in column_ids]
        return f"{df_name}.group_by([{', '.join(group_cols)}], maintain_order=True).agg([{', '.join(aggs)}])"

    elif transform.type == TransformType.SELECT_COLUMNS:
        column_ids = transform.column_ids
        # Update columns in place for subsequent transforms
        all_columns[:] = [str(col) for col in column_ids]
        return f"{df_name}.select({_list_of_strings(column_ids)})"

    elif transform.type == TransformType.SAMPLE_ROWS:
        n = transform.n
        return f"{df_name}.sample_n(n={n})"

    elif transform.type == TransformType.SHUFFLE_ROWS:
        return f"{df_name}.sample(fraction=1.0, shuffle=True)"

    elif transform.type == TransformType.EXPLODE_COLUMNS:
        column_ids = transform.column_ids
        return f"{df_name}.explode({_list_of_strings(column_ids)})"

    elif transform.type == TransformType.EXPAND_DICT:
        column_id = _as_literal(transform.column_id)
        return f"{df_name}.unnest({column_id})"

    elif transform.type == TransformType.UNIQUE:
        column_ids = transform.column_ids
        return f"{df_name}.unique(subset={_list_of_strings(column_ids)}, keep={_as_literal(transform.keep)})"

    elif transform.type == TransformType.PIVOT:
        if not transform.index_column_ids:
            index_column_ids = _list_of_strings(
                list(
                    filter(
                        lambda col: (
                            col not in transform.column_ids
                            and col not in transform.value_column_ids
                        ),
                        all_columns,
                    )
                )
            )
        else:
            index_column_ids = _list_of_strings(transform.index_column_ids)

        if not transform.value_column_ids:
            value_column_ids = _list_of_strings(
                list(
                    filter(
                        lambda col: (
                            col not in transform.column_ids
                            and col not in transform.index_column_ids
                        ),
                        all_columns,
                    )
                )
            )
        else:
            value_column_ids = _list_of_strings(transform.value_column_ids)

        args = _args_list(
            f"on={_list_of_strings(transform.column_ids)}",
            f"index={index_column_ids}",
            f"values={value_column_ids}",
            f"aggregate_function={_as_literal(transform.aggregation) if transform.aggregation != 'count' else _as_literal('len')}",
        )
        pivot_code = f"{df_name}.pivot({args}).sort(by={index_column_ids})"
        # Fill NaN with 0 for count/sum aggregations AFTER pivot to match narwhals behavior
        # Only fill numeric columns (not index columns which retain original types)
        fill_null_code = (
            f"{df_name} = {df_name}.with_columns(pl.selectors.numeric().fill_null(0))"
            if transform.aggregation in ["count", "sum"]
            else ""
        )
        lambda_code = (
            f'lambda col,replacements=replacements: f"{transform.value_column_ids[0]}_{{col.translate(replacements)}}_{transform.aggregation}"'
            if len(transform.value_column_ids) == 1
            else f"lambda col, replacements=replacements: f'{{col.translate(replacements)}}_{transform.aggregation}'"
        ) + f" if col not in {index_column_ids} else col"
        rename_code = (
            'replacements = str.maketrans({"{": "", "}": "", \'"\': "", ",": "_"})\n'
            f"{df_name} = {df_name}.rename({lambda_code})"
        )
        code_parts = [pivot_code, fill_null_code, rename_code]
        return "\n".join(part for part in code_parts if part)

    assert_never(transform.type)


def python_print_ibis(
    df_name: str, all_columns: list[str], transform: Transform
) -> str:
    def generate_where_clause(df_name: str, where: FilterCondition) -> str:
        column_id, operator, value = (
            where.column_id,
            where.operator,
            where.value,
        )
        col = f"{df_name}[{_as_literal(column_id)}]"

        result: str
        if operator == "==" or operator == "equals":
            result = f"({col} == {_as_literal(value)})"
        elif operator == "!=" or operator == "does_not_equal":
            result = f"({col} != {_as_literal(value)})"
        elif operator in [">", ">=", "<", "<="]:
            result = f"({col} {operator} {_as_literal(value)})"
        elif operator == "contains":
            result = f"({col}.contains({_as_literal(value)}))"
        elif operator == "regex":
            result = f"({col}.re_search({_as_literal(value)}))"
        elif operator == "starts_with":
            result = f"({col}.startswith({_as_literal(value)}))"
        elif operator == "ends_with":
            result = f"({col}.endswith({_as_literal(value)}))"
        elif operator == "in" or operator == "not_in":
            expr = f"({col}.isin({_list_of_strings(value)}))"
            result = expr if operator == "in" else f"~{expr}"
        elif operator == "is_null":
            result = f"({col}.isnull())"
        elif operator == "is_not_null":
            result = f"({col}.notnull())"
        elif operator == "is_true":
            result = f"({col} == True)"
        elif operator == "is_false":
            result = f"({col} == False)"
        elif operator == "between":
            if isinstance(value, RangeValue):
                result = f"({col}.between({_as_literal(value.min)}, {_as_literal(value.max)}))"
            else:
                raise ValueError(
                    f"between requires RangeValue, got {type(value)}"
                )
        elif operator == "is_empty":
            result = f'({col} == "")'
        else:
            raise ValueError(f"Unknown operator: {operator}")

        if where.negate:
            result = f"~({result})"

        return result

    if transform.type == TransformType.COLUMN_CONVERSION:
        column_id, data_type, errors = (
            transform.column_id,
            transform.data_type,
            transform.errors,
        )
        transform_data_type = _as_literal(data_type).replace("_", "")
        if errors == "ignore":
            return (
                f"{df_name}.mutate("
                f"ibis.coalesce("
                f"{df_name}[{_as_literal(column_id)}].cast(ibis.dtype({transform_data_type})), "
                f"{df_name}[{_as_literal(column_id)}]"
                f").name({_as_literal(column_id)}))"
            )
        else:
            return (
                f"{df_name}.mutate("
                f"{df_name}[{_as_literal(column_id)}]"
                f".cast(ibis.dtype({transform_data_type}))"
                f".name({_as_literal(column_id)}))"
            )

    elif transform.type == TransformType.RENAME_COLUMN:
        column_id, new_column_id = transform.column_id, transform.new_column_id
        return f"{df_name}.rename({{{_as_literal(new_column_id)}: {_as_literal(column_id)}}})"

    elif transform.type == TransformType.SORT_COLUMN:
        column_id, ascending = transform.column_id, transform.ascending
        return f"{df_name}.order_by([{df_name}[{_as_literal(column_id)}].{'asc' if ascending else 'desc'}()])"

    elif transform.type == TransformType.FILTER_ROWS:
        group, operation = transform.where, transform.operation
        expressions = [
            generate_where_clause(df_name, condition)
            for condition in group.children
            if isinstance(
                condition, FilterCondition
            )  # TODO: handle nested FilterGroup (OR, negate)
        ]
        expression = " & ".join(expressions)
        return (
            f"{df_name}.filter({expression})"
            if operation == "keep_rows"
            else f"{df_name}.filter(~({expression}))"
        )

    elif transform.type == TransformType.AGGREGATE:
        agg_dict = []
        for agg_func in transform.aggregations:
            for column_id in transform.column_ids:
                name = f"{column_id}_{agg_func}"
                agg_dict.append(
                    f"'{name}' : {df_name}['{column_id}'].{agg_func}()"
                )
        return f"{df_name}.agg(**{{{','.join(agg_dict)}}})"

    elif transform.type == TransformType.GROUP_BY:
        column_ids, aggregation = transform.column_ids, transform.aggregation
        columns = transform.aggregation_column_ids or all_columns
        aggregation_columns = [col for col in columns if col not in column_ids]
        aggs: list[str] = []
        for column_id in aggregation_columns:
            agg_alias = f"{column_id}_{aggregation}"
            aggs.append(
                f'"{agg_alias}" : {df_name}["{column_id}"].{aggregation}()'
            )
        return f"{df_name}.group_by({_list_of_strings(column_ids)}).aggregate(**{{{','.join(aggs)}}})"

    elif transform.type == TransformType.SELECT_COLUMNS:
        column_ids = transform.column_ids
        return f"{df_name}.select({_list_of_strings(column_ids)})"

    elif transform.type == TransformType.SAMPLE_ROWS:
        n, seed = transform.n, transform.seed
        return f"{df_name}.sample({n} / {df_name}.count().execute(), method='row', seed={seed})"

    elif transform.type == TransformType.SHUFFLE_ROWS:
        return f"{df_name}.order_by(ibis.random())"

    elif transform.type == TransformType.EXPLODE_COLUMNS:
        column_ids = transform.column_ids
        return f"{df_name}.unnest({_list_of_strings(column_ids)})"

    elif transform.type == TransformType.EXPAND_DICT:
        column_id = transform.column_id
        return f"{df_name}.unpack({_as_literal(column_id)})"

    elif transform.type == TransformType.UNIQUE:
        column_ids = transform.column_ids
        return f"{df_name}.distinct(on={_list_of_strings(column_ids)}, keep={_as_literal(transform.keep)})"

    elif transform.type == TransformType.PIVOT:
        if not transform.index_column_ids:
            index_column_ids = _list_of_strings(
                list(
                    filter(
                        lambda col: (
                            col not in transform.column_ids
                            and col not in transform.value_column_ids
                        ),
                        all_columns,
                    )
                )
            )
        else:
            index_column_ids = _list_of_strings(transform.index_column_ids)

        if not transform.value_column_ids:
            value_column_ids = _list_of_strings(
                list(
                    filter(
                        lambda col: (
                            col not in transform.column_ids
                            and col not in transform.index_column_ids
                        ),
                        all_columns,
                    )
                )
            )
        else:
            value_column_ids = _list_of_strings(transform.value_column_ids)

        column_ids = transform.column_ids
        agg_func = transform.aggregation

        pivot_code = (
            f"{df_name}.pivot_wider("
            f"names_from={_list_of_strings(column_ids)}, "
            f"id_cols={_list_of_strings(index_column_ids)}, "
            f"values_from={_list_of_strings(value_column_ids)}, "
            f"names_prefix={_as_literal(value_column_ids[0]) if len(value_column_ids) == 1 else _as_literal('')}, "
            f"values_agg={_as_literal(agg_func)})"
        )

        rename_code = f'{df_name} = {df_name}.rename(**{{f"{{col}}_{agg_func}": col for col in {df_name}.columns if col not in {_list_of_strings(index_column_ids)}}})'
        return f"{pivot_code}\n{rename_code}"

    assert_never(transform.type)


def _as_literal(value: Any) -> str:
    if isinstance(value, str):
        # escape backslashes
        value = value.replace("\\", "\\\\")
        # convert newlines to spaces
        value = value.replace("\n", " ")
        # convert \r to spaces
        value = value.replace("\r", " ")
        # escape double quotes
        value = value.replace('"', '\\"')
        # remove null bytes
        value = value.replace("\x00", "")
        return f'"{value}"'
    if value == "inf":
        return "float('inf')"
    if value == "-inf":
        return "float('-inf')"
    return f"{value}"


def _list_of_strings(value: list[Any] | Any) -> str:
    if isinstance(value, list):
        return f"[{', '.join(_as_literal(v) for v in value)}]"
    return _as_literal(value)


def _args_list(*args: str) -> str:
    return ", ".join(arg for arg in args if arg)
