import re
from dataclasses import dataclass, field

from core.languages.python.optimizer.diff_patches_utils.diff import Diff, DiffError


@dataclass
class Chunk:
    orig_index: int = -1  # Line number in the *original* file block where the change starts
    del_lines: list[str] = field(default_factory=list)
    ins_lines: list[str] = field(default_factory=list)


@dataclass
class PatchAction:
    path: str
    chunks: list[Chunk] = field(default_factory=list)


@dataclass
class Patch:
    actions: dict[str, PatchAction] = field(default_factory=dict)
    fuzz: int = 0  # Track fuzziness used during parsing


UPDATE_FILE_PREFIX = "*** Update File: "  # followed by file path


def _norm(line: str) -> str:
    """Strip CR so comparisons work for both LF and CRLF input."""
    return line.rstrip("\r")


def find_context_core(lines: list[str], context: list[str], start: int) -> tuple[int, int]:
    """Finds context block, returns start index and fuzz level."""  # noqa: D401
    if not context:
        return start, 0

    # Exact match
    for i in range(start, len(lines) - len(context) + 1):
        if lines[i : i + len(context)] == context:
            return i, 0
    # Rstrip match
    norm_context = [s.rstrip() for s in context]
    for i in range(start, len(lines) - len(context) + 1):
        if [s.rstrip() for s in lines[i : i + len(context)]] == norm_context:
            return i, 1  # Fuzz level 1
    # Strip match
    norm_context_strip = [s.strip() for s in context]
    for i in range(start, len(lines) - len(context) + 1):
        if [s.strip() for s in lines[i : i + len(context)]] == norm_context_strip:
            return i, 100  # Fuzz level 100
    return -1, 0


def find_context(lines: list[str], context: list[str], start: int, eof: bool) -> tuple[int, int]:  # noqa: FBT001
    """Finds context, handling EOF marker."""  # noqa: D401
    if eof:
        # If EOF marker, first try matching at the very end
        if len(lines) >= len(context):
            new_index, fuzz = find_context_core(lines, context, len(lines) - len(context))
            if new_index != -1:
                return new_index, fuzz
        # If not found at end, search from `start` as fallback
        new_index, fuzz = find_context_core(lines, context, start)
        return new_index, fuzz + 10_000  # Add large fuzz penalty if EOF wasn't at end
    # Normal case: search from `start`
    return find_context_core(lines, context, start)


def peek_next_section(lines: list[str], index: int) -> tuple[list[str], list[Chunk], int, bool]:
    """Parse one section (context, -, + lines) of an Update block.
    Returns: (context_lines, chunks_in_section, next_index, is_eof).
    """  # noqa: D205
    context_lines: list[str] = []
    del_lines: list[str] = []
    ins_lines: list[str] = []
    chunks: list[Chunk] = []
    mode = "keep"  # Start by expecting context lines
    start_index = index

    while index < len(lines):
        line = lines[index]
        norm_line = _norm(line)

        # Check for section terminators
        if norm_line.startswith(
            (
                "@@",
                "*** End Patch",
                "*** Update File:",
                "*** Delete File:",
                "*** Add File:",
                "*** End of File",  # Special terminator
            )
        ):
            break
        if norm_line == "***":  # Legacy/alternative terminator? Handle just in case.
            break
        if norm_line.startswith("***"):  # Invalid line
            msg = f"Invalid patch line found in update section: {line}"
            raise DiffError(msg)

        index += 1
        last_mode = mode

        # Determine line type and strip prefix
        if line.startswith("+"):
            mode = "add"
            line_content = line[1:]
        elif line.startswith("-"):
            mode = "delete"
            line_content = line[1:]
        elif line.startswith(" "):
            mode = "keep"
            line_content = line[1:]
        elif line.strip() == "":  # Treat blank lines in patch as context ' '
            mode = "keep"
            line_content = ""  # Keep it as a blank line
        else:
            # Assume lines without prefix are context if format is loose,
            # but strict format requires ' '. Raise error for strictness.
            msg = f"Invalid line prefix in update section: {line}"
            raise DiffError(msg)

        # If mode changes from add/delete back to keep, finalize the previous chunk
        if mode == "keep" and last_mode != "keep":
            if del_lines or ins_lines:
                chunks.append(
                    Chunk(
                        # orig_index is relative to the start of the *context* block found
                        orig_index=len(context_lines) - len(del_lines),
                        del_lines=del_lines,
                        ins_lines=ins_lines,
                    )
                )
            del_lines, ins_lines = [], []

        # Collect lines based on mode
        if mode == "delete":
            del_lines.append(line_content)
            context_lines.append(line_content)  # Deleted lines are part of the original context
        elif mode == "add":
            ins_lines.append(line_content)
        elif mode == "keep":
            context_lines.append(line_content)

    # Finalize any pending chunk at the end of the section
    if del_lines or ins_lines:
        chunks.append(Chunk(orig_index=len(context_lines) - len(del_lines), del_lines=del_lines, ins_lines=ins_lines))

    # Check for EOF marker
    is_eof = False
    if index < len(lines) and _norm(lines[index]) == "*** End of File":
        index += 1
        is_eof = True

    if index == start_index and not is_eof:  # Should not happen if patch is well-formed
        # raise DiffError("Empty patch section found.")
        return context_lines, [], index + 1, is_eof

    return context_lines, chunks, index, is_eof


