import re
import shutil
import tempfile
from pathlib import Path

import bpy

IMAGE_NAME_SUFFIXES = (
    ".jpeg",
    ".jpg",
    ".png",
    ".webp",
    ".tiff",
    ".tif",
    ".exr",
    ".hdr",
    ".bmp",
    ".tga",
)


VIDEO_NAME_SUFFIXES = (
    ".mp4",
    ".mov",
    ".avi",
    ".mkv",
    ".webm",
    ".m4v",
    ".mts",
    ".mpg",
    ".mpeg",
)


def is_image_empty(candidate):
    return bool(
        candidate
        and getattr(candidate, "type", None) == "EMPTY"
        and getattr(candidate, "empty_display_type", None) == "IMAGE"
        and getattr(candidate, "data", None) is not None
    )


def require_image_empty(object_name):
    """Return a named Image Empty or raise when it is no longer available."""
    source_object = bpy.data.objects.get(object_name)
    if not is_image_empty(source_object):
        raise RuntimeError("The source Image Empty is no longer available")
    return source_object


def image_pixels(image):
    """Read one Blender Image into a flat float32 pixel array."""
    import numpy as np

    pixels = np.empty(len(image.pixels), dtype=np.float32)
    image.pixels.foreach_get(pixels)
    return pixels


def image_rgba(image):
    """Read one Blender Image into a top-down RGBA float32 array."""
    import numpy as np

    width, height = (int(value) for value in image.size)
    if width <= 0 or height <= 0:
        raise ValueError("The source image has no readable dimensions")
    return np.flipud(image_pixels(image).reshape((height, width, 4)))



def is_animated_image(image):
    return getattr(image, "source", None) in {"MOVIE", "SEQUENCE"}


def video_total_frames(image):
    """Return the total animated frame count, or 1 for a still image."""
    source = getattr(image, "source", None)
    if source == "MOVIE":
        return int(getattr(image, "frame_duration", 1) or 1)
    if source == "SEQUENCE":
        return len(collect_sequence_frames(image)) or 1
    return 1


def image_user_settings(source_object):
    """Return the Empty's image user settings as a four-item tuple."""
    image_user = getattr(source_object, "image_user", None)
    if image_user is None:
        return 1, 0, 0, False
    return (
        int(getattr(image_user, "frame_start", 1) or 1),
        int(getattr(image_user, "frame_offset", 0) or 0),
        int(getattr(image_user, "frame_duration", 0) or 0),
        bool(getattr(image_user, "use_cyclic", False)),
    )


def extract_frame_range(source_object):
    """Return the clamped (start_frame, frame_count) range extracted from the source."""
    total_frames = video_total_frames(getattr(source_object, "data", None))
    if total_frames <= 1:
        return 1, total_frames
    _frame_start, frame_offset, image_frames, _use_cyclic = (
        image_user_settings(source_object)
    )
    frame_count = (
        min(image_frames, total_frames)
        if image_frames >= 1
        else total_frames
    )
    start_frame = min(1 + frame_offset, total_frames - frame_count + 1)
    start_frame = max(start_frame, 1)
    return start_frame, frame_count


def default_extract_frames(source_object):
    """Return the Empty's Image Frames value, or the full frame count when unset."""
    return extract_frame_range(source_object)[1]


def source_start_frame(source_object_name):
    source_object = bpy.data.objects.get(source_object_name)
    return extract_frame_range(source_object)[0] if source_object else 1


def image_path(image):
    filepath_from_user = getattr(image, "filepath_from_user", None)
    filepath = (
        filepath_from_user()
        if callable(filepath_from_user)
        else getattr(image, "filepath", "")
    )
    if not filepath:
        return Path()
    return Path(bpy.path.abspath(filepath)).expanduser()


def image_is_packed(image):
    if getattr(image, "packed_file", None) is not None:
        return True
    return bool(getattr(image, "packed_files", ()))


def cleanup_image_input(path, temporary):
    """Remove a prepared input and its temporary container when requested."""
    if not temporary or not path:
        return
    path = Path(path)
    try:
        if path.is_dir():
            shutil.rmtree(path, ignore_errors=True)
        else:
            path.unlink(missing_ok=True)
            path.parent.rmdir()
    except OSError:
        pass


