"""Fix synthesis using program synthesis and fix templates.

This module provides template-based fix generation for common flakiness patterns.
It uses AST transformations to apply fixes safely and correctly.

Key features:
- Built-in templates for common flakiness patterns
- AST-based code transformations
- Pattern matching for error-to-template mapping
- Confidence scoring for fix quality
"""

from __future__ import annotations

import ast
import re
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

from autodebug.llm_first.models import FailureContext, FlakinessCategory
try:
    import sqlparse
except Exception:
    sqlparse = None


def _insert_order_by_clause(sql: str, order_expr: str) -> str:
    """Insert ORDER BY consistently while preserving trailing semicolons."""
    has_semicolon = sql.rstrip().endswith(";")
    base_sql = sql.rstrip()
    if has_semicolon:
        base_sql = base_sql[:-1].rstrip()

    limit_match = re.search(r'\s+(LIMIT|OFFSET)\s+', base_sql, re.IGNORECASE)
    if limit_match:
        insert_pos = limit_match.start()
        updated = f"{base_sql[:insert_pos]} ORDER BY {order_expr}{base_sql[insert_pos:]}"
    else:
        updated = f"{base_sql} ORDER BY {order_expr}"

    return f"{updated};" if has_semicolon else updated

# =============================================================================
# Data Models
# =============================================================================


@dataclass
class FixTemplate:
    """Template for generating a fix.

    Attributes:
        name: Unique identifier for the template
        description: Human-readable description of what the fix does
        pattern: Regex or AST pattern to detect the issue
        transformation: Description of the fix transformation
        applicable_errors: List of error types/messages this template applies to
        category: The flakiness category this template addresses
        confidence: Confidence score (0-1) that this fix will work
    """

    name: str
    description: str
    pattern: str
    transformation: str
    applicable_errors: List[str]
    category: FlakinessCategory
    confidence: float = 0.8

    # Optional fields for more advanced matching
    ast_pattern: Optional[Dict[str, Any]] = None
    required_imports: List[str] = field(default_factory=list)
    context_variables: List[str] = field(default_factory=list)

    def matches_error(self, error_type: str, error_message: str) -> bool:
        """Check if this template applies to the given error.

        Args:
            error_type: The type of error (e.g., "AssertionError")
            error_message: The error message

        Returns:
            True if this template might fix the error
        """
        error_lower = error_type.lower()
        message_lower = error_message.lower()

        for applicable in self.applicable_errors:
            applicable_lower = applicable.lower()
            if applicable_lower in error_lower or applicable_lower in message_lower:
                return True
        return False


class TransformationType(str, Enum):
    """Types of AST transformations."""

    ADD_STATEMENT = "add_statement"
    WRAP_EXPRESSION = "wrap_expression"
    REPLACE_CALL = "replace_call"
    ADD_DECORATOR = "add_decorator"
    ADD_IMPORT = "add_import"
    MODIFY_ARGUMENT = "modify_argument"


@dataclass
class TransformationResult:
    """Result of applying a transformation."""

    success: bool
    transformed_code: str
    original_code: str
    changes_made: List[str] = field(default_factory=list)
    error: Optional[str] = None


# =============================================================================
# AST Transformers
# =============================================================================


