"""AST to GAMS expression converter.

This module converts IR expression AST nodes to GAMS syntax strings.
Handles operator precedence, function calls, and all AST node types including MultiplierRef.

Also handles index aliasing for sum expressions to avoid GAMS Error 125
("Set is under control already") when an equation's domain index is reused
in a nested sum.
"""

from src.ir.ast import (
    Binary,
    Call,
    Const,
    DollarConditional,
    EquationRef,
    Expr,
    IndexOffset,
    MultiplierRef,
    ParamRef,
    SetMembershipTest,
    Sum,
    SymbolRef,
    Unary,
    VarRef,
)

# Operator precedence levels (higher = tighter binding)
PRECEDENCE = {
    "or": 1,
    "and": 2,
    "=": 3,
    "<>": 3,
    "<": 3,
    ">": 3,
    "<=": 3,
    ">=": 3,
    "+": 4,
    "-": 4,
    "*": 5,
    "/": 5,
    "^": 6,  # Power operator (highest precedence)
    "**": 6,  # Alternative power operator syntax (same precedence as ^)
}


def _format_numeric(value: int | float) -> str:
    """Format numeric value for GAMS, avoiding unnecessary decimals.

    Args:
        value: Numeric value to format

    Returns:
        Formatted string representation

    Examples:
        >>> _format_numeric(3.0)
        '3'
        >>> _format_numeric(3.14)
        '3.14'
        >>> _format_numeric(5)
        '5'
    """
    if isinstance(value, (int, float)) and value == int(value):
        return str(int(value))
    return str(value)


def _quote_indices(indices: tuple[str, ...]) -> list[str]:
    """Quote element labels in index tuples for GAMS syntax.

    This function uses a heuristic to distinguish between domain variables
    (set indices like 'i', 'j', 'nodes') and element labels (specific values like 'i1', 'H2').

    Heuristic:
    - All-lowercase identifier-like names (letters/underscores only) are domain variables → unquoted
    - Names containing digits, uppercase letters, or special chars are element labels → quoted
    - Indices that arrive already quoted (e.g., "demand") are always element labels → quoted

    This addresses GAMS compilation errors from inconsistent quoting (Error 171,
    Error 340) while preserving correct behavior for indexed equations.

    Also handles indices that may already be quoted from the parser (e.g., "demand")
    by stripping existing quotes before re-quoting to avoid double-quoting.

    Args:
        indices: Tuple of index identifiers

    Returns:
        List of appropriately quoted/unquoted indices

    Examples:
        >>> _quote_indices(("i",))
        ['i']
        >>> _quote_indices(("i1",))
        ['"i1"']
        >>> _quote_indices(("i", "j"))
        ['i', 'j']
        >>> _quote_indices(("nodes",))
        ['nodes']
        >>> _quote_indices(("flow_var",))
        ['flow_var']
        >>> _quote_indices(("H", "H2", "H2O"))
        ['"H"', '"H2"', '"H2O"']
        >>> _quote_indices(("c",))
        ['c']
        >>> _quote_indices(("OH",))
        ['"OH"']
        >>> _quote_indices(('"demand"',))
        ['"demand"']
        >>> _quote_indices(('"x"', '"y"'))
        ['"x"', '"y"']
    """
    result = []
    for idx in indices:
        # Strip ALL layers of quotes to handle double-quoted indices like ""cost""
        # This can happen when indices pass through multiple transformations
        was_quoted = False
        idx_clean = idx
        while True:
            if idx_clean.startswith('"') and idx_clean.endswith('"') and len(idx_clean) >= 2:
                idx_clean = idx_clean[1:-1]
                was_quoted = True
            elif idx_clean.startswith("'") and idx_clean.endswith("'") and len(idx_clean) >= 2:
                idx_clean = idx_clean[1:-1]
                was_quoted = True
            else:
                break

        # If it was quoted, it's an element label - always quote it (single layer)
        if was_quoted:
            result.append(f'"{idx_clean}"')
        # All-lowercase identifier (letters and underscores only) = domain variable, don't quote
        # This handles patterns like sum(i, ...), x(nodes), flow(years)
        # Element labels typically contain digits (i1, H2) or uppercase (H, OH, H2O)
        elif idx_clean.replace("_", "").isalpha() and idx_clean.islower():
            result.append(idx_clean)
        else:
            # Everything else is an element label, quote it for consistency
            result.append(f'"{idx_clean}"')
    return result


