from __future__ import annotations

import re

import attrs

from codeflash_api.diff._base import Diff, DiffError

UPDATE_FILE_PREFIX = "*** Update File: "


@attrs.define
class Chunk:
    """
    A single add/delete change at a specific line index.
    """

    orig_index: int = -1
    del_lines: list[str] = attrs.Factory(list)
    ins_lines: list[str] = attrs.Factory(list)


@attrs.define
class PatchAction:
    """
    All chunks for one file.
    """

    path: str = ""
    chunks: list[Chunk] = attrs.Factory(list)


@attrs.define
class Patch:
    """
    Parsed patch with actions per file.
    """

    actions: dict[str, PatchAction] = attrs.Factory(dict)
    fuzz: int = 0


def _norm(line: str) -> str:
    """
    Strip CR for LF/CRLF compatibility.
    """
    return line.replace("\r", "")


def find_context_core(
    lines: list[str],
    context: list[str],
    start: int,
) -> tuple[int, int]:
    """
    Find *context* in *lines* starting at *start* with 3-tier fuzz.
    """
    ctx_len = len(context)
    end = len(lines) - ctx_len + 1

    for i in range(start, end):
        if lines[i : i + ctx_len] == context:
            return (i, 0)

    rstrip_ctx = [s.rstrip() for s in context]
    for i in range(start, end):
        if [s.rstrip() for s in lines[i : i + ctx_len]] == rstrip_ctx:
            return (i, 1)

    strip_ctx = [s.strip() for s in context]
    for i in range(start, end):
        if [s.strip() for s in lines[i : i + ctx_len]] == strip_ctx:
            return (i, 100)

    return (-1, 0)