class OrderByTransformer(ast.NodeTransformer):
    """Add ORDER BY to SQL query strings."""

    def __init__(self):
        self.modified = False
        self.changes: List[str] = []

    def visit_Call(self, node: ast.Call) -> ast.Call:
        """Visit function calls to find SQL execution methods."""
        self.generic_visit(node)

        # Check for common SQL execution patterns
        if isinstance(node.func, ast.Attribute):
            method_name = node.func.attr
            if method_name in ('execute', 'executemany', 'query', 'raw', 'text'):
                # Check if first argument is a string (SQL query)
                if node.args and isinstance(node.args[0], ast.Constant):
                    sql = node.args[0].value
                    if isinstance(sql, str):
                        new_sql = self._add_order_by(sql)
                        if new_sql != sql:
                            node.args[0] = ast.Constant(value=new_sql)
                            self.modified = True
                            self.changes.append("Added ORDER BY to query")

        return node

    def _add_order_by(self, sql: str) -> str:
        """Add ORDER BY clause if missing from SELECT query."""
        sql_upper = sql.upper().strip()

        # Skip if not a SELECT or already has ORDER BY
        if not sql_upper.startswith('SELECT'):
            return sql
        if 'ORDER BY' in sql_upper:
            return sql

        # Prefer token-aware handling when sqlparse is available.
        if sqlparse is not None:
            try:
                parsed = sqlparse.parse(sql)
                if parsed:
                    normalized = parsed[0].normalized
                    if "ORDER BY" in normalized:
                        return sql
            except Exception:
                pass

        # Find a suitable column to order by using lightweight heuristics.
        match = re.search(r'SELECT\s+(.+?)\s+FROM', sql, re.IGNORECASE | re.DOTALL)
        if match:
            columns = match.group(1)
            order_column = None

            for col in ['id', 'created_at', 'updated_at', 'timestamp', 'name']:
                if col in columns.lower() or '*' in columns:
                    order_column = col
                    break

            if order_column:
                return _insert_order_by_clause(sql, order_column)

        return sql


class AwaitTransformer(ast.NodeTransformer):
    """Add await to async calls that are missing it."""

    def __init__(self, async_functions: Optional[Set[str]] = None):
        self.modified = False
        self.changes: List[str] = []
        self.async_functions = async_functions or {
            'fetch', 'get', 'post', 'put', 'delete', 'patch',
            'sleep', 'gather', 'wait', 'wait_for',
            'read', 'write', 'connect', 'close',
            'execute', 'fetchone', 'fetchall', 'fetchmany',
            'create', 'save', 'update', 'find', 'find_one',
        }
        self.in_async_context = False

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef:
        """Track when we're inside an async function."""
        old_context = self.in_async_context
        self.in_async_context = True
        self.generic_visit(node)
        self.in_async_context = old_context
        return node

    def visit_Expr(self, node: ast.Expr) -> Union[ast.Expr, ast.Await]:
        """Check expression statements for unawaited async calls."""
        if not self.in_async_context:
            return node

        # Check if this is a call that should be awaited
        if isinstance(node.value, ast.Call):
            call = node.value
            func_name = self._get_call_name(call)

            if func_name and func_name in self.async_functions:
                # Wrap in await
                new_node = ast.Expr(
                    value=ast.Await(value=call)
                )
                ast.copy_location(new_node, node)
                self.modified = True
                self.changes.append(f"Added await to {func_name}()")
                return new_node

        return node

    def _get_call_name(self, call: ast.Call) -> Optional[str]:
        """Extract function name from a Call node."""
        if isinstance(call.func, ast.Name):
            return call.func.id
        elif isinstance(call.func, ast.Attribute):
            return call.func.attr
        return None


class RetryTransformer(ast.NodeTransformer):
    """Add retry decorator to functions with network calls."""

    def __init__(self, retry_on: Optional[List[str]] = None):
        self.modified = False
        self.changes: List[str] = []
        self.retry_on = retry_on or ['requests', 'httpx', 'aiohttp', 'urllib']

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
        """Add retry decorator to functions with network calls."""
        self.generic_visit(node)

        # Check if function already has retry decorator
        for decorator in node.decorator_list:
            if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Attribute):
                if decorator.func.attr in ('retry', 'retrying'):
                    return node
            elif isinstance(decorator, ast.Name) and decorator.id == 'retry':
                return node

        # Check if function contains network calls
        if self._has_network_calls(node):
            # Add retry decorator
            retry_decorator = ast.Call(
                func=ast.Attribute(
                    value=ast.Name(id='tenacity', ctx=ast.Load()),
                    attr='retry',
                    ctx=ast.Load()
                ),
                args=[],
                keywords=[
                    ast.keyword(arg='stop', value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id='tenacity', ctx=ast.Load()),
                            attr='stop_after_attempt',
                            ctx=ast.Load()
                        ),
                        args=[ast.Constant(value=3)],
                        keywords=[]
                    )),
                    ast.keyword(arg='wait', value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id='tenacity', ctx=ast.Load()),
                            attr='wait_exponential',
                            ctx=ast.Load()
                        ),
                        args=[],
                        keywords=[
                            ast.keyword(arg='multiplier', value=ast.Constant(value=1)),
                            ast.keyword(arg='max', value=ast.Constant(value=10))
                        ]
                    ))
                ]
            )
            node.decorator_list.insert(0, retry_decorator)
            self.modified = True
            self.changes.append(f"Added retry decorator to {node.name}()")

        return node

    def _has_network_calls(self, node: ast.FunctionDef) -> bool:
        """Check if function contains network calls."""
        for child in ast.walk(node):
            if isinstance(child, ast.Call):
                if isinstance(child.func, ast.Attribute):
                    if isinstance(child.func.value, ast.Name):
                        if child.func.value.id in self.retry_on:
                            return True
        return False


