from __future__ import annotations

import re

import attrs

from codeflash_api.diff._base import Diff

REPLACE_IN_FILE_TAGS_RE = re.compile(
    r"<replace_in_file>(.*?)</replace_in_file>",
    re.DOTALL | re.MULTILINE,
)
MULTI_REPLACE_IN_FILE_RE = re.compile(
    r"<replace_in_file>\s*<path>(.*?)</path>"
    r"\s*<diff>([\s\S]*?)</diff>\s*</replace_in_file>",
    re.MULTILINE,
)

SEARCH_MARKER = "<<<<<<< SEARCH"
DELIMITER = "======="
REPLACE_MARKER = ">>>>>>> REPLACE"


@attrs.frozen
class SearchReplaceBlock:
    """
    One SEARCH/REPLACE pair.
    """

    search: str
    replace: str


def parse_diff(diff: str) -> list[SearchReplaceBlock]:
    """
    Parse SEARCH/REPLACE blocks from *diff*.
    """
    if not diff or not diff.strip():
        msg = "Empty or invalid diff string provided"
        raise ValueError(msg)

    lines = diff.splitlines(keepends=True)
    blocks: list[SearchReplaceBlock] = []
    search_start: int | None = None
    delimiter_idx: int | None = None

    for i, line in enumerate(lines):
        stripped = line.rstrip()
        if stripped == SEARCH_MARKER:
            search_start = i
        elif stripped == DELIMITER and search_start is not None:
            delimiter_idx = i
        elif (
            stripped == REPLACE_MARKER
            and delimiter_idx is not None
            and search_start is not None
        ):
            search_content = "".join(
                lines[search_start + 1 : delimiter_idx]
            ).rstrip()
            replace_content = "".join(lines[delimiter_idx + 1 : i]).rstrip()
            blocks.append(
                SearchReplaceBlock(
                    search=search_content,
                    replace=replace_content,
                )
            )
            search_start = None
            delimiter_idx = None

    if search_start is not None and delimiter_idx is None:
        msg = "Invalid diff format: Missing '=======' marker"
        raise ValueError(msg)
    if delimiter_idx is not None:
        msg = "Invalid diff format: Missing '>>>>>>> REPLACE' marker"
        raise ValueError(msg)
    if not blocks:
        msg = "No valid SEARCH/REPLACE blocks found in the diff"
        raise ValueError(msg)

    return blocks


def group_diff_patches_by_path(
    replace_tags_str: str,
) -> dict[str, str]:
    """
    Group multiple diffs by file path from XML tag content.
    """
    matches = MULTI_REPLACE_IN_FILE_RE.findall(replace_tags_str)
    result: dict[str, str] = {}
    for path, diff_content in matches:
        clean_path = path.strip()
        if clean_path in result:
            result[clean_path] += "\n" + diff_content
        else:
            result[clean_path] = diff_content
    return result


def extract_patches(content: str) -> dict[str, str]:
    """
    Extract all SEARCH/REPLACE patches from *content*.
    """
    matches = REPLACE_IN_FILE_TAGS_RE.findall(content)
    if not matches:
        return {}
    wrapped = f"<replace_in_file>{matches[0]}</replace_in_file>"
    return group_diff_patches_by_path(wrapped)


def find_with_whitespace_flexibility(
    search: str, content: str
) -> tuple[int, int] | None:
    """
    Match *search* in *content* with flexible whitespace.
    """
    parts = re.split(r"(\s+)", search)
    pattern_parts: list[str] = []
    for part in parts:
        if not part:
            continue
        if re.match(r"^\s+$", part):
            pattern_parts.append(r"\s+")
        else:
            pattern_parts.append(re.escape(part))

    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:
    """
    Parse and apply SEARCH/REPLACE blocks to *content*.
    """
    try:
        blocks = parse_diff(diff_str)
    except ValueError:
        return content

    for block in blocks:
        if not block.search and block.replace:
            content += block.replace
            continue
        if block.search:
            idx = content.find(block.search)
            if idx >= 0:
                content = content.replace(block.search, block.replace, 1)
            else:
                span = find_with_whitespace_flexibility(block.search, content)
                if span:
                    start, end = span
                    content = content[:start] + block.replace + content[end:]

    return content


class SearchAndReplaceDiff(Diff):
    """
    Apply SEARCH/REPLACE style diffs.
    """

    def run(self) -> dict[str, str]:
        """
        Apply patches to source files.
        """
        patches = extract_patches(self.content)
        result: dict[str, str] = {}

        if (
            not self.match_files_when_having_single_patch
            and len(self.source_code) == 1
            and len(patches) == 1
        ):
            src_path = next(iter(self.source_code))
            diff_content = next(iter(patches.values()))
            result[src_path] = apply_patches(
                diff_content, self.source_code[src_path]
            )
            return result

        for path, code in self.source_code.items():
            if path in patches:
                result[path] = apply_patches(patches[path], code)
            else:
                result[path] = code

        return result