def _needs_parens(parent_op: str | None, child_op: str | None, is_right: bool = False) -> bool:
    """Determine if child expression needs parentheses.

    Args:
        parent_op: Operator of parent expression (None if no parent)
        child_op: Operator of child expression (None if not a binary op)
        is_right: Whether child is the right operand of parent

    Returns:
        True if parentheses are needed
    """
    if parent_op is None or child_op is None:
        return False

    parent_prec = PRECEDENCE.get(parent_op, 0)
    child_prec = PRECEDENCE.get(child_op, 0)

    # Lower precedence always needs parens
    if child_prec < parent_prec:
        return True

    # Equal precedence on right side needs parens for non-associative ops
    # e.g., a - (b - c) vs a - b - c
    if child_prec == parent_prec and is_right:
        # Subtraction and division are left-associative
        if parent_op in ("-", "/", "^"):
            return True

    return False


def expr_to_gams(expr: Expr, parent_op: str | None = None, is_right: bool = False) -> str:
    """Convert an AST expression to GAMS syntax.

    Args:
        expr: Expression AST node
        parent_op: Operator of parent expression (for precedence handling)
        is_right: Whether this is the right operand of parent

    Returns:
        GAMS expression string

    Examples:
        >>> expr_to_gams(Const(3.14))
        '3.14'
        >>> expr_to_gams(VarRef("x", ("i",)))
        'x(i)'
        >>> expr_to_gams(Binary("+", Const(1), Const(2)))
        '1 + 2'
        >>> expr_to_gams(Binary("^", VarRef("x", ()), Const(2)))
        'x ** 2'
    """
    match expr:
        case Const(value):
            return _format_numeric(value)

        case SymbolRef(name):
            return name

        case VarRef() as var_ref:
            if var_ref.indices:
                quoted_indices = _quote_indices(var_ref.indices_as_strings())
                indices_str = ",".join(quoted_indices)
                return f"{var_ref.name}({indices_str})"
            return var_ref.name

        case ParamRef() as param_ref:
            if param_ref.indices:
                quoted_indices = _quote_indices(param_ref.indices_as_strings())
                indices_str = ",".join(quoted_indices)
                return f"{param_ref.name}({indices_str})"
            return param_ref.name

        case MultiplierRef() as mult_ref:
            if mult_ref.indices:
                quoted_indices = _quote_indices(mult_ref.indices_as_strings())
                indices_str = ",".join(quoted_indices)
                return f"{mult_ref.name}({indices_str})"
            return mult_ref.name

        case EquationRef() as eq_ref:
            # Equation attribute access: eq.m, eq.l('1'), etc.
            if eq_ref.indices:
                quoted_indices = _quote_indices(eq_ref.indices_as_strings())
                indices_str = ",".join(quoted_indices)
                return f"{eq_ref.name}.{eq_ref.attribute}({indices_str})"
            return f"{eq_ref.name}.{eq_ref.attribute}"

        case Unary(op, child):
            child_str = expr_to_gams(child, parent_op=op)
            # GAMS unary operators: +, -
            # For unary minus, ALWAYS convert to multiplication form to avoid GAMS Error 445
            # ("More than one operator in a row"). This happens when unary minus follows
            # other operators like ".." (equation definition), "+", "-", etc.
            # Solution: -(expr) becomes ((-1) * (expr)) which GAMS parses correctly.
            if op == "-":
                # Wrap child in parentheses if it's a complex expression (Binary or Call)
                # or a negative constant (to avoid ((-1) * -5) which has two operators in a row)
                # Simple cases like -x become (-1) * x, complex like -(a+b) become (-1) * (a+b)
                is_negative_const = isinstance(child, Const) and child.value < 0
                if isinstance(child, (Binary, Call)) or is_negative_const:
                    return f"((-1) * ({child_str}))"
                else:
                    return f"((-1) * {child_str})"
            # Unary plus can be passed through directly
            return f"{op}{child_str}"

        case Binary(op, left, right):
            # Convert power operator to GAMS syntax
            # Handle both ^ and ** (term collection may generate **)
            if op in ("^", "**"):
                # GAMS uses ** for exponentiation
                left_str = expr_to_gams(left, parent_op=op, is_right=False)
                right_str = expr_to_gams(right, parent_op=op, is_right=True)

                # Determine if we need parentheses for the whole expression
                needs_parens = _needs_parens(parent_op, op, is_right)
                result = f"{left_str} ** {right_str}"
                return f"({result})" if needs_parens else result

            # Special handling for subtraction of negative constants
            # Convert "x - (-5)" to "x + 5" to avoid double operators
            if op == "-" and isinstance(right, Const) and right.value < 0:
                left_str = expr_to_gams(left, parent_op="+", is_right=False)
                # Negate the negative value to get positive
                right_val = -right.value
                right_str = _format_numeric(right_val)
                # Use addition instead
                needs_parens = _needs_parens(parent_op, "+", is_right)
                result = f"{left_str} + {right_str}"
                return f"({result})" if needs_parens else result

            # Other binary operators
            left_str = expr_to_gams(left, parent_op=op, is_right=False)
            right_str = expr_to_gams(right, parent_op=op, is_right=True)

            # Special handling: wrap negative constants in parentheses to avoid GAMS Error 445
            # ("More than one operator in a row").
            #
            # For multiplication/division, GAMS cannot parse negative constants without parens:
            # - Left operand: "a + -1 * b" is invalid → becomes "a + (-1) * b"
            # - Right operand: "y * -1" is invalid → becomes "y * (-1)"
            #
            # Other operators like exponentiation (**) don't currently need this fix.
            if op in ("*", "/") and isinstance(left, Const) and left.value < 0:
                left_str = f"({left_str})"
            if op in ("*", "/") and isinstance(right, Const) and right.value < 0:
                right_str = f"({right_str})"

            # Determine if we need parentheses
            needs_parens = _needs_parens(parent_op, op, is_right)
            result = f"{left_str} {op} {right_str}"
            return f"({result})" if needs_parens else result

        case Sum(index_sets, body):
            # GAMS: sum((i,j), body) or sum(i, body)
            body_str = expr_to_gams(body)
            if len(index_sets) == 1:
                return f"sum({index_sets[0]}, {body_str})"
            indices_str = ",".join(index_sets)
            return f"sum(({indices_str}), {body_str})"

        case Call(func, args):
            # Function calls: exp(x), log(x), sqrt(x), etc.
            args_str = ", ".join(expr_to_gams(arg) for arg in args)
            return f"{func}({args_str})"

        case DollarConditional(value_expr, condition):
            # Dollar conditional: value_expr$condition
            # Evaluates to value_expr if condition is non-zero, otherwise 0
            value_str = expr_to_gams(value_expr)
            condition_str = expr_to_gams(condition)
            # Parenthesize value if it's a complex expression to avoid precedence issues
            if isinstance(value_expr, (Binary, Unary, DollarConditional)):
                value_str = f"({value_str})"
            return f"{value_str}${condition_str}"

        case SetMembershipTest(set_name, indices):
            # Set membership test: set_name(indices)
            # In GAMS conditional context, tests if index combination is in set
            if indices:
                indices_str = ",".join(expr_to_gams(idx) for idx in indices)
                return f"{set_name}({indices_str})"
            return set_name

        case _:
            # Fallback for unknown node types
            raise ValueError(f"Unknown expression type: {type(expr).__name__}")