def image_base_name(image):
    name = image.name
    lower_name = name.lower()
    for suffix in IMAGE_NAME_SUFFIXES + VIDEO_NAME_SUFFIXES:
        suffix_index = lower_name.rfind(suffix)
        if suffix_index <= 0:
            continue
        trailing = name[suffix_index + len(suffix) :]
        if trailing and not (
            trailing.startswith(".") and trailing[1:].isdigit()
        ):
            continue
        return f"{name[:suffix_index]}{trailing}"
    return name


def depth_data_name(source_object):
    return f"{image_base_name(source_object.data)}_depth.exr"


def normal_data_name(source_object):
    return f"{image_base_name(source_object.data)}_normal.png"


def color_data_name(source_object):
    return f"{image_base_name(source_object.data)}_color.png"


def save_packed_image(image, _scene):
    output_path = Path(tempfile.mkdtemp(prefix="anyimage-input-")) / "input.png"
    original_file_format = image.file_format
    try:
        image.file_format = "PNG"
        image.save(filepath=str(output_path), save_copy=True)
    except (OSError, RuntimeError) as error:
        output_path.unlink(missing_ok=True)
        raise RuntimeError("Unable to export the packed image") from error
    finally:
        image.file_format = original_file_format
    if not output_path.is_file():
        raise RuntimeError("Blender did not export the packed image")
    return output_path


def collect_sequence_frames(image):
    """Return the numbered frame files of a SEQUENCE image, or [] when absent."""
    source_path = image_path(image)
    if not is_animated_image(image) or not source_path.is_file():
        return []
    match = re.match(
        r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
        source_path.name,
    )
    if match is None:
        return []
    prefix = match.group("prefix")
    suffix = match.group("suffix").lower()
    frames = []
    for path in source_path.parent.iterdir():
        frame_match = re.match(
            r"^(?P<prefix>.*?)(?P<digits>\d+)(?P<suffix>\.[^.]+)$",
            path.name,
        )
        if (
            frame_match is not None
            and frame_match.group("prefix") == prefix
            and frame_match.group("suffix").lower() == suffix
        ):
            frames.append((int(frame_match.group("digits")), path))
    return [path for _digits, path in sorted(frames)]


def prepare_image_input(image, scene):
    source_path = image_path(image)
    if image_is_packed(image):
        return save_packed_image(image, scene), True
    if getattr(image, "source", None) == "MOVIE":
        if not source_path.is_file():
            raise RuntimeError(
                "The image empty must reference an existing video file"
            )
        return source_path, False
    sequence_frames = collect_sequence_frames(image)
    if sequence_frames:
        temporary_directory = Path(tempfile.mkdtemp(prefix="anyimage-input-"))
        try:
            for index, frame_path in enumerate(sequence_frames):
                suffix = frame_path.suffix.lower() or ".jpg"
                shutil.copy2(
                    frame_path,
                    temporary_directory / f"frame.{index + 1:04d}{suffix}",
                )
        except Exception:
            shutil.rmtree(temporary_directory, ignore_errors=True)
            raise
        return temporary_directory, True
    if not source_path.is_file():
        raise RuntimeError(
            "The image empty must reference an existing image file"
        )
    return source_path, False


def prepare_region_input(image, region_mask):
    """Write one cropped RGBA Region input owned by its temporary directory."""
    import numpy as np

    left, top, right, bottom = region_mask.bounds
    rgba = image_rgba(image)
    cropped = rgba[top:bottom, left:right].copy()
    if cropped.shape[:2] != region_mask.values.shape:
        raise ValueError("The Region does not match the source image")
    cropped[:, :, 3] *= region_mask.values
    cropped[cropped[:, :, 3] <= 0.0, :3] = 0.0
    output_directory = Path(tempfile.mkdtemp(prefix="anyimage-region-"))
    output_path = output_directory / "input.png"
    temporary_image = bpy.data.images.new(
        "AnyImage Region Input",
        width=right - left,
        height=bottom - top,
        alpha=True,
    )
    original_format = temporary_image.file_format
    try:
        temporary_image.colorspace_settings.name = image.colorspace_settings.name
        temporary_image.alpha_mode = "STRAIGHT"
        temporary_image.pixels.foreach_set(np.flipud(cropped).ravel())
        temporary_image.file_format = "PNG"
        temporary_image.save(filepath=str(output_path), save_copy=True)
    except (OSError, RuntimeError):
        cleanup_image_input(output_path, True)
        raise
    finally:
        temporary_image.file_format = original_format
        bpy.data.images.remove(temporary_image, do_unlink=True)
    if not output_path.is_file():
        cleanup_image_input(output_path, True)
        raise RuntimeError("Blender did not export the Region input")
    return output_path