class SortedTransformer(ast.NodeTransformer):
    """Add sorted() to collection comparisons."""

    def __init__(self):
        self.modified = False
        self.changes: List[str] = []

    def visit_Compare(self, node: ast.Compare) -> ast.Compare:
        """Wrap list comparisons with sorted()."""
        self.generic_visit(node)

        # Check for equality comparisons
        if len(node.ops) == 1 and isinstance(node.ops[0], ast.Eq):
            left = node.left
            right = node.comparators[0]

            # Check if comparing lists/sets
            left_is_collection = self._is_collection(left)
            right_is_collection = self._is_collection(right)

            if left_is_collection or right_is_collection:
                # Wrap both sides in sorted() if they're not already
                if not self._is_sorted_call(left):
                    node.left = self._wrap_sorted(left)
                    self.modified = True

                if not self._is_sorted_call(right):
                    node.comparators[0] = self._wrap_sorted(right)
                    self.modified = True

                if self.modified:
                    self.changes.append("Wrapped collection comparison with sorted()")

        return node

    def _is_collection(self, node: ast.AST) -> bool:
        """Check if node represents a collection."""
        if isinstance(node, (ast.List, ast.Set)):
            return True
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name):
                if node.func.id in ('list', 'set', 'dict', 'keys', 'values', 'items'):
                    return True
            elif isinstance(node.func, ast.Attribute):
                if node.func.attr in ('keys', 'values', 'items', 'all', 'filter'):
                    return True
        return False

    def _is_sorted_call(self, node: ast.AST) -> bool:
        """Check if node is already a sorted() call."""
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name) and node.func.id == 'sorted':
                return True
        return False

    def _wrap_sorted(self, node: ast.AST) -> ast.Call:
        """Wrap a node in sorted()."""
        return ast.Call(
            func=ast.Name(id='sorted', ctx=ast.Load()),
            args=[node],
            keywords=[]
        )


class SleepTransformer(ast.NodeTransformer):
    """Add time.sleep for race condition mitigation."""

    def __init__(self, before_pattern: Optional[str] = None, sleep_duration: float = 0.1):
        self.modified = False
        self.changes: List[str] = []
        self.before_pattern = before_pattern or r'assert|expect|verify'
        self.sleep_duration = sleep_duration

    def visit_Expr(self, node: ast.Expr) -> Union[ast.Expr, List[ast.stmt]]:
        """Add sleep before assertions."""
        # Check if this is an assertion-like statement
        if isinstance(node.value, ast.Call):
            func_name = self._get_call_name(node.value)
            if func_name and re.match(self.before_pattern, func_name, re.IGNORECASE):
                # Create sleep statement
                sleep_stmt = ast.Expr(
                    value=ast.Call(
                        func=ast.Attribute(
                            value=ast.Name(id='time', ctx=ast.Load()),
                            attr='sleep',
                            ctx=ast.Load()
                        ),
                        args=[ast.Constant(value=self.sleep_duration)],
                        keywords=[]
                    )
                )
                ast.copy_location(sleep_stmt, node)
                self.modified = True
                self.changes.append(f"Added time.sleep({self.sleep_duration}) before {func_name}")
                return [sleep_stmt, node]

        return node

    def _get_call_name(self, call: ast.Call) -> Optional[str]:
        """Extract function name from a Call node."""
        if isinstance(call.func, ast.Name):
            return call.func.id
        elif isinstance(call.func, ast.Attribute):
            return call.func.attr
        return None