def _make_alias_name(index: str) -> str:
    """Create an alias name for an index to avoid conflicts.

    Uses double underscore suffix to create unique alias names that are
    unlikely to conflict with user-defined sets.

    Args:
        index: Original index name

    Returns:
        Aliased index name

    Examples:
        >>> _make_alias_name("i")
        'i__'
        >>> _make_alias_name("nodes")
        'nodes__'
    """
    return f"{index}__"


def collect_index_aliases(expr: Expr, equation_domain: tuple[str, ...]) -> set[str]:
    """Collect all indices that need aliases due to conflicts with equation domain.

    Recursively traverses the expression tree to find Sum nodes whose indices
    conflict with the equation's domain indices.

    Args:
        expr: Expression to analyze
        equation_domain: Tuple of index names used in the equation's domain

    Returns:
        Set of original index names that need aliases

    Examples:
        >>> from src.ir.ast import Sum, Const
        >>> expr = Sum(("i",), Const(0))
        >>> collect_index_aliases(expr, ("i", "j"))
        {'i'}
        >>> collect_index_aliases(expr, ("k",))
        set()
    """
    if not equation_domain:
        return set()

    domain_set = set(equation_domain)
    aliases_needed: set[str] = set()

    def _collect(e: Expr) -> None:
        match e:
            case Sum(index_sets, body):
                # Check for conflicts
                for idx in index_sets:
                    if idx in domain_set:
                        aliases_needed.add(idx)
                # Recurse into body
                _collect(body)

            case Binary(_, left, right):
                _collect(left)
                _collect(right)

            case Unary(_, child):
                _collect(child)

            case Call(_, args):
                for arg in args:
                    _collect(arg)

            case DollarConditional(value_expr, condition):
                _collect(value_expr)
                _collect(condition)

            case SetMembershipTest() as smt:
                for smt_idx in smt.indices:
                    _collect(smt_idx)

            case VarRef() | ParamRef() | MultiplierRef() | EquationRef() | Const() | SymbolRef():
                # Leaf nodes - no nested sums
                pass

            case _:
                pass

    _collect(expr)
    return aliases_needed