def configure_image_user(
    image_user,
    frame_count=None,
    frame_start=1,
    frame_offset=0,
    use_cyclic=None,
):
    """Map the image to scene frames and enable auto refresh."""
    if image_user is None:
        return
    if hasattr(image_user, "frame_start"):
        image_user.frame_start = frame_start
    if hasattr(image_user, "frame_offset"):
        image_user.frame_offset = frame_offset
    if hasattr(image_user, "use_auto_refresh"):
        image_user.use_auto_refresh = True
    if use_cyclic is not None and hasattr(image_user, "use_cyclic"):
        image_user.use_cyclic = use_cyclic
    if frame_count is not None and hasattr(image_user, "frame_duration"):
        image_user.frame_duration = frame_count


def pack_sequence_if_possible(image, expected_frames):
    """Pack a frame sequence when Blender supports it; otherwise keep external files."""
    if expected_frames <= 1:
        return False
    try:
        image.pack()
    except (OSError, RuntimeError):
        return False
    if len(getattr(image, "packed_files", ())) >= expected_frames:
        return True
    try:
        image.unpack(method="REMOVE")
    except (OSError, RuntimeError):
        pass
    return False


def sorted_frame_paths(directory, prefix=None):
    """Return image files in a directory sorted by their trailing frame number."""
    def frame_key(path):
        match = re.search(r"(\d+)(?=\.[^.]+$)", path.name)
        return (int(match.group(1)) if match is not None else 0, path.name)

    return sorted(
        (
            path
            for path in Path(directory).iterdir()
            if path.is_file()
            and path.suffix.lower() in IMAGE_NAME_SUFFIXES
            and (prefix is None or path.stem.startswith(prefix))
        ),
        key=frame_key,
    )


def force_sequence_source(image):
    """Ensure a loaded image is a SEQUENCE datablock for its numbered frames."""
    if getattr(image, "source", None) == "SEQUENCE":
        return image
    try:
        image.source = "SEQUENCE"
    except (OSError, RuntimeError, TypeError) as error:
        raise RuntimeError(
            "Blender cannot load the frame sequence"
        ) from error
    if getattr(image, "source", None) != "SEQUENCE":
        raise RuntimeError("Blender cannot load the frame sequence")
    return image


def load_normal_result_image(path, source_object):
    """Load, name, configure, and pack a generated normal image."""
    image = bpy.data.images.load(str(path), check_existing=False)
    image.name = normal_data_name(source_object)
    image.colorspace_settings.name = "Non-Color"
    image.pack()
    return image


def create_pixel_result_image(source_image, pixels, output_size, suffix=""):
    """Create a packed Image datablock from edited RGBA pixels."""
    output_width, output_height = (int(value) for value in output_size)
    if output_width <= 0 or output_height <= 0:
        raise ValueError("The edited image dimensions must be positive")
    if len(pixels) != output_width * output_height * 4:
        raise ValueError("The edited pixels do not match the image dimensions")
    image = bpy.data.images.new(
        f"{image_base_name(source_image)}{suffix}",
        width=output_width,
        height=output_height,
        alpha=True,
    )
    try:
        image.colorspace_settings.name = source_image.colorspace_settings.name
        image.alpha_mode = "STRAIGHT"
        image.pixels.foreach_set(pixels)
        image.pack()
        image.update()
    except Exception:
        bpy.data.images.remove(image, do_unlink=True)
        raise
    return image


