# Copyright 2025 the original author or authors.
# <p>
# Licensed under the Moderne Source Available License (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# <p>
# https://docs.moderne.io/licensing/moderne-source-available-license
# <p>
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Pattern matching comparator for comparing pattern and target ASTs."""

from __future__ import annotations

from typing import Dict, Optional, TYPE_CHECKING, cast

from rewrite.java import J
from rewrite.java import tree as j
from rewrite.python import tree as py
from .capture import Capture
from .placeholder import from_placeholder

if TYPE_CHECKING:
    from rewrite.visitor import Cursor


class PatternMatchingComparator:
    """
    Compares a pattern AST against a target AST, extracting captures.

    The comparator walks both trees in parallel, comparing node structures
    and types. When it encounters a placeholder identifier in the pattern,
    it captures the corresponding subtree from the target.

    Examples:
        comparator = PatternMatchingComparator(captures)
        result = comparator.match(pattern_tree, target_tree, cursor)
        if result:
            # result is a dict of capture_name -> captured_value
            x_value = result['x']
    """

    def __init__(
        self,
        captures: Dict[str, Capture],
        lenient_type_matching: bool = True
    ):
        """
        Initialize the comparator.

        Args:
            captures: Dict mapping capture names to Capture objects.
            lenient_type_matching: If True, ignore type differences when comparing.
        """
        self._captures = captures
        self._lenient_type_matching = lenient_type_matching
        self._captured: Dict[str, J] = {}
        self._debug = False

    def match(
        self,
        pattern: J,
        target: J,
        cursor: 'Cursor',
        *,
        debug: bool = False
    ) -> Optional[Dict[str, J]]:
        """
        Match pattern against target, returning captures if successful.

        Args:
            pattern: The pattern AST to match.
            target: The target AST to match against.
            cursor: Cursor at the target's position.
            debug: Enable debug logging.

        Returns:
            Dict of captured values if matched, None otherwise.
        """
        self._captured = {}
        self._debug = debug

        if self._compare(pattern, target, cursor):
            # Validate all required captures were matched
            for name, cap in self._captures.items():
                if name not in self._captured and not cap.variadic:
                    if self._debug:
                        print(f"Required capture '{name}' was not matched")
                    return None

            # Validate constraints
            for name, value in self._captured.items():
                cap = self._captures.get(name)
                if cap and cap.constraint:
                    if not cap.constraint(value):
                        if self._debug:
                            print(f"Capture '{name}' failed constraint")
                        return None

            return self._captured
        return None

    def _compare(
        self,
        pattern: Optional[J],
        target: Optional[J],
        cursor: 'Cursor'
    ) -> bool:
        """
        Recursively compare pattern and target nodes.

        Args:
            pattern: Pattern node (may contain placeholders).
            target: Target node to compare against.
            cursor: Current cursor position.

        Returns:
            True if nodes match.
        """
        # Both None is a match
        if pattern is None and target is None:
            return True

        # Only one None is not a match
        if pattern is None or target is None:
            return False

        # Check if pattern is a capture placeholder
        if isinstance(pattern, j.Identifier):
            capture_name = from_placeholder(pattern.simple_name)
            if capture_name is not None:
                return self._capture_node(capture_name, target)

        # Types must match (unless lenient)
        if type(pattern) != type(target):
            if self._debug:
                print(f"Type mismatch: {type(pattern).__name__} vs {type(target).__name__}")
            return False

        # Dispatch to type-specific comparison
        if isinstance(pattern, j.Identifier):
            return self._compare_identifier(pattern, cast(j.Identifier, target))
        elif isinstance(pattern, j.Literal):
            return self._compare_literal(pattern, cast(j.Literal, target))
        elif isinstance(pattern, j.Empty):
            # Two Empty sentinel nodes always match (used for absent values)
            return True
        elif isinstance(pattern, j.MethodInvocation):
            return self._compare_method_invocation(pattern, cast(j.MethodInvocation, target), cursor)
        elif isinstance(pattern, j.FieldAccess):
            return self._compare_field_access(pattern, cast(j.FieldAccess, target), cursor)
        elif isinstance(pattern, j.Binary):
            return self._compare_binary(pattern, cast(j.Binary, target), cursor)
        elif isinstance(pattern, j.Unary):
            return self._compare_unary(pattern, cast(j.Unary, target), cursor)
        elif isinstance(pattern, j.Assignment):
            return self._compare_assignment(pattern, cast(j.Assignment, target), cursor)
        elif isinstance(pattern, j.Parentheses):
            return self._compare_parentheses(pattern, cast(j.Parentheses, target), cursor)
        elif isinstance(pattern, j.Ternary):
            return self._compare_ternary(pattern, cast(j.Ternary, target), cursor)
        elif isinstance(pattern, j.Return):
            return self._compare_return(pattern, cast(j.Return, target), cursor)
        elif isinstance(pattern, j.ArrayAccess):
            return self._compare_array_access(pattern, cast(j.ArrayAccess, target), cursor)
        elif isinstance(pattern, py.ExpressionStatement):
            return self._compare_expression_statement(pattern, cast(py.ExpressionStatement, target), cursor)
        elif isinstance(pattern, py.Binary):
            return self._compare_python_binary(pattern, cast(py.Binary, target), cursor)
        elif isinstance(pattern, py.CollectionLiteral):
            return self._compare_collection_literal(pattern, cast(py.CollectionLiteral, target), cursor)
        elif isinstance(pattern, py.DictLiteral):
            return self._compare_dict_literal(pattern, cast(py.DictLiteral, target), cursor)
        else:
            # Default: unhandled node type — reject the match to prevent
            # false positives.  If a pattern uses a node type that reaches
            # this branch, a specific comparator method should be added.
            if self._debug:
                print(f"No specific comparison for {type(pattern).__name__}, rejecting match")
            return False

    def _capture_node(self, name: str, target: J) -> bool:
        """
        Capture a target node for a placeholder.

        Args:
            name: The capture name.
            target: The target node to capture.

        Returns:
            True if capture succeeded.
        """
        if name not in self._captures:
            if self._debug:
                print(f"Unknown capture: {name}")
            return False

        cap = self._captures[name]

        # Check if this capture already has a value
        if name in self._captured:
            # Must match the existing captured value
            existing = self._captured[name]
            # For now, require exact same node (by id)
            # TODO: implement semantic equality
            if self._debug:
                print(f"Capture '{name}' already has value, checking match")
            return existing.id == target.id

        # New capture
        self._captured[name] = target
        if self._debug:
            print(f"Captured '{name}': {type(target).__name__}")
        return True

    def _compare_identifier(self, pattern: j.Identifier, target: j.Identifier) -> bool:
        """Compare two identifiers."""
        return pattern.simple_name == target.simple_name

    def _compare_literal(self, pattern: j.Literal, target: j.Literal) -> bool:
        """Compare two literals.

        Uses a two-level comparison:
        1. Value types must match (prevents cross-type false positives like
           None vs b"" where both might serialize to the same representation).
        2. When both values are None (Ellipsis, None keyword, or literals with
           unicode escapes all store value=None), fall back to value_source
           comparison to distinguish them.
        """
        if type(pattern.value) != type(target.value):
            return False
        if pattern.value is None:
            return pattern.value_source == target.value_source
        return pattern.value == target.value

    def _compare_method_invocation(
        self,
        pattern: j.MethodInvocation,
        target: j.MethodInvocation,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two method invocations."""
        # Compare select (receiver)
        if not self._compare(pattern.select, target.select, cursor):
            return False

        # Compare method name
        if not self._compare(pattern.name, target.name, cursor):
            return False

        # Compare arguments
        return self._compare_arguments(pattern.padding.arguments, target.padding.arguments, cursor)

    def _compare_arguments(
        self,
        pattern_args: Optional[j.JContainer],
        target_args: Optional[j.JContainer],
        cursor: 'Cursor'
    ) -> bool:
        """Compare argument containers."""
        if pattern_args is None and target_args is None:
            return True
        if pattern_args is None or target_args is None:
            return False

        pattern_elements = pattern_args.elements
        target_elements = target_args.elements

        # Check for variadic capture
        if len(pattern_elements) == 1:
            pattern_arg = pattern_elements[0]
            if isinstance(pattern_arg, j.Identifier):
                cap_name = from_placeholder(pattern_arg.simple_name)
                if cap_name and self._captures.get(cap_name, Capture(name=cap_name)).variadic:
                    # Variadic capture - capture all target arguments
                    # TODO: implement proper variadic handling
                    return True

        # Non-variadic: must have same number of arguments
        if len(pattern_elements) != len(target_elements):
            if self._debug:
                print(f"Argument count mismatch: {len(pattern_elements)} vs {len(target_elements)}")
            return False

        # Compare each argument
        for p_elem, t_elem in zip(pattern_elements, target_elements):
            if not self._compare(p_elem, t_elem, cursor):
                return False

        return True

    def _compare_field_access(
        self,
        pattern: j.FieldAccess,
        target: j.FieldAccess,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two field accesses."""
        # Compare target (receiver)
        if not self._compare(pattern.target, target.target, cursor):
            return False

        # Compare name
        return self._compare(pattern.name, target.name, cursor)

    def _compare_binary(
        self,
        pattern: j.Binary,
        target: j.Binary,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two Java binary expressions."""
        # Compare operator
        if pattern.operator != target.operator:
            return False

        # Compare operands
        if not self._compare(pattern.left, target.left, cursor):
            return False

        return self._compare(pattern.right, target.right, cursor)

    def _compare_python_binary(
        self,
        pattern: py.Binary,
        target: py.Binary,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two Python binary expressions."""
        # Compare operator
        if pattern.operator != target.operator:
            return False

        # Compare operands
        if not self._compare(pattern.left, target.left, cursor):
            return False

        return self._compare(pattern.right, target.right, cursor)

    def _compare_unary(
        self,
        pattern: j.Unary,
        target: j.Unary,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two unary expressions."""
        if pattern.operator != target.operator:
            return False

        return self._compare(pattern.expression, target.expression, cursor)

    def _compare_assignment(
        self,
        pattern: j.Assignment,
        target: j.Assignment,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two assignments."""
        if not self._compare(pattern.variable, target.variable, cursor):
            return False

        return self._compare(pattern.assignment, target.assignment, cursor)

    def _compare_ternary(
        self,
        pattern: j.Ternary,
        target: j.Ternary,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two ternary (conditional) expressions."""
        if not self._compare(pattern.condition, target.condition, cursor):
            return False

        if not self._compare(pattern.true_part, target.true_part, cursor):
            return False

        return self._compare(pattern.false_part, target.false_part, cursor)

    def _compare_array_access(
        self,
        pattern: j.ArrayAccess,
        target: j.ArrayAccess,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two array/subscript accesses."""
        if not self._compare(pattern.indexed, target.indexed, cursor):
            return False

        return self._compare(pattern.dimension.index, target.dimension.index, cursor)

    def _compare_parentheses(
        self,
        pattern: j.Parentheses,
        target: j.Parentheses,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two parenthesized expressions."""
        return self._compare(pattern.tree, target.tree, cursor)

    def _compare_return(
        self,
        pattern: j.Return,
        target: j.Return,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two return statements."""
        return self._compare(pattern.expression, target.expression, cursor)

    def _compare_expression_statement(
        self,
        pattern: py.ExpressionStatement,
        target: py.ExpressionStatement,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two Python expression statements."""
        return self._compare(pattern.expression, target.expression, cursor)

    def _compare_collection_literal(
        self,
        pattern: py.CollectionLiteral,
        target: py.CollectionLiteral,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two Python collection literals."""
        # Compare kind (list, set, tuple)
        if pattern.kind != target.kind:
            return False

        # Compare elements
        pattern_elements = pattern.elements
        target_elements = target.elements

        if len(pattern_elements) != len(target_elements):
            return False

        for p_elem, t_elem in zip(pattern_elements, target_elements):
            if not self._compare(p_elem, t_elem, cursor):
                return False

        return True

    def _compare_dict_literal(
        self,
        pattern: py.DictLiteral,
        target: py.DictLiteral,
        cursor: 'Cursor'
    ) -> bool:
        """Compare two Python dict literals."""
        pattern_elements = pattern.elements
        target_elements = target.elements

        if len(pattern_elements) != len(target_elements):
            return False

        for p_elem, t_elem in zip(pattern_elements, target_elements):
            if not self._compare(p_elem, t_elem, cursor):
                return False

        return True
