import re
from dataclasses import dataclass, field
from difflib import SequenceMatcher
from typing import Callable, DefaultDict, Literal, Optional

TOLERANCE_TYPES = Literal["SILENT", "WARNING", "ERROR"]


class SearchReplaceMatchError(Exception):
    def __init__(self, message: str):
        message = f"""
{message}
---
Last edit failed, no changes are applied. None of the search/replace blocks applied in the last tool call.
Recommendations:
- Retry immediately with same "percentage_to_change" using search replace blocks fixing above error.
- If you are still unsure you may re-read the file and then proceed accordingly.

If your search failed due to updates in the file content, remember that these
are the changes that the user made and you should preserve them
by updating your replace blocks too.
"""
        super().__init__(message)


@dataclass
class Tolerance:
    line_process: Callable[[str], str]
    severity_cat: TOLERANCE_TYPES
    score_multiplier: float
    error_name: str


@dataclass
class TolerancesHit(Tolerance):
    count: int


@dataclass
class FileEditOutput:
    original_content: list[str]
    orig_search_blocks: list[list[str]]
    edited_with_tolerances: list[
        tuple[slice, list[TolerancesHit], list[str]]
    ]  # Need not be equal to orig_search_blocks when early exit

    def replace_or_throw(
        self,
        max_errors: int,
    ) -> tuple[list[str], set[str]]:
        new_lines = list[str]()
        last_idx = 0
        errors = []
        warnings = set[str]()
        info = set[str]()
        score = 0.0
        for (span, tolerances, replace_with), search_ in zip(
            self.edited_with_tolerances, self.orig_search_blocks
        ):
            for tol in tolerances:
                score += tol.count * tol.score_multiplier
                if tol.count > 0:
                    if tol.severity_cat == "WARNING":
                        warnings.add(tol.error_name)
                    elif tol.severity_cat == "ERROR":
                        search__ = "\n".join(search_)
                        errors.append(f"""
Got error while processing the following search block:
---
```
{search__}
```
---
Error:
{tol.error_name}
---
                                  """)
                    else:
                        info.add(tol.error_name)
                    if len(errors) >= max_errors:
                        raise SearchReplaceMatchError("\n".join(errors))
            if last_idx < span.start:
                new_lines.extend(self.original_content[last_idx : span.start])

            new_lines.extend(replace_with)
            last_idx = span.stop

        if last_idx < len(self.original_content):
            new_lines.extend(self.original_content[last_idx:])

        if errors:
            raise SearchReplaceMatchError("\n".join(errors))

        if score > 1000:
            display = (list(warnings) + list(info))[:max_errors]
            raise SearchReplaceMatchError(
                "Too many warnings generated, not applying the edits\n"
                + "\n".join(display)
            )

        return new_lines, set(warnings)

    @staticmethod
    def get_best_match(
        outputs: list["FileEditOutput"],
    ) -> list["FileEditOutput"]:
        best_hits: list[FileEditOutput] = []
        best_score = float("-inf")
        assert outputs
        for output in outputs:
            hit_score = 0.0
            for _, tols, _ in output.edited_with_tolerances:
                for tol in tols:
                    hit_score += tol.count * tol.score_multiplier
            if not best_hits:
                best_hits.append(output)
                best_score = hit_score
            else:
                if hit_score < best_score:
                    best_hits = [output]
                    best_score = hit_score
                elif abs(hit_score - best_score) < 1e-3:
                    best_hits.append(output)
        return best_hits


def line_process_max_space_tolerance(line: str) -> str:
    line = line.strip()
    return re.sub(r"\s", "", line)


REMOVE_INDENTATION = (
    "Warning: matching without considering indentation (leading spaces)."
)
REMOVE_LINE_NUMS = "Warning: you gave search/replace blocks with leading line numbers, do not give them from the next time."

COMMON_MISTAKE_TRANSLATION = str.maketrans(
    {
        "‘": "'",
        "’": "'",
        "‚": ",",
        "‛": "'",
        "′": "'",
        "“": '"',
        "”": '"',
        "‟": '"',
        "″": '"',
        "‹": "<",
        "›": ">",
        "‐": "-",
        "‑": "-",
        "‒": "-",
        "–": "-",
        "—": "-",
        "―": "-",
        "−": "-",
        "…": "...",
    }
)


def remove_leading_linenums(string: str) -> str:
    return re.sub(r"^\d+ ", "", string).rstrip()


def normalize_common_mistakes(string: str) -> str:
    """Normalize unicode chars which are commonly confused by their ascii variants"""
    return string.translate(COMMON_MISTAKE_TRANSLATION).rstrip()