class ContextManagerTransformer(ast.NodeTransformer):
    """Add context manager (with statement) for resource cleanup."""

    def __init__(self, resource_patterns: Optional[Dict[str, str]] = None):
        self.modified = False
        self.changes: List[str] = []
        self.resource_patterns = resource_patterns or {
            'open': 'file',
            'connect': 'connection',
            'create_connection': 'connection',
            'get_connection': 'connection',
            'Session': 'session',
        }

    def visit_Assign(self, node: ast.Assign) -> Union[ast.Assign, ast.With]:
        """Convert resource assignments to with statements."""
        # Check if this is a resource creation
        if isinstance(node.value, ast.Call):
            func_name = self._get_call_name(node.value)
            if func_name in self.resource_patterns:
                # Get the variable name
                if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
                    var_name = node.targets[0].id
                    resource_type = self.resource_patterns[func_name]

                    # Create with statement
                    with_stmt = ast.With(
                        items=[
                            ast.withitem(
                                context_expr=node.value,
                                optional_vars=ast.Name(id=var_name, ctx=ast.Store())
                            )
                        ],
                        body=[ast.Pass()]  # Placeholder - caller should fill in
                    )
                    ast.copy_location(with_stmt, node)
                    self.modified = True
                    self.changes.append(f"Converted {func_name}() to context manager")
                    return with_stmt

        return node

    def _get_call_name(self, call: ast.Call) -> Optional[str]:
        """Extract function name from a Call node."""
        if isinstance(call.func, ast.Name):
            return call.func.id
        elif isinstance(call.func, ast.Attribute):
            return call.func.attr
        return None


class MockTransformer(ast.NodeTransformer):
    """Add mock decorators for external dependencies."""

    def __init__(self, mock_targets: Optional[List[str]] = None):
        self.modified = False
        self.changes: List[str] = []
        self.mock_targets = mock_targets or [
            'requests.get', 'requests.post', 'requests.put', 'requests.delete',
            'httpx.get', 'httpx.post', 'httpx.AsyncClient',
            'aiohttp.ClientSession',
            'datetime.datetime.now', 'datetime.datetime.utcnow',
            'time.time', 'time.sleep',
        ]

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
        """Add mock decorators to test functions."""
        self.generic_visit(node)

        # Only process test functions
        if not node.name.startswith('test_'):
            return node

        # Check for external calls that should be mocked
        targets_found = self._find_mock_targets(node)

        for target in targets_found:
            # Check if already mocked
            if self._has_mock_decorator(node, target):
                continue

            # Add mock.patch decorator
            mock_decorator = ast.Call(
                func=ast.Attribute(
                    value=ast.Attribute(
                        value=ast.Name(id='unittest', ctx=ast.Load()),
                        attr='mock',
                        ctx=ast.Load()
                    ),
                    attr='patch',
                    ctx=ast.Load()
                ),
                args=[ast.Constant(value=target)],
                keywords=[]
            )
            node.decorator_list.append(mock_decorator)
            self.modified = True
            self.changes.append(f"Added mock.patch for {target}")

        return node

    def _find_mock_targets(self, node: ast.FunctionDef) -> List[str]:
        """Find external calls that should be mocked."""
        found = []
        for child in ast.walk(node):
            if isinstance(child, ast.Call):
                call_name = self._get_full_call_name(child)
                if call_name in self.mock_targets:
                    found.append(call_name)
        return list(set(found))

    def _has_mock_decorator(self, node: ast.FunctionDef, target: str) -> bool:
        """Check if function already has mock for target."""
        for decorator in node.decorator_list:
            if isinstance(decorator, ast.Call):
                if decorator.args and isinstance(decorator.args[0], ast.Constant):
                    if decorator.args[0].value == target:
                        return True
        return False

    def _get_full_call_name(self, call: ast.Call) -> str:
        """Get full dotted name of a call."""
        parts = []
        node = call.func
        while isinstance(node, ast.Attribute):
            parts.append(node.attr)
            node = node.value
        if isinstance(node, ast.Name):
            parts.append(node.id)
        return '.'.join(reversed(parts))


