"""Aider-style SEARCH/REPLACE format for applying code fixes.

Based on Aider's proven approach that works better with LLMs than unified diffs:
- No line numbers (LLMs struggle with them)
- Unique context makes search unambiguous
- Layered matching: exact → whitespace → fuzzy → AST
- Small blocks are better than large ones

Format:
    path/to/file.py
    <<<<<<< SEARCH
    old code here
    =======
    new code here
    >>>>>>> REPLACE

Usage:
    blocks = parse_search_replace(llm_response)
    for block in blocks:
        result = apply_block(file_content, block)
        if result.success:
            file_content = result.new_content
"""

from __future__ import annotations

import ast
import difflib
import re
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import List, Optional, Tuple


class MatchType(Enum):
    """Type of match used to apply the block."""
    EXACT = "exact"
    WHITESPACE = "whitespace"
    FUZZY = "fuzzy"
    AST = "ast"
    FAILED = "failed"


@dataclass
class SearchReplaceBlock:
    """A single SEARCH/REPLACE block."""
    file_path: str
    search: str
    replace: str
    original_text: str = ""  # The raw block text

    def __post_init__(self):
        # Normalize line endings
        self.search = self.search.replace("\r\n", "\n")
        self.replace = self.replace.replace("\r\n", "\n")


@dataclass
class ApplyResult:
    """Result of applying a SEARCH/REPLACE block."""
    success: bool
    match_type: MatchType
    new_content: str
    error_message: Optional[str] = None
    match_location: Optional[Tuple[int, int]] = None  # (start_line, end_line)


def parse_search_replace(text: str) -> List[SearchReplaceBlock]:
    """Parse LLM response containing SEARCH/REPLACE blocks.

    Handles multiple blocks in a single response.
    Flexible parsing to handle minor formatting variations.

    Args:
        text: Raw LLM response text

    Returns:
        List of parsed blocks
    """
    blocks: List[SearchReplaceBlock] = []

    # Pattern for SEARCH/REPLACE blocks
    # Flexible: handles variations in markers
    pattern = r'''
        ^(.+?)[\s]*\n            # File path on its own line
        <{6,7}\s*SEARCH[\s]*\n   # <<<<<<< SEARCH
        (.*?)                    # Search content (non-greedy)
        \n={6,7}[\s]*\n          # =======
        (.*?)                    # Replace content (non-greedy)
        \n>{6,7}\s*REPLACE       # >>>>>>> REPLACE
    '''

    for match in re.finditer(pattern, text, re.MULTILINE | re.DOTALL | re.VERBOSE):
        file_path = match.group(1).strip()
        search = match.group(2)
        replace = match.group(3)

        # Skip if file path looks like code
        if file_path.startswith(("def ", "class ", "import ", "from ")):
            continue

        blocks.append(SearchReplaceBlock(
            file_path=file_path,
            search=search,
            replace=replace,
            original_text=match.group(0),
        ))

    return blocks


def parse_search_replace_relaxed(text: str) -> List[SearchReplaceBlock]:
    """More relaxed parsing for edge cases.

    Tries multiple patterns and strategies.
    """
    blocks = parse_search_replace(text)
    if blocks:
        return blocks

    # Try alternative patterns
    alt_patterns = [
        # Without file path on separate line
        r'```(\w+)?\n(.+?)\n<<<<<<< SEARCH\n(.*?)\n=======\n(.*?)\n>>>>>>> REPLACE\n```',
        # With different markers
        r'FILE:\s*(.+?)\n<<<\n(.*?)\n===\n(.*?)\n>>>',
    ]

    for pattern in alt_patterns:
        for match in re.finditer(pattern, text, re.DOTALL):
            groups = match.groups()
            if len(groups) >= 3:
                file_path = groups[0] if len(groups) == 4 else groups[0]
                search = groups[-2]
                replace = groups[-1]
                blocks.append(SearchReplaceBlock(
                    file_path=file_path.strip() if file_path else "unknown",
                    search=search,
                    replace=replace,
                ))

    return blocks


