import re

from pydantic import ValidationError

from core.languages.python.optimizer.context_utils.constants import (
    MULTI_REPLACE_IN_FILE_TAGS_REGEX,
    REPLACE_IN_FILE_TAGS_REGEX,
)
from core.languages.python.optimizer.diff_patches_utils.diff import Diff


class SearchReplaceBlock:
    def __init__(self, search, replace):
        self.search = search
        self.replace = replace

    @classmethod
    def from_block(self, search: str, replace: str) -> "SearchReplaceBlock":
        """Creates a SearchReplaceBlock instance from search and replace strings."""
        return self(search=search, replace=replace)


def parse_diff(diff: str) -> list[SearchReplaceBlock]:
    """Parses the diff string into a list of SearchReplaceBlock."""
    if not diff or not diff.strip():
        raise ValueError("Empty or invalid diff string provided")

    blocks: list[SearchReplaceBlock] = []
    lines = diff.splitlines(keepends=True)
    n = len(lines)
    idx = 0

    # Precompute the "marker" strings for efficiency
    search_marker = "<<<<<<< SEARCH"
    delimiter_marker = "======="
    replace_marker = ">>>>>>> REPLACE"

    while idx < n:
        line_stripped = lines[idx].strip()
        if line_stripped.startswith(search_marker):
            idx += 1
            search_start = idx
            # Find delimiter_marker line
            while idx < n and lines[idx].strip() != delimiter_marker:
                idx += 1
            search_end = idx

            if idx >= n:
                raise ValueError("Invalid diff format: Missing '=======' marker")
            idx += 1
            replace_start = idx
            while idx < n and not lines[idx].strip().startswith(replace_marker):
                idx += 1
            replace_end = idx

            if idx >= n:
                raise ValueError("Invalid diff format: Missing '>>>>>>> REPLACE' marker")

            search_content = "".join(lines[search_start:search_end]).rstrip()
            replace_content = "".join(lines[replace_start:replace_end]).rstrip()

            try:
                block = SearchReplaceBlock.from_block(search=search_content, replace=replace_content)
                blocks.append(block)
            except ValidationError as ve:
                raise ValueError(f"Invalid block format: {ve}")

        idx += 1

    if not blocks:
        raise ValueError("No valid SEARCH/REPLACE blocks found in the diff")

    return blocks


def group_diff_patches_by_path(replace_tags_str: str) -> dict[str, str]:
    matches = MULTI_REPLACE_IN_FILE_TAGS_REGEX.findall(replace_tags_str)
    file_to_diffs = {}

    current_file = None
    current_diff = ""

    for path, diff in matches:
        if path != current_file:
            if current_file:
                file_to_diffs[current_file] = current_diff
            current_file = path
            current_diff = diff
        else:
            current_diff += diff

    if current_file:
        file_to_diffs[current_file] = current_diff

    return file_to_diffs


def extract_patches(content: str) -> dict[str, str]:
    """Return a mapping between the file path and the search/replace diff content."""
    matches = REPLACE_IN_FILE_TAGS_REGEX.findall(content)
    replace_tags = ""
    if matches and len(matches) != 0:
        replace_tags = f"<replace_in_file>{matches[0]}</replace_in_file>"

    if replace_tags == "":
        return {}

    return group_diff_patches_by_path(replace_tags)


def find_with_whitespace_flexibility(search: str, content: str) -> tuple[int, int] | None:
    """Find the search string in content with whitespace flexibility.

    This handles cases where whitespace differs between the search block and the actual content
    (e.g., different indentation, trailing whitespace, tabs vs spaces).

    Returns (start_idx, end_idx) of the match in the original content, or None if not found.
    """
    # Split by whitespace sequences, preserving both whitespace and non-whitespace parts
    parts = re.split(r"(\s+)", search)
    pattern_parts = []

    for part in parts:
        if not part:
            continue
        if re.match(r"^\s+$", part):
            # This is a whitespace sequence - make it flexible (match any whitespace)
            pattern_parts.append(r"\s+")
        else:
            # This is non-whitespace - escape it for literal matching
            pattern_parts.append(re.escape(part))

    if not pattern_parts:
        return None

    pattern = "".join(pattern_parts)

    match = re.search(pattern, content, re.MULTILINE | re.DOTALL)
    if match:
        return match.start(), match.end()
    return None


def apply_patches(diff_str: str, content: str) -> str:
    try:
        patch_blocks = parse_diff(diff_str)
    except ValueError:
        return content

    for block in patch_blocks:
        if not block.search:
            if block.replace:
                # a replacement block without a search, then just add the replace block
                content += block.replace
            continue

        if block.replace:
            # check for exact match (most of the time it will find the exact match when we get the diff from the LLM)
            start_char_idx = content.find(block.search)
            if start_char_idx != -1:
                end_char_idx = start_char_idx + len(block.search)
                content = f"{content[:start_char_idx]}{block.replace}{content[end_char_idx:]}"
            else:
                # Fallback: try whitespace-flexible matching
                match_result = find_with_whitespace_flexibility(block.search, content)
                if match_result:
                    start_char_idx, end_char_idx = match_result
                    content = f"{content[:start_char_idx]}{block.replace}{content[end_char_idx:]}"
    return content


class SearchAndReplaceDiff(Diff):
    def run(self) -> dict[str, str]:
        file_to_diff = extract_patches(self.content)
        file_to_original_code = self.source_code
        new_file_to_code: dict[str, str] = {}

        # in deprecated single context we don't know the file name so we ignore it completely
        single_patch = len(file_to_original_code.keys()) == 1 and len(file_to_diff.keys()) == 1
        if single_patch and not self.match_files_when_having_single_patch:
            diffs = list(file_to_diff.values())
            if diffs:
                original_content = next(iter(file_to_original_code.values()))
                return {"file": apply_patches(diffs[0], original_content)}

        for file_name, original_content in file_to_original_code.items():
            diff = file_to_diff.get(file_name)
            if diff is None:
                # no change
                new_file_to_code[file_name] = original_content
            else:
                new_file_to_code[file_name] = apply_patches(diff, original_content)

        return new_file_to_code