class RandomSeedTransformer(ast.NodeTransformer):
    """Add random seed for reproducibility."""

    def __init__(self, seed: int = 42):
        self.modified = False
        self.changes: List[str] = []
        self.seed = seed
        self.needs_seed = False
        self.random_modules = {'random', 'np', 'numpy', 'torch', 'tf', 'tensorflow'}

    def visit_Module(self, node: ast.Module) -> ast.Module:
        """Add seed at module level if random is used."""
        self.generic_visit(node)

        if self.needs_seed:
            # Find first import or first statement
            insert_pos = 0
            for i, stmt in enumerate(node.body):
                if isinstance(stmt, (ast.Import, ast.ImportFrom)):
                    insert_pos = i + 1
                else:
                    break

            # Create seed statements
            seed_stmts = self._create_seed_statements()
            for stmt in reversed(seed_stmts):
                node.body.insert(insert_pos, stmt)

            self.modified = True
            self.changes.append(f"Added random seed ({self.seed}) at module level")

        return node

    def visit_Call(self, node: ast.Call) -> ast.Call:
        """Check for random usage."""
        self.generic_visit(node)

        if isinstance(node.func, ast.Attribute):
            if isinstance(node.func.value, ast.Name):
                module = node.func.value.id
                method = node.func.attr

                # Check for random methods
                if module in self.random_modules:
                    if method in ('random', 'randint', 'choice', 'shuffle', 'sample',
                                  'uniform', 'gauss', 'normal', 'rand', 'randn'):
                        self.needs_seed = True

        return node

    def _create_seed_statements(self) -> List[ast.stmt]:
        """Create seed setting statements."""
        stmts = []

        # random.seed()
        stmts.append(ast.Expr(
            value=ast.Call(
                func=ast.Attribute(
                    value=ast.Name(id='random', ctx=ast.Load()),
                    attr='seed',
                    ctx=ast.Load()
                ),
                args=[ast.Constant(value=self.seed)],
                keywords=[]
            )
        ))

        return stmts


# =============================================================================
# Built-in Fix Templates
# =============================================================================