def replace_empty_image(source_object, result_image, *, placement_bounds=None):
    """Replace one Empty image, retaining shared source images."""
    source_image = source_object.data
    source_name = source_image.name
    original_fake_user = source_image.use_fake_user
    image_user = getattr(source_object, "image_user", None)
    frame_count = (
        video_total_frames(result_image)
        if is_animated_image(result_image)
        else None
    )
    alignment = (
        empty_display_alignment(
            source_object,
            tuple(source_image.size),
            tuple(result_image.size),
            placement_bounds,
        )
        if placement_bounds is not None
        else None
    )
    if frame_count is not None:
        (
            original_frame_start,
            _original_frame_offset,
            _original_frame_duration,
            original_use_cyclic,
        ) = image_user_settings(source_object)

    source_image.use_fake_user = True
    try:
        source_object.data = result_image
        if frame_count is not None:
            configure_image_user(
                image_user,
                frame_count,
                frame_start=original_frame_start,
                use_cyclic=original_use_cyclic,
            )
        if alignment is not None:
            source_object.empty_display_size = alignment[0]
            source_object.empty_image_offset = alignment[1]
    except Exception:
        source_object.data = source_image
        source_image.use_fake_user = original_fake_user
        raise

    source_image.use_fake_user = original_fake_user
    if not original_fake_user and int(source_image.users) == 0:
        bpy.data.images.remove(source_image, do_unlink=True)
        result_image.name = source_name
    return result_image



def load_color_result_image(
    source_object,
    output_directory,
    *,
    frame_prefix="image.",
):
    output_directory = Path(output_directory)
    color_paths = sorted_frame_paths(output_directory, frame_prefix)
    if not color_paths:
        raise RuntimeError(
            f"Color result has no images: {output_directory}"
        )
    source_image = source_object.data
    source_color_space = source_image.colorspace_settings.name
    animated = is_animated_image(source_image)
    try:
        result_image = bpy.data.images.load(
            str(color_paths[0]),
            check_existing=False,
        )
        result_image.alpha_mode = "CHANNEL_PACKED"
        result_image.colorspace_settings.name = source_color_space
        if animated:
            force_sequence_source(result_image)
    except (OSError, RuntimeError, TypeError, ValueError) as error:
        if "result_image" in locals():
            bpy.data.images.remove(result_image, do_unlink=True)
        raise RuntimeError("Unable to load the color result") from error

    try:
        result_name = color_data_name(source_object)
        if animated:
            pack_sequence_if_possible(result_image, len(color_paths))
        else:
            result_image.pack()
        result_image.name = result_name
    except (OSError, RuntimeError, TypeError, ValueError) as error:
        bpy.data.images.remove(result_image, do_unlink=True)
        raise RuntimeError("Unable to prepare the color result") from error
    return result_image


def image_empty_bounds(source_object):
    image_width, image_height = source_object.data.size
    if image_width <= 0 or image_height <= 0:
        raise ValueError("The source image has no readable dimensions")

    aspect = image_width / image_height
    display_size = float(source_object.empty_display_size)
    if aspect >= 1.0:
        width = display_size
        height = display_size / aspect
    else:
        width = display_size * aspect
        height = display_size

    offset = source_object.empty_image_offset
    x_min = float(offset[0]) * width
    y_min = float(offset[1]) * height
    return (
        x_min,
        x_min + width,
        y_min,
        y_min + height,
    )


def empty_display_alignment(
    source_object,
    source_size,
    output_size,
    placement_bounds,
):
    """Return the display size and offset for an edited Image Empty."""
    source_width, source_height = source_size
    output_width, output_height = output_size
    left, top, right, bottom = placement_bounds
    if (
        source_width <= 0
        or source_height <= 0
        or output_width <= 0
        or output_height <= 0
        or right <= left
        or bottom <= top
    ):
        raise ValueError("The edited image placement is invalid")
    x_min, x_max, y_min, y_max = image_empty_bounds(source_object)
    local_width = x_max - x_min
    local_height = y_max - y_min
    desired_x_min = x_min + left * local_width / source_width
    desired_x_max = x_min + right * local_width / source_width
    desired_y_max = y_max - top * local_height / source_height
    desired_y_min = y_max - bottom * local_height / source_height
    desired_width = desired_x_max - desired_x_min
    desired_height = desired_y_max - desired_y_min
    display_size = max(desired_width, desired_height)
    aspect = output_width / output_height
    if aspect >= 1.0:
        display_width = display_size
        display_height = display_size / aspect
    else:
        display_width = display_size * aspect
        display_height = display_size
    center_x = (desired_x_min + desired_x_max) * 0.5
    center_y = (desired_y_min + desired_y_max) * 0.5
    return display_size, (
        center_x / display_width - 0.5,
        center_y / display_height - 0.5,
    )