"""Expression optimizer for YARA rules."""

from __future__ import annotations

from typing import Any

from yaraast.ast.base import YaraFile
from yaraast.ast.expressions import (
    BinaryExpression,
    BooleanLiteral,
    Expression,
    Identifier,
    IntegerLiteral,
    ParenthesesExpression,
    UnaryExpression,
)
from yaraast.ast.rules import Rule
from yaraast.visitor.base import ASTTransformer

_SENTINEL = object()


def _fold_boolean(left: BooleanLiteral, right: BooleanLiteral, operator: str):
    """Fold constant boolean expressions. Returns result or _SENTINEL."""
    if operator == "and":
        return BooleanLiteral(value=left.value and right.value)
    if operator == "or":
        return BooleanLiteral(value=left.value or right.value)
    return _SENTINEL


_ARITHMETIC_OPS: dict[str, Any] = {
    "+": lambda a, b: a + b,
    "-": lambda a, b: a - b,
    "*": lambda a, b: a * b,
    "&": lambda a, b: a & b,
    "|": lambda a, b: a | b,
    "^": lambda a, b: a ^ b,
    "<<": lambda a, b: a << b,
    ">>": lambda a, b: a >> b,
}

_COMPARISON_OPS: dict[str, Any] = {
    "==": lambda a, b: a == b,
    "!=": lambda a, b: a != b,
    "<": lambda a, b: a < b,
    ">": lambda a, b: a > b,
    "<=": lambda a, b: a <= b,
    ">=": lambda a, b: a >= b,
}


def _fold_arithmetic(left_val: int, right_val: int, operator: str):
    """Fold constant integer arithmetic. Returns result or _SENTINEL."""
    if operator in _ARITHMETIC_OPS:
        return IntegerLiteral(value=_ARITHMETIC_OPS[operator](left_val, right_val))
    if operator == "/" and right_val != 0:
        return IntegerLiteral(value=int(left_val / right_val))
    if operator == "%" and right_val != 0:
        return IntegerLiteral(value=int(left_val - int(left_val / right_val) * right_val))
    return _SENTINEL


def _fold_comparison(left_val: int, right_val: int, operator: str):
    """Fold constant integer comparisons. Returns result or _SENTINEL."""
    fn = _COMPARISON_OPS.get(operator)
    if fn is not None:
        return BooleanLiteral(value=fn(left_val, right_val))
    return _SENTINEL


def _simplify_identity(node: BinaryExpression) -> Expression | None:
    """Simplify identity/annihilation patterns. Returns simplified or None."""
    left, right, op = node.left, node.right, node.operator

    # x + 0 = x, 0 + x = x
    if op == "+" and isinstance(right, IntegerLiteral) and right.value == 0:
        return left
    if op == "+" and isinstance(left, IntegerLiteral) and left.value == 0:
        return right

    # x * 1 = x, 1 * x = x
    if op == "*" and isinstance(right, IntegerLiteral) and right.value == 1:
        return left
    if op == "*" and isinstance(left, IntegerLiteral) and left.value == 1:
        return right

    # x * 0 = 0, 0 * x = 0
    if op == "*" and isinstance(right, IntegerLiteral) and right.value == 0:
        return IntegerLiteral(value=0)
    if op == "*" and isinstance(left, IntegerLiteral) and left.value == 0:
        return IntegerLiteral(value=0)

    return None


def _simplify_boolean_short_circuit(node: BinaryExpression) -> tuple[Expression | None, int]:
    """Simplify boolean short-circuit patterns. Returns (result, opt_count)."""
    count = 0

    if isinstance(node.left, BooleanLiteral):
        if node.operator == "and":
            count += 1
            if not node.left.value:
                return BooleanLiteral(value=False), count
            return node.right, count
        if node.operator == "or":
            count += 1
            if node.left.value:
                return BooleanLiteral(value=True), count
            return node.right, count

    if isinstance(node.right, BooleanLiteral):
        if node.operator == "and":
            count += 1
            if not node.right.value:
                return BooleanLiteral(value=False), count
            return node.left, count
        if node.operator == "or":
            count += 1
            if node.right.value:
                return BooleanLiteral(value=True), count
            return node.left, count

    return None, 0