def resolve_index_conflicts(expr: Expr, equation_domain: tuple[str, ...]) -> Expr:
    """Resolve index conflicts by replacing conflicting sum indices with aliases.

    Creates a new expression tree where any Sum index that conflicts with the
    equation's domain is replaced with an aliased version (e.g., "i" -> "i__").

    Also replaces any references to the original index within the sum body
    (in VarRef, ParamRef, MultiplierRef indices) with the aliased version.

    Args:
        expr: Expression to transform
        equation_domain: Tuple of index names used in the equation's domain

    Returns:
        New expression with conflicts resolved

    Examples:
        >>> from src.ir.ast import Sum, VarRef, Const
        >>> expr = Sum(("i",), VarRef("x", ("i",)))
        >>> result = resolve_index_conflicts(expr, ("i",))
        >>> # Result: Sum(("i__",), VarRef("x", ("i__",)))
    """
    if not equation_domain:
        return expr

    domain_set = set(equation_domain)

    def _resolve(e: Expr, active_aliases: dict[str, str]) -> Expr:
        """Recursively resolve conflicts, tracking active alias substitutions."""
        match e:
            case Sum(index_sets, body):
                # Check which indices conflict and need aliasing
                sum_indices: list[str] = []
                new_aliases = dict(active_aliases)  # Copy current aliases

                for idx in index_sets:
                    if idx in domain_set:
                        # This index conflicts - use alias
                        alias = _make_alias_name(idx)
                        sum_indices.append(alias)
                        new_aliases[idx] = alias
                    else:
                        sum_indices.append(idx)

                # Recurse into body with updated aliases
                new_body = _resolve(body, new_aliases)
                return Sum(tuple(sum_indices), new_body)

            case Binary(op, left, right):
                new_left = _resolve(left, active_aliases)
                new_right = _resolve(right, active_aliases)
                return Binary(op, new_left, new_right)

            case Unary(op, child):
                new_child = _resolve(child, active_aliases)
                return Unary(op, new_child)

            case Call(func, args):
                new_args = tuple(_resolve(arg, active_aliases) for arg in args)
                return Call(func, new_args)

            case DollarConditional(value_expr, condition):
                new_value = _resolve(value_expr, active_aliases)
                new_condition = _resolve(condition, active_aliases)
                return DollarConditional(new_value, new_condition)

            case SetMembershipTest(set_name, indices):
                new_indices = tuple(_resolve(idx, active_aliases) for idx in indices)
                return SetMembershipTest(set_name, new_indices)

            case VarRef() as var_ref:
                if var_ref.indices and active_aliases:
                    # Replace any aliased indices (only strings, not IndexOffset)
                    var_indices: tuple[str | IndexOffset, ...] = tuple(
                        active_aliases.get(idx, idx) if isinstance(idx, str) else idx
                        for idx in var_ref.indices
                    )
                    if var_indices != var_ref.indices:
                        return VarRef(var_ref.name, var_indices)
                return var_ref

            case ParamRef() as param_ref:
                if param_ref.indices and active_aliases:
                    param_indices: tuple[str | IndexOffset, ...] = tuple(
                        active_aliases.get(idx, idx) if isinstance(idx, str) else idx
                        for idx in param_ref.indices
                    )
                    if param_indices != param_ref.indices:
                        return ParamRef(param_ref.name, param_indices)
                return param_ref

            case MultiplierRef() as mult_ref:
                if mult_ref.indices and active_aliases:
                    mult_indices: tuple[str | IndexOffset, ...] = tuple(
                        active_aliases.get(idx, idx) if isinstance(idx, str) else idx
                        for idx in mult_ref.indices
                    )
                    if mult_indices != mult_ref.indices:
                        return MultiplierRef(mult_ref.name, mult_indices)
                return mult_ref

            case EquationRef() as eq_ref:
                if eq_ref.indices and active_aliases:
                    eq_indices: tuple[str | IndexOffset, ...] = tuple(
                        active_aliases.get(idx, idx) if isinstance(idx, str) else idx
                        for idx in eq_ref.indices
                    )
                    if eq_indices != eq_ref.indices:
                        return EquationRef(eq_ref.name, eq_indices, eq_ref.attribute)
                return eq_ref

            case Const() | SymbolRef():
                # Leaf nodes with no indices
                return e

            case _:
                return e

    return _resolve(expr, {})