def parse_patch_text(lines: list[str], current_files: dict[str, str]) -> Patch:
    """Parse patch content lines into a Patch object.

    Adapted from the Parser class in apply_patch.py.
    """
    patch = Patch()
    index = 0
    total_fuzz_count = 0

    while index < len(lines):
        line = lines[index]
        norm_line = _norm(line)

        # Allow blank lines between actions
        if not norm_line.strip():
            index += 1
            continue

        # ---------- UPDATE ---------- #
        if norm_line.startswith(UPDATE_FILE_PREFIX):
            index += 1

            path = norm_line[len(UPDATE_FILE_PREFIX) :].strip()
            if not path:
                raise DiffError("Update File action missing path.")

            if path not in current_files:
                msg = f"Update File Error - missing file content for: {path}"
                raise DiffError(msg)

            file_content = current_files[path]

            existing_action = patch.actions.get(path)
            if existing_action is not None:
                new_action, index, fuzz = _parse_update_file_sections(lines, index, file_content)
                existing_action.chunks.extend(new_action.chunks)
                total_fuzz_count += fuzz
            else:
                # First UPDATE block for this file
                action, index, fuzz = _parse_update_file_sections(lines, index, file_content)
                action.path = path
                patch.actions[path] = action
                total_fuzz_count += fuzz
            continue

        # If we are here, the line is unexpected
        msg = f"Unknown or misplaced line while parsing patch: {line}"
        raise DiffError(msg)

    patch.fuzz = total_fuzz_count
    return patch


def _parse_update_file_sections(lines: list[str], index: int, file_content: str) -> tuple[PatchAction, int, int]:
    """Parse all sections (@@, context, -, +) for a single Update File action."""
    action = PatchAction(path="")
    orig_lines = file_content.splitlines()  # Use splitlines for consistency
    current_file_index = 0  # Track position in original file content
    total_fuzz = 0

    while index < len(lines):
        # Handle @@ scope lines
        scope_lines = []
        while index < len(lines) and _norm(lines[index]).startswith("@@"):
            scope_line_content = lines[index][len("@@") :].strip()
            if scope_line_content:  # Ignore empty @@ lines?
                scope_lines.append(scope_line_content)
            index += 1

        # Find the scope in the original file if specified
        if scope_lines:
            # Simple scope finding: search from current position
            found_scope = False
            temp_index = current_file_index
            while temp_index < len(orig_lines):
                # Check if all scope lines match sequentially from temp_index
                match = True
                for i, scope in enumerate(scope_lines):
                    if temp_index + i >= len(orig_lines) or _norm(orig_lines[temp_index + i]).strip() != scope:
                        match = False
                        break
                if match:
                    current_file_index = temp_index + len(scope_lines)
                    found_scope = True
                    break
                temp_index += 1

            if not found_scope:
                # Try fuzzy scope matching (strip whitespace)
                temp_index = current_file_index
                while temp_index < len(orig_lines):
                    match = True
                    for i, scope in enumerate(scope_lines):
                        if (
                            temp_index + i >= len(orig_lines)
                            or _norm(orig_lines[temp_index + i]).strip() != scope.strip()
                        ):
                            match = False
                            break
                    if match:
                        current_file_index = temp_index + len(scope_lines)
                        found_scope = True
                        total_fuzz += 1  # Add fuzz for scope match difference
                        break
                    temp_index += 1

            if not found_scope:
                scope_txt = "\n".join(scope_lines)
                msg = f"Could not find scope context:\n{scope_txt} in \n{file_content}"
                raise DiffError(msg)

        # Peek and parse the next context/change section
        context_block, chunks_in_section, next_index, is_eof = peek_next_section(lines, index)

        # Find where this context block appears in the original file
        found_index, fuzz = find_context(orig_lines, context_block, current_file_index, is_eof)
        total_fuzz += fuzz

        if found_index == -1:
            ctx_txt = "\n".join(context_block)
            marker = "*** End of File" if is_eof else ""
            msg = f"Could not find patch context '{marker}' starting near line {current_file_index}:\n{ctx_txt}"
            raise DiffError(msg)

        # Adjust chunk original indices to be absolute within the file
        for chunk in chunks_in_section:
            # chunk.orig_index from peek is relative to context_block start
            # We need it relative to the file start
            chunk.orig_index += found_index
            action.chunks.append(chunk)

        # Advance file index past the matched context block
        current_file_index = found_index + len(context_block)
        # Advance line index past the processed section in the patch
        index = next_index

    return action, index, total_fuzz