BUILTIN_TEMPLATES: List[FixTemplate] = [
    FixTemplate(
        name="add_order_by",
        description="Add ORDER BY clause to SQL queries that lack deterministic ordering",
        pattern=r"SELECT\s+.+?\s+FROM\s+.+?(?!.*ORDER\s+BY)",
        transformation="Add ORDER BY <column> to ensure deterministic result ordering",
        applicable_errors=[
            "AssertionError", "assertion failed", "order", "sequence",
            "list index", "index out of range", "not equal",
        ],
        category=FlakinessCategory.ORDERING,
        confidence=0.85,
        required_imports=[],
    ),
    FixTemplate(
        name="add_await",
        description="Add missing await keyword to async function calls",
        pattern=r"(?<!await\s)(fetch|get|post|put|delete|read|write|connect)\s*\(",
        transformation="Add await keyword before async call",
        applicable_errors=[
            "coroutine", "was never awaited", "RuntimeWarning",
            "asyncio", "async", "awaitable",
        ],
        category=FlakinessCategory.TIMING,
        confidence=0.9,
        required_imports=['asyncio'],
    ),
    FixTemplate(
        name="add_retry",
        description="Add retry decorator to handle transient failures",
        pattern=r"(requests|httpx|aiohttp)\.(get|post|put|delete|patch)\s*\(",
        transformation="Add @tenacity.retry decorator with exponential backoff",
        applicable_errors=[
            "ConnectionError", "TimeoutError", "HTTPError",
            "connection refused", "timeout", "network",
            "503", "502", "504", "429",
        ],
        category=FlakinessCategory.NETWORK,
        confidence=0.8,
        required_imports=['tenacity'],
    ),
    FixTemplate(
        name="add_sorted",
        description="Wrap collection comparisons with sorted() for order-independent comparison",
        pattern=r"(assert|expect|verify).*==.*(\[|\{|list|set|dict\.keys|\.values|\.items)",
        transformation="Wrap both sides of comparison with sorted()",
        applicable_errors=[
            "AssertionError", "not equal", "!=", "order",
            "list differs", "sets are not equal",
        ],
        category=FlakinessCategory.ORDERING,
        confidence=0.9,
    ),
    FixTemplate(
        name="add_sleep",
        description="Add small delay to allow async operations to complete",
        pattern=r"(assert|expect|verify)\s*\(",
        transformation="Add time.sleep(0.1) before assertion",
        applicable_errors=[
            "AssertionError", "not yet", "pending", "race condition",
            "timing", "intermittent", "flaky",
        ],
        category=FlakinessCategory.TIMING,
        confidence=0.6,
        required_imports=['time'],
        context_variables=['sleep_duration'],
    ),
    FixTemplate(
        name="add_context_manager",
        description="Use context manager (with statement) for proper resource cleanup",
        pattern=r"(\w+)\s*=\s*(open|connect|Session|get_connection)\s*\(",
        transformation="Convert to 'with ... as var:' pattern",
        applicable_errors=[
            "ResourceWarning", "unclosed", "leaked",
            "connection pool", "too many open files",
            "file descriptor", "resource exhausted",
        ],
        category=FlakinessCategory.RESOURCE,
        confidence=0.85,
    ),
    FixTemplate(
        name="add_mock",
        description="Add mock decorator for external dependencies",
        pattern=r"(requests|httpx|datetime\.datetime\.(now|utcnow)|time\.time)\s*[.\(]",
        transformation="Add @unittest.mock.patch decorator",
        applicable_errors=[
            "network", "connection", "external service",
            "time", "date", "timezone",
        ],
        category=FlakinessCategory.NETWORK,
        confidence=0.75,
        required_imports=['unittest.mock'],
    ),
    FixTemplate(
        name="fix_random_seed",
        description="Add random seed for reproducible random behavior",
        pattern=r"random\.(random|randint|choice|shuffle|sample|uniform|gauss)",
        transformation="Add random.seed(42) at module/test setup",
        applicable_errors=[
            "random", "non-deterministic", "flaky",
            "different results", "inconsistent",
        ],
        category=FlakinessCategory.STATE,
        confidence=0.9,
        required_imports=['random'],
    ),
]


# =============================================================================
# Main Synthesizer Class
# =============================================================================