DEFAULT_TOLERANCES = [
    Tolerance(
        line_process=str.rstrip,
        severity_cat="SILENT",
        score_multiplier=1,
        error_name="",
    ),
    Tolerance(
        line_process=str.lstrip,
        severity_cat="WARNING",
        score_multiplier=10,
        error_name=REMOVE_INDENTATION,
    ),
    Tolerance(
        line_process=remove_leading_linenums,
        severity_cat="WARNING",
        score_multiplier=5,
        error_name=REMOVE_LINE_NUMS,
    ),
    Tolerance(
        line_process=normalize_common_mistakes,
        severity_cat="WARNING",
        score_multiplier=5,
        error_name="Warning: matching after normalizing commonly confused characters (quotes, dashes, ellipsis).",
    ),
    Tolerance(
        line_process=line_process_max_space_tolerance,
        severity_cat="WARNING",
        score_multiplier=50,
        error_name="Warning: matching after removing all spaces in lines.",
    ),
]


def fix_line_nums(
    matched_lines: list[str], searched_lines: list[str], replaced_lines: list[str]
) -> list[str]:
    return [remove_leading_linenums(line) for line in replaced_lines]


def fix_indentation(
    matched_lines: list[str], searched_lines: list[str], replaced_lines: list[str]
) -> list[str]:
    if not matched_lines or not searched_lines or not replaced_lines:
        return replaced_lines

    def get_indentation(line: str) -> str:
        match = re.match(r"^(\s*)", line)
        assert match
        return match.group(0)

    matched_indents = [get_indentation(line) for line in matched_lines if line.strip()]
    searched_indents = [
        get_indentation(line) for line in searched_lines if line.strip()
    ]
    if len(matched_indents) != len(searched_indents):
        return replaced_lines
    diffs: list[int] = [
        len(searched) - len(matched)
        for matched, searched in zip(matched_indents, searched_indents)
    ]
    if not all(diff == diffs[0] for diff in diffs):
        return replaced_lines
    if diffs[0] == 0:
        return replaced_lines

    # At this point we have same number of non-empty lines and the same indentation difference
    # We can now adjust the indentation of the replaced lines
    def adjust_indentation(line: str, diff: int) -> str:
        if diff < 0:
            return matched_indents[0][:-diff] + line
        return line[diff:]

    if diffs[0] > 0:
        if not (all(not line[: diffs[0]].strip() for line in replaced_lines)):
            return replaced_lines
    return [adjust_indentation(line, diffs[0]) for line in replaced_lines]


def remove_leading_trailing_empty_lines(lines: list[str]) -> list[str]:
    start = 0
    end = len(lines) - 1
    if end < start:
        return lines
    while not lines[start].strip():
        start += 1
        if start >= len(lines):
            break
    while not lines[end].strip():
        end -= 1
        if end < 0:
            break
    return lines[start : end + 1]