def extract_update_sections(content: str) -> dict[str, str]:
    """Return a mapping between the file path and the update diff content.

    Note that this will only return the pure update section only without begin and end patch markers,
    for example

    *** Update File: file.py
    @@
    - a + b
    + a - b
    @@
    - print("Hello")
    + print("Hi")

    """
    # Extract content between Begin and End Patch markers
    match = re.search(r"\*\*\* Begin Patch\s*(.*?)\s*\*\*\* End Patch", content, re.DOTALL)
    if not match:
        return {}

    patch_body = match.group(1).strip()
    sections = re.split(r"^\*\*\* Update File:\s*", patch_body, flags=re.MULTILINE)
    file_to_diff = {}

    # Skip anything before the first "*** Update File:" (e.g., stray lines)
    for section in sections[1:]:
        lines = section.strip().splitlines()
        if not lines:
            continue

        file_path = lines[0].strip()
        diff_content = "\n".join(lines[1:]).strip()

        if diff_content:
            file_to_diff[file_path] = f"*** Update File: {file_path}\n{diff_content}"

    return file_to_diff


def apply_update(text: str, action: PatchAction, path: str) -> str:
    """Applies UPDATE chunks to the given text content."""  # noqa: D401
    if action.chunks == []:
        return text

    orig_lines = text.splitlines()  # Use splitlines to handle endings consistently
    dest_lines: list[str] = []
    current_orig_line_idx = 0  # Tracks index in orig_lines processed so far

    # Sort chunks by their original index to apply them sequentially
    sorted_chunks = sorted(action.chunks, key=lambda c: c.orig_index)

    for chunk in sorted_chunks:
        # chunk.orig_index is the absolute line number where the change starts
        # (where the first deleted line was, or where inserted lines go if no deletes)
        chunk_start_index = chunk.orig_index

        if chunk_start_index < current_orig_line_idx:
            # This indicates overlapping chunks or incorrect indices from parsing
            msg = (
                f"{path}: Overlapping or out-of-order chunk detected."
                f" Current index {current_orig_line_idx}, chunk starts at {chunk_start_index}."
            )
            raise DiffError(msg)

        # Add lines from original file between the last chunk and this one
        dest_lines.extend(orig_lines[current_orig_line_idx:chunk_start_index])

        # Verify that the lines to be deleted actually match the original file content
        # (The parser should have used find_context, but double-check here)
        num_del = len(chunk.del_lines)
        actual_deleted_lines = orig_lines[chunk_start_index : chunk_start_index + num_del]

        # Use the same normalization as find_context_core for comparison robustness
        norm_chunk_del = [_norm(s).strip() for s in chunk.del_lines]
        norm_actual_del = [_norm(s).strip() for s in actual_deleted_lines]

        if norm_chunk_del != norm_actual_del:
            # This indicates the context matching failed or the file changed since parsing
            # Provide detailed error message
            expected_str = "\n".join(f"- {s}" for s in chunk.del_lines)
            actual_str = "\n".join(f"  {s}" for s in actual_deleted_lines)
            msg = (
                f"{path}: Mismatch applying patch near line {chunk_start_index + 1}.\n"
                f"Expected lines to remove:\n{expected_str}\n"
                f"Found lines in file:\n{actual_str}"
            )
            raise DiffError(msg)

        # Add the inserted lines from the chunk
        dest_lines.extend(chunk.ins_lines)

        # Advance the original line index past the lines processed (deleted lines)
        current_orig_line_idx = chunk_start_index + num_del

    # Add any remaining lines from the original file after the last chunk
    dest_lines.extend(orig_lines[current_orig_line_idx:])

    # Join lines and ensure a single trailing newline
    result = "\n".join(dest_lines)
    if result or orig_lines:  # Add newline unless result is empty and original was empty
        result += "\n"
    return result


class V4ADiff(Diff):
    def run(self) -> dict[str, str]:
        file_to_diff = extract_update_sections(self.content)
        file_to_original_code = self.source_code
        new_file_to_code: dict[str, str] = {}
        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:
                patch = parse_patch_text(diff.strip().splitlines(), file_to_original_code)
                actions = patch.actions[file_name]
                new_file_to_code[file_name] = apply_update(original_content, actions, file_name)

        return new_file_to_code