class FixSynthesizer:
    """Synthesize fixes using templates and AST transformations.

    This class provides template-based fix generation for common flakiness
    patterns. It uses AST transformations for safe code modifications.

    Example usage:
        synthesizer = FixSynthesizer()

        # Get templates matching an error
        templates = synthesizer.match_error_to_template("AssertionError", "order mismatch")

        # Apply a template
        if templates:
            fixed_code = synthesizer.apply_template(templates[0], source_code)
    """

    def __init__(
        self,
        templates: Optional[List[FixTemplate]] = None,
        custom_transformers: Optional[Dict[str, Callable]] = None,
    ):
        """Initialize the synthesizer.

        Args:
            templates: List of fix templates (defaults to built-in templates)
            custom_transformers: Additional AST transformers keyed by template name
        """
        self.templates = templates or BUILTIN_TEMPLATES.copy()
        self._template_index: Dict[str, FixTemplate] = {t.name: t for t in self.templates}
        self._last_used_regex_fallback: bool = False

        # Register transformers
        self._transformers: Dict[str, Callable[..., ast.NodeTransformer]] = {
            "add_order_by": lambda: OrderByTransformer(),
            "add_await": lambda: AwaitTransformer(),
            "add_retry": lambda: RetryTransformer(),
            "add_sorted": lambda: SortedTransformer(),
            "add_sleep": lambda: SleepTransformer(),
            "add_context_manager": lambda: ContextManagerTransformer(),
            "add_mock": lambda: MockTransformer(),
            "fix_random_seed": lambda: RandomSeedTransformer(),
        }

        if custom_transformers:
            self._transformers.update(custom_transformers)

    def synthesize_from_template(
        self,
        template: FixTemplate,
        context: FailureContext,
    ) -> Optional[str]:
        """Generate a fix using a template and failure context.

        Args:
            template: The fix template to apply
            context: Context about the failure

        Returns:
            Fixed code string, or None if synthesis fails
        """
        if not context.code_snippet and not context.surrounding_code:
            return None

        source_code = context.surrounding_code or context.code_snippet
        if not source_code:
            return None

        try:
            return self.apply_template(template, source_code)
        except Exception:
            return None

    def available_templates(self) -> List[FixTemplate]:
        """Get list of all available fix templates.

        Returns:
            List of FixTemplate objects
        """
        return self.templates.copy()

    def match_error_to_template(
        self,
        error_type: str,
        error_message: str,
    ) -> List[FixTemplate]:
        """Find templates that might fix the given error.

        Args:
            error_type: The type of error (e.g., "AssertionError")
            error_message: The error message

        Returns:
            List of matching templates, sorted by confidence
        """
        matches = []

        for template in self.templates:
            if template.matches_error(error_type, error_message):
                matches.append(template)

        # Sort by confidence descending
        matches.sort(key=lambda t: t.confidence, reverse=True)
        return matches

    def apply_template(
        self,
        template: FixTemplate,
        source_code: str,
    ) -> str:
        """Apply a fix template to source code.

        Args:
            template: The fix template to apply
            source_code: The source code to transform

        Returns:
            Transformed source code

        Raises:
            ValueError: If template is unknown or transformation fails
        """
        if template.name not in self._transformers:
            # Fall back to regex-based transformation
            self._last_used_regex_fallback = True
            return self._apply_regex_transformation(template, source_code)

        try:
            # Parse the source code
            tree = ast.parse(source_code)

            # Create and apply transformer
            transformer_factory = self._transformers[template.name]
            transformer = transformer_factory()
            new_tree = transformer.visit(tree)

            # Fix missing locations
            ast.fix_missing_locations(new_tree)

            # Convert back to source
            self._last_used_regex_fallback = False
            return ast.unparse(new_tree)

        except SyntaxError:
            # If AST parsing fails, try regex transformation
            self._last_used_regex_fallback = True
            return self._apply_regex_transformation(template, source_code)

    def adjusted_confidence(self, base_confidence: float) -> float:
        """Reduce confidence when regex fallback was used."""
        if self._last_used_regex_fallback:
            return max(0.0, min(1.0, base_confidence * 0.75))
        return max(0.0, min(1.0, base_confidence))

    def _apply_regex_transformation(
        self,
        template: FixTemplate,
        source_code: str,
    ) -> str:
        """Apply transformation using regex pattern matching.

        Args:
            template: The fix template
            source_code: Source code to transform

        Returns:
            Transformed source code
        """
        if template.name == "add_order_by":
            return self._add_order_by_regex(source_code)
        elif template.name == "add_await":
            return self._add_await_regex(source_code)
        elif template.name == "add_sorted":
            return self._add_sorted_regex(source_code)
        elif template.name == "fix_random_seed":
            return self._add_random_seed_regex(source_code)
        else:
            # Return unchanged if no regex fallback
            return source_code

    def _add_order_by_regex(self, code: str) -> str:
        """Add ORDER BY using regex."""
        def add_order_by(match: re.Match) -> str:
            sql = match.group(0)
            if 'ORDER BY' not in sql.upper():
                # Add ORDER BY before LIMIT/OFFSET or at end
                if re.search(r'\bLIMIT\b', sql, re.IGNORECASE):
                    updated = re.sub(
                        r'(\s+)(LIMIT\b)',
                        r' ORDER BY id\1\2',
                        sql,
                        flags=re.IGNORECASE
                    )
                    if sql.rstrip().endswith(";") and not updated.rstrip().endswith(";"):
                        return f"{updated.rstrip()};"
                    return updated
                else:
                    return _insert_order_by_clause(sql, "id")
            return sql

        return re.sub(
            r"SELECT\s+.+?\s+FROM\s+\w+(?:\s+WHERE\s+.+?)?(?:\s+LIMIT\s+\d+)?",
            add_order_by,
            code,
            flags=re.IGNORECASE | re.DOTALL
        )

    def _add_await_regex(self, code: str) -> str:
        """Add await keyword using regex."""
        async_funcs = r'(fetch|get|post|put|delete|read|write|connect)'
        return re.sub(
            rf'(?<!await\s)({async_funcs}\s*\()',
            r'await \1',
            code
        )

    def _add_sorted_regex(self, code: str) -> str:
        """Add sorted() wrapper using regex."""
        # Match assertions with list/set comparisons
        def wrap_comparison(match: re.Match) -> str:
            full = match.group(0)
            if 'sorted(' not in full:
                # Simple heuristic: wrap each side of == with sorted()
                return re.sub(
                    r'(\w+)\s*==\s*(\w+)',
                    r'sorted(\1) == sorted(\2)',
                    full
                )
            return full

        return re.sub(
            r'assert\s+.+?\s*==\s*.+?(?=\n|$)',
            wrap_comparison,
            code
        )

    def _add_random_seed_regex(self, code: str) -> str:
        """Add random seed using regex."""
        # Check if seed is already set
        if 'random.seed(' in code:
            return code

        # Check if random is used
        if not re.search(r'random\.(random|randint|choice|shuffle|sample)', code):
            return code

        # Add seed after imports
        import_match = re.search(r'^import\s+random\s*$', code, re.MULTILINE)
        if import_match:
            insert_pos = import_match.end()
            return f"{code[:insert_pos]}\nrandom.seed(42)  # For reproducibility{code[insert_pos:]}"

        # Add import and seed at beginning
        return f"import random\nrandom.seed(42)  # For reproducibility\n\n{code}"

    def get_template(self, name: str) -> Optional[FixTemplate]:
        """Get a template by name.

        Args:
            name: Template name

        Returns:
            FixTemplate if found, None otherwise
        """
        return self._template_index.get(name)

    def add_template(self, template: FixTemplate) -> None:
        """Add a new template.

        Args:
            template: Template to add
        """
        self.templates.append(template)
        self._template_index[template.name] = template

    def register_transformer(
        self,
        name: str,
        transformer_factory: Callable[..., ast.NodeTransformer],
    ) -> None:
        """Register a custom AST transformer.

        Args:
            name: Template name this transformer handles
            transformer_factory: Factory function that creates the transformer
        """
        self._transformers[name] = transformer_factory

    def get_transformation_result(
        self,
        template: FixTemplate,
        source_code: str,
    ) -> TransformationResult:
        """Apply template and get detailed transformation result.

        Args:
            template: The fix template to apply
            source_code: The source code to transform

        Returns:
            TransformationResult with details about the transformation
        """
        try:
            transformed = self.apply_template(template, source_code)
            return TransformationResult(
                success=transformed != source_code,
                transformed_code=transformed,
                original_code=source_code,
                changes_made=[template.transformation] if transformed != source_code else [],
            )
        except Exception as e:
            return TransformationResult(
                success=False,
                transformed_code=source_code,
                original_code=source_code,
                error=str(e),
            )

    def suggest_fixes(
        self,
        context: FailureContext,
        max_suggestions: int = 3,
    ) -> List[Tuple[FixTemplate, str]]:
        """Suggest fixes for a failure context.

        Args:
            context: The failure context
            max_suggestions: Maximum number of suggestions to return

        Returns:
            List of (template, fixed_code) tuples
        """
        suggestions = []

        # Find matching templates
        templates = self.match_error_to_template(
            context.error_type,
            context.error_message,
        )

        for template in templates[:max_suggestions * 2]:  # Try more to account for failures
            fixed = self.synthesize_from_template(template, context)
            if fixed and fixed != (context.surrounding_code or context.code_snippet):
                suggestions.append((template, fixed))
                if len(suggestions) >= max_suggestions:
                    break

        return suggestions