@dataclass
class FileEditInput:
    file_lines: list[str]
    file_line_offset: int
    search_replace_blocks: list[tuple[list[str], list[str]]]
    search_replace_offset: int
    tolerances: list["Tolerance"] = field(default_factory=lambda: DEFAULT_TOLERANCES)

    def edit_file(self) -> list[FileEditOutput]:
        n_file_lines = len(self.file_lines)
        n_blocks = len(self.search_replace_blocks)

        # Boundary conditions
        no_match_output = FileEditOutput(
            original_content=self.file_lines,
            orig_search_blocks=[x[0] for x in self.search_replace_blocks],
            edited_with_tolerances=[
                (
                    slice(0, 0),
                    [
                        TolerancesHit(
                            line_process=lambda x: x,
                            severity_cat="ERROR",
                            score_multiplier=float("inf"),
                            error_name="The blocks couldn't be matched, maybe the sequence of search blocks was incorrect?",
                            count=max(1, len(search_lines)),
                        )
                        for search_lines, _ in self.search_replace_blocks[
                            self.search_replace_offset :
                        ]
                    ],
                    [],
                )
            ],
        )
        if (
            self.file_line_offset >= n_file_lines
            and self.search_replace_offset < n_blocks
        ):
            return [no_match_output]
        elif self.file_line_offset >= n_file_lines:
            return [
                FileEditOutput(
                    self.file_lines,
                    [x[0] for x in self.search_replace_blocks],
                    [(slice(0, 0), [], [])],
                )
            ]
        elif self.search_replace_offset >= n_blocks:
            return [
                FileEditOutput(
                    self.file_lines,
                    [x[0] for x in self.search_replace_blocks],
                    [(slice(0, 0), [], [])],
                )
            ]

        # search for first block
        first_block = self.search_replace_blocks[self.search_replace_offset]
        replace_by = first_block[1]

        # Try exact match
        matches = match_exact(self.file_lines, self.file_line_offset, first_block[0])

        all_outputs = list[list[tuple[slice, list[TolerancesHit], list[str]]]]()

        if not matches:
            # Try tolerances
            matches_with_tolerances = match_with_tolerance(
                self.file_lines, self.file_line_offset, first_block[0], self.tolerances
            )
            if not matches_with_tolerances:
                # Try with no empty lines
                matches_with_tolerances = match_with_tolerance_empty_line(
                    self.file_lines,
                    self.file_line_offset,
                    first_block[0],
                    self.tolerances,
                )
                replace_by = remove_leading_trailing_empty_lines(first_block[1])

                if not matches_with_tolerances:
                    # Report edit distance
                    sim_match, sim_sim, sim_context = (
                        find_least_edit_distance_substring(
                            self.file_lines, self.file_line_offset, first_block[0]
                        )
                    )
                    if sim_match:
                        matches_with_tolerances = [
                            (
                                sim_match,
                                [
                                    TolerancesHit(
                                        lambda x: x,
                                        "ERROR",
                                        float("inf"),
                                        "Couldn't find match. Here's the latest snippet from the file which might be relevant for you to consider:\n```"
                                        + sim_context
                                        + "\n```",
                                        int(len(first_block[0]) // sim_sim),
                                    )
                                ],
                            )
                        ]

        else:
            matches_with_tolerances = [(match, []) for match in matches]

        for match, tolerances in matches_with_tolerances:
            if any(
                tolerance.error_name == REMOVE_LINE_NUMS for tolerance in tolerances
            ):
                replace_by = fix_line_nums(
                    self.file_lines[match.start : match.stop],
                    first_block[0],
                    replace_by,
                )
            if any(
                tolerance.error_name == REMOVE_INDENTATION for tolerance in tolerances
            ):
                replace_by = fix_indentation(
                    self.file_lines[match.start : match.stop],
                    first_block[0],
                    replace_by,
                )

            file_edit_input = FileEditInput(
                self.file_lines,
                match.stop,
                self.search_replace_blocks,
                self.search_replace_offset + 1,
                self.tolerances,
            )

            if any(tolerance.severity_cat == "ERROR" for tolerance in tolerances):
                # Exit early
                all_outputs.append(
                    [
                        (match, tolerances, replace_by),
                    ]
                )
            else:
                remaining_output = file_edit_input.edit_file()
                for rem_output in remaining_output:
                    all_outputs.append(
                        [
                            (match, tolerances, replace_by),
                            *rem_output.edited_with_tolerances,
                        ]
                    )

        if not all_outputs:
            return [no_match_output]

        return [
            FileEditOutput(
                self.file_lines, [x[0] for x in self.search_replace_blocks], output
            )
            for output in all_outputs
        ]


def find_contiguous_match(search_line_positions: list[set[int]]) -> list[slice]:
    n_search_lines = len(search_line_positions)

    def search_in_dictionary(search_offset: int, search_index: int) -> bool:
        if search_offset >= n_search_lines:
            return True

        if search_index in search_line_positions[search_offset]:
            return search_in_dictionary(search_offset + 1, search_index + 1)
        return False

    matched_slices = []
    for index in search_line_positions[0]:
        if search_in_dictionary(1, index + 1):
            matched_slices.append(slice(index, index + n_search_lines, 1))
    return matched_slices


def match_exact(
    content: list[str], content_offset: int, search: list[str]
) -> list[slice]:
    n_search_lines = len(search)
    n_content = len(content) - content_offset
    if n_search_lines > n_content:
        return []
    if n_search_lines == 0:
        return []
    if n_content == 0:
        return []
    content_positions = DefaultDict[str, set[int]](set)
    for i in range(content_offset, n_content):
        content_positions[content[i]].add(i)
    search_line_positions = [content_positions[line] for line in search]

    matched_slices = find_contiguous_match(search_line_positions)

    return matched_slices


def match_with_tolerance(
    content: list[str],
    content_offset: int,
    search: list[str],
    tolerances: list[Tolerance],
) -> list[tuple[slice, list[TolerancesHit]]]:
    n_search_lines = len(search)
    n_content = len(content) - content_offset
    if n_search_lines > n_content:
        return []
    if n_search_lines == 0:
        return []
    if n_content == 0:
        return []
    content_positions = DefaultDict[str, set[int]](set)
    for i in range(content_offset, n_content):
        content_positions[content[i]].add(i)
    search_line_positions = [content_positions[line] for line in search]

    tolerance_index_by_content_line: list[dict[int, int]] = [
        {} for _ in range(len(search))
    ]
    for tidx, tolerance in enumerate(tolerances):
        content_positions = DefaultDict[str, set[int]](set)
        for i in range(content_offset, n_content):
            line = content[i]
            content_positions[tolerance.line_process(line)].add(i)
        for i, line in enumerate(search):
            new_lines = content_positions[tolerance.line_process(line)]
            new_indices = new_lines - search_line_positions[i]
            search_line_positions[i].update(new_indices)
            tolerance_index_by_content_line[i].update(
                {idx: tidx for idx in new_indices}
            )
    matched_slices = find_contiguous_match(search_line_positions)

    tolerances_counts: list[list[TolerancesHit]] = [
        [
            TolerancesHit(
                line_process=tol.line_process,
                severity_cat=tol.severity_cat,
                score_multiplier=tol.score_multiplier,
                count=0,
                error_name=tol.error_name,
            )
            for tol in tolerances
        ]
        for _ in range(len(matched_slices))
    ]
    for sidx, slice in enumerate(matched_slices):
        for search_idx, content_idx in enumerate(
            range(slice.start, slice.stop, slice.step)
        ):
            if content_idx in tolerance_index_by_content_line[search_idx]:
                tolerances_counts[sidx][
                    tolerance_index_by_content_line[search_idx][content_idx]
                ].count += 1

    # Remove 0 counts
    tolerances_counts = [[x for x in y if x.count > 0] for y in tolerances_counts]

    return list(zip(matched_slices, tolerances_counts))


def match_with_tolerance_empty_line(
    content: list[str],
    content_offset: int,
    search: list[str],
    tolerances: list[Tolerance],
) -> list[tuple[slice, list[TolerancesHit]]]:
    new_content = list[str]()
    new_to_original = dict[int, int]()
    for i in range(content_offset, len(content)):
        line = content[i]
        if line.strip():
            new_to_original[len(new_content)] = i
            new_content.append(line)

    search = [line for line in search if line.strip()]

    matches_with_tolerancs = match_with_tolerance(new_content, 0, search, tolerances)

    new_matches_with_tolerances = list[tuple[slice, list[TolerancesHit]]]()
    for matches, tolerance_counts in matches_with_tolerancs:
        matches = slice(
            new_to_original[matches.start], new_to_original[matches.stop - 1] + 1, 1
        )
        new_matches_with_tolerances.append((matches, tolerance_counts))
    return new_matches_with_tolerances


def find_least_edit_distance_substring(
    orig_content_lines: list[str], offset: int, find_lines: list[str]
) -> tuple[Optional[slice], float, str]:
    # Prepare content lines, stripping whitespace and keeping track of original indices
    content_lines = [
        orig_content_lines[i].strip() for i in range(offset, len(orig_content_lines))
    ]
    new_to_original_indices = {}
    new_content_lines = []
    for i, line in enumerate(content_lines):
        if not line:
            continue
        new_content_lines.append(line)
        new_to_original_indices[len(new_content_lines) - 1] = i
    content_lines = new_content_lines

    # Prepare find lines, removing empty lines
    find_lines = [line.strip() for line in find_lines if line.strip()]

    # Initialize variables for best match tracking
    max_similarity = 0.0
    min_edit_distance_lines = None
    context_lines = []

    # For each possible starting position in content
    for i in range(max(1, len(content_lines) - len(find_lines) + 1)):
        # Calculate similarity for the block starting at position i
        block_similarity = 0.0
        for j in range(len(find_lines)):
            if (i + j) < len(content_lines):
                # Use SequenceMatcher for more efficient similarity calculation
                similarity = SequenceMatcher(
                    None, content_lines[i + j], find_lines[j]
                ).ratio()
                block_similarity += similarity

        # If this block is more similar than previous best
        if block_similarity > max_similarity:
            max_similarity = block_similarity
            # Map back to original line indices
            orig_start_index = new_to_original_indices[i]
            orig_end_index = (
                new_to_original_indices.get(
                    i + len(find_lines) - 1, len(orig_content_lines) - 1
                )
                + 1
            )
            # Get the original lines
            min_edit_distance_lines = slice(
                orig_start_index + offset, orig_end_index + offset
            )
            # Get context (10 lines before and after)
            context_lines = orig_content_lines[
                max(0, orig_start_index - 10 + offset) : (orig_end_index + 10 + offset)
            ]

    return (
        min_edit_distance_lines,
        max_similarity,
        "\n".join(context_lines),
    )
