#!/usr/bin/env python3
import argparse
import os
import re
import sys
import shutil

try:
    import pyperclip
except ImportError:
    pyperclip = None


PATCH_BEGIN = "*** Begin Patch"
PATCH_END = "*** End Patch"
PATCH_UPDATE_FILE = "*** Update File:"


def read_patch_input(source: str = None):
    """Read patch from file, stdin, clipboard, or prompt."""
    if source:
        with open(source, "r", encoding="utf-8") as f:
            data = f.read()
    elif not sys.stdin.isatty():
        data = sys.stdin.read()
    else:
        clip = pyperclip.paste().strip() if pyperclip else ""
        if clip:
            print("Using patch from clipboard.")
            data = clip
        else:
            if pyperclip:
                print("No valid patch detected in clipboard.")
            print("Please paste your patch below, then press Ctrl-D (Linux/macOS) or Ctrl-Z (Windows) when done:\n")
            try:
                data = sys.stdin.read()
            except KeyboardInterrupt:
                sys.exit("\nAborted.")
    return data.strip()


def detect_patch_format(data: str):
    """Return (has_header, has_update_line) flags."""
    has_header = PATCH_BEGIN in data and PATCH_END in data
    has_update_line = PATCH_UPDATE_FILE in data
    return has_header, has_update_line


def parse_patch(data: str, filename_override=None):
    """Extract the target filename and list of hunks."""
    match = re.search(r"\*\*\* Update File:\s*(.+)", data)
    filename = match.group(1).strip() if match else filename_override

    if not filename:
        filename = input("No target file specified. Please enter the filename to patch: ").strip()
        if not filename:
            sys.exit("Error: Target filename required.")

    # Extract hunks (between @@ markers)
    hunks = []
    current_hunk = []
    in_hunk = False
    for line in data.splitlines():
        if line.startswith("@@"):
            if current_hunk:
                hunks.append(current_hunk)
            current_hunk = []
            in_hunk = True
        elif in_hunk:
            current_hunk.append(line)
    if current_hunk:
        hunks.append(current_hunk)
    return filename, hunks


def make_backup(path: str):
    """Create sequential .bak backups if needed."""
    base_backup = path + ".bak"
    if not os.path.exists(base_backup):
        shutil.copy2(path, base_backup)
        return base_backup

    i = 1
    while True:
        backup_name = f"{base_backup}{i}"
        if not os.path.exists(backup_name):
            shutil.copy2(path, backup_name)
            return backup_name
        i += 1


def find_and_replace_block(content, old_block, new_block, context_before=None, context_after=None, filename=None):
    """Find a full old_block with optional context and replace with new_block."""
    pattern_parts = []
    if context_before:
        pattern_parts.append(re.escape(context_before))
    pattern_parts.append(re.escape(old_block))
    if context_after:
        pattern_parts.append(re.escape(context_after))

    pattern = "(?s)" + ".*?".join(pattern_parts)
    match = re.search(pattern, content)
    if not match:
        # Try looser match
        match = re.search(re.escape(old_block), content)
        if not match:
            sys.exit(
                f"Error: Could not find patch target block in '{filename}' containing:\n{old_block.strip()[:200]}"
            )

    start, end = match.span()
    return content[:start] + content[start:end].replace(old_block, new_block, 1) + content[end:]


def apply_hunk_to_content(content: str, hunk, filename):
    """Apply a hunk (multi-line context aware)."""
    minus_lines, plus_lines, context_lines = [], [], []

    for line in hunk:
        if line.startswith("-"):
            minus_lines.append(line[1:])
        elif line.startswith("+"):
            plus_lines.append(line[1:])
        elif line.startswith(" "):
            context_lines.append(line[1:])

    old_block = "\n".join(minus_lines)
    new_block = "\n".join(plus_lines)
    context_before = "\n".join(context_lines[:3]) if context_lines else None
    context_after = "\n".join(context_lines[-3:]) if context_lines else None

    return find_and_replace_block(content, old_block, new_block, context_before, context_after, filename)


def apply_patch_to_file(filename, hunks):
    """Apply patch hunks to file contents."""
    with open(filename, "r", encoding="utf-8") as f:
        content = f.read()

    for hunk in hunks:
        content = apply_hunk_to_content(content, hunk, filename)

    return content


def main():
    parser = argparse.ArgumentParser(description="Apply ChatGPT or raw unified diffs to files.")
    parser.add_argument("patchfile", nargs="?", help="File containing the patch (optional)")
    parser.add_argument("-i", "--in-place", action="store_true", help="Replace original file (create backup)")
    parser.add_argument("-t", "--target", help="Apply patch to this file instead of the one mentioned in patch")

    args = parser.parse_args()

    data = read_patch_input(args.patchfile)
    has_header, has_update_line = detect_patch_format(data)

    # If header present but update line missing, prompt for file
    patch_file, hunks = parse_patch(data, filename_override=args.target)

    if not os.path.exists(patch_file):
        sys.exit(f"Error: Target file '{patch_file}' not found.")

    result = apply_patch_to_file(patch_file, hunks)

    if args.in_place:
        backup = make_backup(patch_file)
        with open(patch_file, "w", encoding="utf-8") as f:
            f.write(result)
        print(f"Patched {patch_file} (backup: {backup})")
    else:
        sys.stdout.write(result)


if __name__ == "__main__":
    main()