from typing import TYPE_CHECKING

from pbi_parsers import dax

from pbi_core.ssas.model_tables.partition.enums import PartitionType

if TYPE_CHECKING:
    from .table import Table


def update_table_references(dax_expr: str, old_name: str, new_name: str) -> str | None:
    ast = dax.to_ast(dax_expr)
    if ast is None:
        return None

    ast_nodes = ast.find_all(
        (
            dax.exprs.ColumnExpression,
            dax.exprs.TableExpression,
            dax.exprs.HierarchyExpression,
        ),
    )
    # this will eventually be a function to handle escaping properly
    dax_new_name = f"'{new_name}'"
    for node in ast_nodes:
        if isinstance(node, dax.exprs.ColumnExpression):
            if dax.utils.get_inner_text(node.table) == old_name:
                node.table = dax.Token.from_str(dax_new_name)
        elif isinstance(node, dax.exprs.TableExpression):
            if dax.utils.get_inner_text(node.name) == old_name:
                node.name = dax.Token.from_str(dax_new_name)
        elif isinstance(node, dax.exprs.HierarchyExpression):  # noqa: SIM102
            if dax.utils.get_inner_text(node.table) == old_name:
                node.table = dax.Token.from_str(dax_new_name)
    return dax.Formatter(ast).format()


def fix_dax(table: "Table", new_name: str) -> None:
    """Fix DAX expressions for a table rename.

    Args:
        table (Table): The table being renamed.
        new_name (str): The new name for the table.

    """
    for p in table.tabular_model.partitions.find_all({"type": PartitionType.CALCULATED}):
        new_query = update_table_references(p.query_definition, table.name, new_name)
        if new_query is not None and new_query != p.query_definition:
            p.query_definition = new_query
    for m in table.tabular_model.measures:
        if isinstance(m.expression, str):
            m.expression = update_table_references(m.expression, table.name, new_name) or m.expression
        if f := m.format_string_definition():
            f.expression = update_table_references(f.expression, table.name, new_name) or f.expression
    for c in table.tabular_model.columns:
        if isinstance(c.expression, str):
            c.expression = update_table_references(c.expression, table.name, new_name) or c.expression