class ExpressionOptimizer(ASTTransformer):
    """Optimizes expressions in YARA rules."""

    def __init__(self) -> None:
        super().__init__()
        self.optimization_count = 0

    def optimize(self, node: Expression | YaraFile) -> Expression | tuple[YaraFile, int]:
        """Optimize an expression or YaraFile.

        Args:
            node: Expression to optimize or YaraFile containing rules to optimize

        Returns:
            If Expression: optimized Expression
            If YaraFile: tuple of (optimized YaraFile, optimization count)
        """
        if isinstance(node, YaraFile):
            self.optimization_count = 0
            optimized_rules = []
            for rule in node.rules:
                optimized_rule = self._optimize_rule(rule)
                optimized_rules.append(optimized_rule)

            optimized_file = YaraFile(
                imports=node.imports,
                includes=node.includes,
                rules=optimized_rules,
            )
            return optimized_file, self.optimization_count
        # Single expression optimization
        return self.visit(node)

    def _optimize_rule(self, rule: Rule) -> Rule:
        """Optimize expressions in a rule."""
        if rule.condition:
            # Reset count for this rule to track optimizations
            optimized_condition = self.visit(rule.condition)
            # Count is incremented by visit methods for each optimization
            rule.condition = optimized_condition
        return rule

    def visit_binary_expression(self, node: BinaryExpression) -> Expression:
        """Visit BinaryExpression and optimize."""
        # Optimize children first
        node.left = self.visit(node.left)
        node.right = self.visit(node.right)

        # Constant folding for boolean operations
        if isinstance(node.left, BooleanLiteral) and isinstance(node.right, BooleanLiteral):
            result = _fold_boolean(node.left, node.right, node.operator)
            if result is not _SENTINEL:
                self.optimization_count += 1
                return result

        # Constant folding for integer arithmetic and comparisons
        if isinstance(node.left, IntegerLiteral) and isinstance(node.right, IntegerLiteral):
            left_val = node.left.value
            right_val = node.right.value

            result = _fold_arithmetic(left_val, right_val, node.operator)
            if result is not _SENTINEL:
                self.optimization_count += 1
                return result

            result = _fold_comparison(left_val, right_val, node.operator)
            if result is not _SENTINEL:
                self.optimization_count += 1
                return result

        # Identity operations
        identity = _simplify_identity(node)
        if identity is not None:
            return identity

        # Boolean simplifications
        short_circuit, count = _simplify_boolean_short_circuit(node)
        if short_circuit is not None:
            self.optimization_count += count
            return short_circuit

        return node

    def visit_unary_expression(self, node: UnaryExpression) -> Expression:
        """Visit UnaryExpression and optimize."""
        # Optimize operand first
        node.operand = self.visit(node.operand)

        # Constant folding
        if node.operator == "not" and isinstance(node.operand, BooleanLiteral):
            self.optimization_count += 1
            return BooleanLiteral(value=not node.operand.value)

        if node.operator == "-" and isinstance(node.operand, IntegerLiteral):
            self.optimization_count += 1
            return IntegerLiteral(value=-node.operand.value)

        if node.operator == "~" and isinstance(node.operand, IntegerLiteral):
            self.optimization_count += 1
            return IntegerLiteral(value=~node.operand.value)

        # Double negation elimination
        if (
            node.operator == "not"
            and isinstance(node.operand, UnaryExpression)
            and node.operand.operator == "not"
        ):
            self.optimization_count += 1
            return node.operand.operand

        return node

    def visit_parentheses_expression(self, node: ParenthesesExpression) -> Expression:
        """Visit ParenthesesExpression and optimize."""
        # Optimize inner expression
        inner = self.visit(node.expression)

        # Remove unnecessary parentheses around literals and identifiers
        if isinstance(inner, BooleanLiteral | IntegerLiteral | Identifier):
            return inner

        node.expression = inner
        return node

    # Pass-through methods for other expression types
    def visit_boolean_literal(self, node: BooleanLiteral) -> BooleanLiteral:
        return node

    def visit_integer_literal(self, node: IntegerLiteral) -> IntegerLiteral:
        return node

    def visit_identifier(self, node: Identifier) -> Identifier:
        return node

    def visit_string_identifier(self, node: Any) -> Any:
        return node

    def visit_string_count(self, node: Any) -> Any:
        return node

    def visit_string_offset(self, node: Any) -> Any:
        return node

    def visit_string_length(self, node: Any) -> Any:
        return node

    def visit_double_literal(self, node: Any) -> Any:
        return node

    def visit_string_literal(self, node: Any) -> Any:
        return node

    def visit_array_access(self, node: Any) -> Any:
        if hasattr(node, "array"):
            node.array = self.visit(node.array)
        if hasattr(node, "index"):
            node.index = self.visit(node.index)
        return node

    def visit_member_access(self, node: Any) -> Any:
        if hasattr(node, "object"):
            node.object = self.visit(node.object)
        return node

    def visit_function_call(self, node: Any) -> Any:
        if hasattr(node, "arguments"):
            node.arguments = [self.visit(arg) for arg in node.arguments]
        return node

    def visit_range_expression(self, node: Any) -> Any:
        if hasattr(node, "low"):
            node.low = self.visit(node.low)
        if hasattr(node, "high"):
            node.high = self.visit(node.high)
        return node

    def visit_set_expression(self, node: Any) -> Any:
        if hasattr(node, "elements"):
            # Optimize elements first
            node.elements = [self.visit(elem) for elem in node.elements]

            # Remove duplicates from integer literals
            seen = set()
            unique_elements = []
            duplicates_removed = 0

            for elem in node.elements:
                # Create a hashable representation for comparison
                if isinstance(elem, IntegerLiteral):
                    key = ("int", elem.value)
                elif isinstance(elem, BooleanLiteral):
                    key = ("bool", elem.value)
                elif hasattr(elem, "name"):
                    key = ("name", elem.name)
                else:
                    # For other types, keep them all (can't easily detect duplicates)
                    unique_elements.append(elem)
                    continue

                if key not in seen:
                    seen.add(key)
                    unique_elements.append(elem)
                else:
                    duplicates_removed += 1

            if duplicates_removed > 0:
                self.optimization_count += duplicates_removed
                node.elements = unique_elements

        return node

    def visit_for_expression(self, node: Any) -> Any:
        if hasattr(node, "iterable"):
            node.iterable = self.visit(node.iterable)
        if hasattr(node, "body"):
            node.body = self.visit(node.body)
        return node

    def visit_of_expression(self, node: Any) -> Any:
        if hasattr(node, "quantifier"):
            node.quantifier = self.visit(node.quantifier)
        if hasattr(node, "string_set"):
            node.string_set = self.visit(node.string_set)
        return node

    def visit_at_expression(self, node: Any) -> Any:
        if hasattr(node, "offset"):
            node.offset = self.visit(node.offset)
        return node

    def visit_in_expression(self, node: Any) -> Any:
        if hasattr(node, "range"):
            node.range = self.visit(node.range)
        return node

    # Rule-level methods (not used for expression optimization)
    def visit_yara_file(self, node: Any) -> Any:
        return node

    def visit_rule(self, node: Any) -> Any:
        return node

    def visit_import(self, node: Any) -> Any:
        return node

    def visit_include(self, node: Any) -> Any:
        return node

    def visit_tag(self, node: Any) -> Any:
        return node

    def visit_meta(self, node: Any) -> Any:
        return node

    def visit_plain_string(self, node: Any) -> Any:
        return node

    def visit_hex_string(self, node: Any) -> Any:
        return node

    def visit_regex_string(self, node: Any) -> Any:
        return node


def optimize_expression(expr: Expression) -> Expression:
    """Convenience function to optimize an expression."""
    optimizer = ExpressionOptimizer()
    return optimizer.optimize(expr)
