"""Dead code elimination for YARA rules."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from yaraast.ast.expressions import BooleanLiteral, Identifier, StringIdentifier
from yaraast.visitor.base import ASTTransformer

if TYPE_CHECKING:
    from yaraast.ast.base import ASTNode, YaraFile
    from yaraast.ast.rules import Rule


class DeadCodeEliminator(ASTTransformer):
    """Eliminates dead code from YARA rules."""

    def __init__(self) -> None:
        super().__init__()
        self.used_strings: set[str] = set()
        self.used_rules: set[str] = set()
        self.in_condition = False
        self.current_rule: str | None = None
        self.elimination_count = 0

    def eliminate(self, ast: YaraFile) -> tuple[YaraFile, int]:
        """Eliminate dead code from AST.

        Returns:
            Tuple of (optimized YaraFile, number of eliminations)
        """
        # Reset state
        self.used_strings.clear()
        self.used_rules.clear()
        self.in_condition = False
        self.current_rule = None
        self.elimination_count = 0

        # First pass: collect used strings and rules
        self._collect_usage(ast)

        # Count strings to be removed
        for rule in ast.rules:
            if rule.strings:
                for string_def in rule.strings:
                    if string_def.identifier not in self.used_strings:
                        self.elimination_count += 1

        # Count rules with always-false conditions
        for rule in ast.rules:
            if (
                rule.condition
                and isinstance(rule.condition, BooleanLiteral)
                and not rule.condition.value
            ):  # false condition
                self.elimination_count += 1

        # Second pass: eliminate unused code
        optimized_ast = self.visit(ast)

        return optimized_ast, self.elimination_count

    def _collect_usage(self, ast: YaraFile) -> None:
        """Collect usage information."""
        for rule in ast.rules:
            self.current_rule = rule.name
            self.in_condition = True
            if rule.condition:
                self._collect_from_expression(rule.condition)
            self.in_condition = False

    def _collect_from_expression(self, expr: ASTNode) -> None:
        """Collect usage from expression."""
        if isinstance(expr, StringIdentifier):
            self.used_strings.add(expr.name)
        elif isinstance(expr, Identifier):
            # Could be a rule reference
            self.used_rules.add(expr.name)

        # Recursively collect from children
        for child in expr.children():
            self._collect_from_expression(child)

    def visit_yara_file(self, node: YaraFile) -> YaraFile:
        """Visit YaraFile and remove unused rules."""
        kept_rules = []

        for rule in node.rules:
            # Skip rules with always-false conditions
            if (
                rule.condition
                and isinstance(rule.condition, BooleanLiteral)
                and not rule.condition.value
            ):  # false condition
                continue  # Remove this rule

            # Keep only used rules (or all if we can't determine)
            if self.used_rules:
                # Determine if the rule is private (internal helper)
                is_private = False
                if hasattr(rule, "modifiers") and isinstance(
                    rule.modifiers,
                    list | tuple,
                ):
                    is_private = any(
                        getattr(m, "modifier_type", None) and m.modifier_type.value == "private"
                        for m in rule.modifiers
                    )
                # Remove only private rules that nobody references
                if not is_private or rule.name in self.used_rules:
                    kept_rules.append(self.visit(rule))
            else:
                # If no usage info, keep all rules but optimize them
                kept_rules.append(self.visit(rule))

        node.rules = kept_rules
        return node

    def _is_referenced_by_other_rules(self, rule_name: str) -> bool:
        """Check if this rule is referenced by any other rule."""
        return rule_name in self.used_rules

    def _has_external_references(self, rule: Rule) -> bool:
        """Check if rule references other rules."""
        if rule.condition:
            return self._contains_rule_reference(rule.condition)
        return False

    def _contains_rule_reference(self, expr: ASTNode) -> bool:
        """Check if expression contains rule reference."""
        if isinstance(expr, Identifier) and expr.name not in [
            "true",
            "false",
            "any",
            "all",
            "them",
        ]:
            # Could be a rule reference
            return True

        return any(self._contains_rule_reference(child) for child in expr.children())

    def visit_rule(self, node: Rule) -> Rule:
        """Visit Rule and remove unused strings."""
        self.current_rule = node.name

        # Remove unused strings
        if node.strings and self.used_strings:
            kept_strings = []
            for string_def in node.strings:
                if string_def.identifier in self.used_strings:
                    kept_strings.append(string_def)
            node.strings = kept_strings

        # Optimize condition
        if node.condition:
            self.in_condition = True
            node.condition = self.visit(node.condition)
            self.in_condition = False

        return node

    def visit_boolean_literal(self, node: BooleanLiteral) -> BooleanLiteral:
        """Visit BooleanLiteral - detect always true/false conditions."""
        # In a real implementation, we might eliminate rules with always-false conditions
        return node

    def visit_string_identifier(self, node: StringIdentifier) -> StringIdentifier:
        """Visit StringIdentifier - track usage."""
        if self.in_condition:
            self.used_strings.add(node.name)
        return node

    def visit_string_wildcard(self, node: StringIdentifier) -> StringIdentifier:
        """Visit StringWildcard node."""
        pass

    def visit_identifier(self, node: Identifier) -> Identifier:
        """Visit Identifier - track potential rule usage."""
        if self.in_condition and node.name not in [
            "true",
            "false",
            "any",
            "all",
            "them",
        ]:
            self.used_rules.add(node.name)
        return node

    # Pass-through methods for other node types
    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 visit_string_count(self, node: Any) -> Any:
        if self.in_condition and hasattr(node, "string_id"):
            self.used_strings.add(node.string_id)
        return node

    def visit_string_offset(self, node: Any) -> Any:
        if self.in_condition and hasattr(node, "string_id"):
            self.used_strings.add(node.string_id)
        return node

    def visit_string_length(self, node: Any) -> Any:
        if self.in_condition and hasattr(node, "string_id"):
            self.used_strings.add(node.string_id)
        return node

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

        # Constant folding for boolean literals
        if isinstance(node.left, BooleanLiteral) and isinstance(
            node.right,
            BooleanLiteral,
        ):
            if node.operator == "and":
                return BooleanLiteral(value=node.left.value and node.right.value)
            if node.operator == "or":
                return BooleanLiteral(value=node.left.value or node.right.value)

        # Simplifications
        if isinstance(node.left, BooleanLiteral):
            if node.operator == "and" and not node.left.value:
                return BooleanLiteral(value=False)
            if node.operator == "or" and node.left.value:
                return BooleanLiteral(value=True)

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

        return node

    def visit_unary_expression(self, node: Any) -> Any:
        """Visit UnaryExpression and optimize if possible."""
        node.operand = self.visit(node.operand)

        # Optimize not on boolean literal
        if node.operator == "not" and isinstance(node.operand, BooleanLiteral):
            return BooleanLiteral(value=not node.operand.value)

        return node

    def eliminate_dead_code(self, rule: Rule) -> Rule:
        """Eliminate dead code from a single rule.

        This is a simplified version that just removes unused strings.
        """
        # First collect used strings from this rule
        self.used_strings.clear()
        self.current_rule = rule.name
        self.in_condition = True

        if rule.condition:
            self._collect_from_expression(rule.condition)

        # Remove unused strings
        if rule.strings and self.used_strings:
            kept_strings = []
            for string_def in rule.strings:
                if string_def.identifier in self.used_strings:
                    kept_strings.append(string_def)
            rule.strings = kept_strings

        return rule


def eliminate_dead_code(ast: YaraFile) -> YaraFile:
    """Convenience function to eliminate dead code."""
    eliminator = DeadCodeEliminator()
    return eliminator.eliminate(ast)