class LayeredMatcher:
    """Apply SEARCH/REPLACE blocks with increasingly fuzzy matching."""

    def __init__(self, fuzzy_threshold: float = 0.8):
        """Initialize matcher.

        Args:
            fuzzy_threshold: Minimum similarity for fuzzy matching (0-1)
        """
        self.fuzzy_threshold = fuzzy_threshold

    def apply_block(
        self,
        file_content: str,
        block: SearchReplaceBlock,
    ) -> ApplyResult:
        """Apply a SEARCH/REPLACE block using layered matching.

        Tries in order:
        1. Exact match
        2. Whitespace-normalized match
        3. Fuzzy line-by-line match
        4. AST-based semantic match (for Python)

        Args:
            file_content: Current file content
            block: The block to apply

        Returns:
            ApplyResult with new content if successful
        """
        # Level 1: Exact match
        result = self._try_exact_match(file_content, block)
        if result.success:
            return result

        # Level 2: Whitespace normalized
        result = self._try_whitespace_match(file_content, block)
        if result.success:
            return result

        # Level 3: Fuzzy match
        result = self._try_fuzzy_match(file_content, block)
        if result.success:
            return result

        # Level 4: AST match (Python only)
        if block.file_path.endswith(".py"):
            result = self._try_ast_match(file_content, block)
            if result.success:
                return result

        return ApplyResult(
            success=False,
            match_type=MatchType.FAILED,
            new_content=file_content,
            error_message=f"Could not find match for SEARCH block in {block.file_path}",
        )

    def _try_exact_match(
        self,
        file_content: str,
        block: SearchReplaceBlock,
    ) -> ApplyResult:
        """Try exact string match."""
        if block.search in file_content:
            new_content = file_content.replace(block.search, block.replace, 1)

            # Find match location for reporting
            start_pos = file_content.find(block.search)
            start_line = file_content[:start_pos].count("\n") + 1
            end_line = start_line + block.search.count("\n")

            return ApplyResult(
                success=True,
                match_type=MatchType.EXACT,
                new_content=new_content,
                match_location=(start_line, end_line),
            )

        return ApplyResult(
            success=False,
            match_type=MatchType.FAILED,
            new_content=file_content,
        )

    def _try_whitespace_match(
        self,
        file_content: str,
        block: SearchReplaceBlock,
    ) -> ApplyResult:
        """Try matching with normalized whitespace."""
        def normalize(text: str) -> str:
            # Normalize each line individually to preserve structure
            lines = text.split("\n")
            return "\n".join(line.strip() for line in lines)

        normalized_search = normalize(block.search)
        normalized_content = normalize(file_content)

        if normalized_search in normalized_content:
            # Find the actual location in original content
            # by matching line by line
            search_lines = block.search.strip().split("\n")
            content_lines = file_content.split("\n")

            for i in range(len(content_lines) - len(search_lines) + 1):
                window = content_lines[i:i + len(search_lines)]
                if all(
                    a.strip() == b.strip()
                    for a, b in zip(window, search_lines)
                ):
                    # Found match at line i
                    # Preserve original indentation from first line
                    original_indent = len(content_lines[i]) - len(content_lines[i].lstrip())
                    indent = " " * original_indent

                    # Apply replacement with original indentation
                    replace_lines = block.replace.split("\n")
                    indented_replace = []
                    for j, line in enumerate(replace_lines):
                        if j == 0:
                            indented_replace.append(indent + line.lstrip())
                        elif line.strip():
                            # Detect relative indent in replacement
                            rel_indent = len(line) - len(line.lstrip())
                            base_indent = len(replace_lines[0]) - len(replace_lines[0].lstrip())
                            extra = max(0, rel_indent - base_indent)
                            indented_replace.append(indent + " " * extra + line.strip())
                        else:
                            indented_replace.append("")

                    # Replace lines
                    new_lines = (
                        content_lines[:i] +
                        indented_replace +
                        content_lines[i + len(search_lines):]
                    )

                    return ApplyResult(
                        success=True,
                        match_type=MatchType.WHITESPACE,
                        new_content="\n".join(new_lines),
                        match_location=(i + 1, i + len(search_lines)),
                    )

        return ApplyResult(
            success=False,
            match_type=MatchType.FAILED,
            new_content=file_content,
        )

    def _try_fuzzy_match(
        self,
        file_content: str,
        block: SearchReplaceBlock,
    ) -> ApplyResult:
        """Try fuzzy line-by-line matching."""
        search_lines = block.search.strip().split("\n")
        content_lines = file_content.split("\n")

        # Look for the best matching window
        best_match: Optional[Tuple[int, float]] = None

        for i in range(len(content_lines) - len(search_lines) + 1):
            window = content_lines[i:i + len(search_lines)]
            window_text = "\n".join(window)
            search_text = "\n".join(search_lines)

            similarity = difflib.SequenceMatcher(
                None, window_text, search_text
            ).ratio()

            if similarity >= self.fuzzy_threshold:
                if best_match is None or similarity > best_match[1]:
                    best_match = (i, similarity)

        if best_match:
            i, similarity = best_match

            # Get original indentation
            original_indent = len(content_lines[i]) - len(content_lines[i].lstrip())
            indent = " " * original_indent

            # Apply replacement with indentation
            replace_lines = block.replace.split("\n")
            indented_replace = []
            for j, line in enumerate(replace_lines):
                if line.strip():
                    if j == 0:
                        indented_replace.append(indent + line.strip())
                    else:
                        # Try to preserve relative indentation
                        indented_replace.append(indent + "    " + line.strip())
                else:
                    indented_replace.append("")

            new_lines = (
                content_lines[:i] +
                indented_replace +
                content_lines[i + len(search_lines):]
            )

            return ApplyResult(
                success=True,
                match_type=MatchType.FUZZY,
                new_content="\n".join(new_lines),
                match_location=(i + 1, i + len(search_lines)),
            )

        return ApplyResult(
            success=False,
            match_type=MatchType.FAILED,
            new_content=file_content,
        )

    def _try_ast_match(
        self,
        file_content: str,
        block: SearchReplaceBlock,
    ) -> ApplyResult:
        """Try AST-based matching for Python code.

        Finds functions/classes with matching names and structure.
        """
        try:
            file_tree = ast.parse(file_content)
            search_tree = ast.parse(block.search)
        except SyntaxError:
            return ApplyResult(
                success=False,
                match_type=MatchType.FAILED,
                new_content=file_content,
            )

        # Find the main node in search (function or class)
        search_node = None
        for node in ast.walk(search_tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                search_node = node
                break

        if not search_node:
            return ApplyResult(
                success=False,
                match_type=MatchType.FAILED,
                new_content=file_content,
            )

        # Find matching node in file
        for node in ast.walk(file_tree):
            if type(node) == type(search_node) and node.name == search_node.name:
                # Found matching function/class
                start_line = node.lineno
                end_line = node.end_lineno or start_line

                lines = file_content.split("\n")

                # Get indentation from original
                original_line = lines[start_line - 1]
                indent = len(original_line) - len(original_line.lstrip())

                # Indent the replacement
                replace_lines = block.replace.split("\n")
                indented = []
                for j, line in enumerate(replace_lines):
                    if line.strip():
                        indented.append(" " * indent + line.lstrip())
                    else:
                        indented.append("")

                new_lines = lines[:start_line - 1] + indented + lines[end_line:]

                return ApplyResult(
                    success=True,
                    match_type=MatchType.AST,
                    new_content="\n".join(new_lines),
                    match_location=(start_line, end_line),
                )

        return ApplyResult(
            success=False,
            match_type=MatchType.FAILED,
            new_content=file_content,
        )


def apply_blocks_to_file(
    file_path: Path,
    blocks: List[SearchReplaceBlock],
    dry_run: bool = False,
) -> List[ApplyResult]:
    """Apply multiple SEARCH/REPLACE blocks to a file.

    Args:
        file_path: Path to the file
        blocks: Blocks to apply (filtered to this file)
        dry_run: If True, don't write changes

    Returns:
        List of results for each block
    """
    if not file_path.exists():
        return [ApplyResult(
            success=False,
            match_type=MatchType.FAILED,
            new_content="",
            error_message=f"File not found: {file_path}",
        ) for _ in blocks]

    content = file_path.read_text()
    matcher = LayeredMatcher()
    results = []

    for block in blocks:
        result = matcher.apply_block(content, block)
        results.append(result)

        if result.success:
            content = result.new_content

    if not dry_run and any(r.success for r in results):
        file_path.write_text(content)

    return results


def blocks_to_unified_diff(
    file_path: str,
    original_content: str,
    blocks: List[SearchReplaceBlock],
) -> str:
    """Convert SEARCH/REPLACE blocks to unified diff format.

    Useful for preview and for tools that expect unified diffs.

    Args:
        file_path: Path to file
        original_content: Original file content
        blocks: Blocks to apply

    Returns:
        Unified diff string
    """
    matcher = LayeredMatcher()
    content = original_content

    for block in blocks:
        if block.file_path == file_path:
            result = matcher.apply_block(content, block)
            if result.success:
                content = result.new_content

    original_lines = original_content.splitlines(keepends=True)
    new_lines = content.splitlines(keepends=True)

    diff = difflib.unified_diff(
        original_lines,
        new_lines,
        fromfile=f"a/{file_path}",
        tofile=f"b/{file_path}",
    )

    return "".join(diff)