def find_context(
    lines: list[str],
    context: list[str],
    start: int,
    *,
    eof: bool,
) -> tuple[int, int]:
    """
    Find *context* with special EOF marker handling.
    """
    if eof:
        idx, fuzz = find_context_core(
            lines, context, len(lines) - len(context)
        )
        if idx >= 0:
            return (idx, fuzz)
        idx, fuzz = find_context_core(lines, context, start)
        if idx >= 0:
            return (idx, fuzz + 10_000)
        return (-1, 0)
    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 context/change section from an UPDATE block.
    """
    context_lines: list[str] = []
    chunks: list[Chunk] = []
    del_lines: list[str] = []
    ins_lines: list[str] = []
    mode = "keep"
    is_eof = False

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

        if line.startswith("@@"):
            break
        if line.startswith("*** "):
            tag = line.strip()
            if tag == "*** End of File":
                is_eof = True
                i += 1
            break

        if line.startswith(" ") or line == "":
            if mode in ("add", "delete"):
                chunks.append(
                    Chunk(
                        orig_index=len(context_lines) - len(del_lines),
                        del_lines=del_lines,
                        ins_lines=ins_lines,
                    )
                )
                del_lines = []
                ins_lines = []
                mode = "keep"

            ctx = line[1:] if line.startswith(" ") else ""
            context_lines.append(ctx)
            i += 1

        elif line.startswith("-"):
            if mode == "add":
                chunks.append(
                    Chunk(
                        orig_index=len(context_lines) - len(del_lines),
                        del_lines=del_lines,
                        ins_lines=ins_lines,
                    )
                )
                del_lines = []
                ins_lines = []
            mode = "delete"
            stripped = line[1:]
            del_lines.append(stripped)
            context_lines.append(stripped)
            i += 1

        elif line.startswith("+"):
            mode = "add"
            ins_lines.append(line[1:])
            i += 1

        else:
            msg = f"Invalid patch line found: {line!r}"
            raise DiffError(msg)

    if mode in ("add", "delete"):
        chunks.append(
            Chunk(
                orig_index=len(context_lines) - len(del_lines),
                del_lines=del_lines,
                ins_lines=ins_lines,
            )
        )

    return (context_lines, chunks, i, is_eof)


def _parse_update_file_sections(
    lines: list[str],
    index: int,
    file_content: str,
) -> tuple[PatchAction, int, int]:
    """
    Parse all @@ scope markers and sections for one file.
    """
    action = PatchAction()
    total_fuzz = 0
    file_lines = [_norm(ln) for ln in file_content.splitlines()]
    current_file_idx = 0

    i = index
    while i < len(lines):
        line = _norm(lines[i])
        if not line.strip():
            i += 1
            continue
        if line.startswith("*** ") and not line.startswith("*** End of File"):
            break

        if line.startswith("@@"):
            scope_lines: list[str] = []
            while i < len(lines) and _norm(lines[i]).startswith("@@"):
                scope_lines.append(_norm(lines[i]).strip().strip("@").strip())
                i += 1

            found = False
            for scope in scope_lines:
                for j in range(current_file_idx, len(file_lines)):
                    if file_lines[j].strip() == scope:
                        current_file_idx = j
                        found = True
                        break
                if found:
                    break

            if not found:
                for scope in scope_lines:
                    for j in range(current_file_idx, len(file_lines)):
                        if file_lines[j].strip() == scope.strip():
                            current_file_idx = j
                            total_fuzz += 1
                            found = True
                            break
                    if found:
                        break

            if not found:
                msg = f"Could not find scope: {scope_lines!r}"
                raise DiffError(msg)
            continue

        context_block, chunks, next_i, is_eof = peek_next_section(lines, i)
        i = next_i

        if not context_block and not chunks:
            continue

        if context_block:
            found_idx, fuzz = find_context(
                file_lines,
                context_block,
                current_file_idx,
                eof=is_eof,
            )
            if found_idx < 0:
                msg = (
                    "Could not find context block"
                    f" starting with: {context_block[:3]!r}"
                )
                raise DiffError(msg)
            total_fuzz += fuzz
            for chunk in chunks:
                chunk.orig_index += found_idx
            current_file_idx = found_idx + len(context_block)
        else:
            for chunk in chunks:
                chunk.orig_index += current_file_idx

        action.chunks.extend(chunks)

    return (action, i, total_fuzz)


def parse_patch_text(
    lines: list[str],
    current_files: dict[str, str],
) -> Patch:
    """
    Parse full patch content into a Patch object.
    """
    patch = Patch()
    i = 0
    while i < len(lines):
        line = _norm(lines[i])
        if not line.strip():
            i += 1
            continue

        norm_line = line.strip()
        if norm_line.startswith(UPDATE_FILE_PREFIX.strip()):
            path = norm_line[len(UPDATE_FILE_PREFIX.strip()) :].strip()
            if path not in current_files:
                msg = f"File not found: {path!r}"
                raise DiffError(msg)
            i += 1
            action, i, fuzz = _parse_update_file_sections(
                lines, i, current_files[path]
            )
            action.path = path
            if path in patch.actions:
                patch.actions[path].chunks.extend(action.chunks)
            else:
                patch.actions[path] = action
            patch.fuzz += fuzz
        else:
            msg = f"Unknown or misplaced line: {line!r}"
            raise DiffError(msg)

    return patch


_PATCH_BODY_RE = re.compile(
    r"\*\*\* Begin Patch\s*(.*?)\s*\*\*\* End Patch",
    re.DOTALL,
)
_UPDATE_SPLIT_RE = re.compile(r"^\*\*\* Update File:\s*", re.MULTILINE)


def extract_update_sections(
    content: str,
) -> dict[str, str]:
    """
    Extract UPDATE sections from patch content.
    """
    match = _PATCH_BODY_RE.search(content)
    if not match:
        return {}

    body = match.group(1)
    parts = _UPDATE_SPLIT_RE.split(body)

    result: dict[str, str] = {}
    for part in parts[1:]:
        if not part.strip():
            continue
        first_nl = part.find("\n")
        if first_nl < 0:
            continue
        file_path = part[:first_nl].strip()
        diff_content = part[first_nl + 1 :]
        result[file_path] = f"{UPDATE_FILE_PREFIX}{file_path}\n{diff_content}"

    return result


def apply_update(
    text: str,
    action: PatchAction,
    path: str,
) -> str:
    """
    Apply UPDATE chunks to file content.
    """
    orig_lines = text.splitlines()
    output: list[str] = []
    current = 0

    sorted_chunks = sorted(action.chunks, key=lambda c: c.orig_index)
    for chunk in sorted_chunks:
        if chunk.orig_index < current:
            msg = f"Overlapping chunks in {path} at line {chunk.orig_index}"
            raise DiffError(msg)

        output.extend(orig_lines[current : chunk.orig_index])

        expected = [_norm(s).strip() for s in chunk.del_lines]
        actual = [
            _norm(s).strip()
            for s in orig_lines[
                chunk.orig_index : chunk.orig_index + len(chunk.del_lines)
            ]
        ]
        if expected != actual:
            msg = (
                f"Deleted lines mismatch in {path}"
                f" at line {chunk.orig_index}:"
                f" expected {expected!r},"
                f" got {actual!r}"
            )
            raise DiffError(msg)

        output.extend(chunk.ins_lines)
        current = chunk.orig_index + len(chunk.del_lines)

    output.extend(orig_lines[current:])

    result = "\n".join(output)
    if result or text:
        result += "\n"
    return result


class V4ADiff(Diff):
    """
    Apply V4A unified diff patches.
    """

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

        for path, code in self.source_code.items():
            if path in sections:
                diff_text = sections[path].strip()
                patch = parse_patch_text(
                    diff_text.splitlines(),
                    {path: code},
                )
                if path in patch.actions:
                    result[path] = apply_update(
                        code, patch.actions[path], path
                    )
                else:
                    result[path] = code
            else:
                result[path] = code

        return result
