import shutil
import tempfile
import time
from pathlib import Path


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_SUFFIXES = {
    ".mp4",
    ".mov",
    ".avi",
    ".mkv",
    ".webm",
    ".m4v",
    ".mts",
    ".mpg",
    ".mpeg",
}
SCENE_FILENAME = "scene.glb"


class TaskCancelled(RuntimeError):
    pass


def check_cancelled(cancel_check):
    if cancel_check is not None and cancel_check():
        raise TaskCancelled("Task cancelled")


def collect_images(directory):
    directory = Path(directory)
    return sorted(
        path
        for path in directory.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    )


def first_frames(images, maximum_frames, start_frame=1):
    """Return up to maximum_frames image paths in order from start_frame."""
    if maximum_frames <= 0:
        return images
    start_index = max(start_frame - 1, 0)
    return images[start_index:start_index + maximum_frames]


def extract_video_frames(
    video_path,
    output_directory,
    maximum_frames,
    av,
    start_frame=1,
    cancel_check=None,
):
    output_directory = Path(output_directory)
    output_directory.mkdir(parents=True, exist_ok=True)
    try:
        container = av.open(str(video_path))
    except Exception as error:
        raise RuntimeError(f"Unable to open video: {video_path}") from error
    image_paths = []
    skipped = 0
    try:
        stream = container.streams.video[0]
        for frame in container.decode(stream):
            check_cancelled(cancel_check)
            if skipped < start_frame - 1:
                skipped += 1
                continue
            image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
            frame.to_image().save(image_path, format="JPEG")
            image_paths.append(image_path)
            if len(image_paths) >= maximum_frames:
                break
    finally:
        container.close()

    if not image_paths:
        raise RuntimeError(f"Unable to read any video frame from: {video_path}")
    return image_paths


def prepare_input(
    input_path,
    video_frames,
    temporary_root,
    start_frame=1,
    cancel_check=None,
):
    check_cancelled(cancel_check)
    input_path = Path(input_path)
    if input_path.is_dir():
        images = collect_images(input_path)
        if not images:
            raise RuntimeError(f"No image frames found in: {input_path}")
        return first_frames(images, video_frames, start_frame), None
    suffix = input_path.suffix.lower()
    if suffix in IMAGE_SUFFIXES:
        return [input_path], None
    if suffix not in VIDEO_SUFFIXES:
        supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
        raise RuntimeError(f"Unsupported input type; expected one of: {supported}")

    import av

    temporary_directory = Path(
        tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
    )
    try:
        images = extract_video_frames(
            input_path,
            temporary_directory,
            video_frames,
            av,
            start_frame=start_frame,
            cancel_check=cancel_check,
        )
    except Exception:
        shutil.rmtree(temporary_directory, ignore_errors=True)
        raise
    return images, temporary_directory


def prepare_background_input(
    input_path,
    video_frames,
    temporary_root,
    start_frame=1,
    cancel_check=None,
):
    """Resolve a BEN2 input to frame images; returns (images, temporary_directory)."""
    input_path = Path(input_path)
    if input_path.is_dir():
        images = collect_images(input_path)
        if not images:
            raise RuntimeError(f"No image frames found in: {input_path}")
        return first_frames(images, video_frames, start_frame), None
    if input_path.suffix.lower() in IMAGE_SUFFIXES:
        return [input_path], None
    return prepare_input(
        input_path,
        video_frames,
        temporary_root,
        start_frame=start_frame,
        cancel_check=cancel_check,
    )


def copy_frames(images, output_directory):
    """Copy a multi-frame input as numbered frame files for the Blender sequence."""
    if len(images) <= 1:
        return
    for index, image_path in enumerate(images):
        suffix = Path(image_path).suffix.lower() or ".jpg"
        destination = (
            Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
        )
        if Path(image_path).resolve() == destination.resolve():
            continue
        shutil.copy2(image_path, destination)


def extract_frames(context, parameters):
    """Copy selected source frames into the current job directory."""
    context.directory.mkdir(parents=True, exist_ok=True)
    images, temporary_directory = prepare_input(
        parameters["input"],
        parameters.get("video_frames", 16),
        context.directory,
        start_frame=parameters.get("start_frame", 1),
        cancel_check=context.check_cancelled,
    )
    try:
        copy_frames(images, context.directory)
    finally:
        if temporary_directory is not None:
            shutil.rmtree(temporary_directory, ignore_errors=True)
    message = f"Extracted {len(images)} frame(s)"
    context.progress(1.0, message)
    return {"frames": len(images)}


def exported_paths(output_directory, export_format):
    output_directory = Path(output_directory)
    if export_format == "glb":
        scene_path = output_directory / SCENE_FILENAME
        return [scene_path] if scene_path.is_file() else []
    if export_format == "gs_ply":
        return sorted((output_directory / "gs_ply").glob("*.ply"))
    if export_format == "depth_vis":
        return sorted((output_directory / "depth_vis").glob("*.jpg"))
    if export_format == "mesh":
        result_path = output_directory / "mesh.glb"
        return [result_path] if result_path.is_file() else []
    if export_format == "mesh_data":
        result_path = output_directory / "mesh.npz"
        return [result_path] if result_path.is_file() else []
    if export_format == "mini_npz":
        result_path = output_directory / "results.npz"
        if result_path.is_file():
            return [result_path]
        result_path = (
            output_directory / "exports" / "mini_npz" / "results.npz"
        )
        return [result_path] if result_path.is_file() else []
    return []


def npz_is_complete(path):
    """Return True only when the npz archive can be fully read."""
    try:
        import numpy as np

        with np.load(path, allow_pickle=False) as archive:
            return "depth" in archive.files
    except Exception:
        return False


def wait_for_exported_paths(
    output_directory,
    export_format,
    timeout=30.0,
    cancel_check=None,
):
    # DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整，
    # 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        check_cancelled(cancel_check)
        result_paths = exported_paths(output_directory, export_format)
        if result_paths:
            if (
                export_format in ("mini_npz", "npz")
                and not npz_is_complete(result_paths[0])
            ):
                time.sleep(0.1)
                continue
            return result_paths
        time.sleep(0.05)
    return []


def move_exported_file(
    source,
    destination,
    timeout=10.0,
    cancel_check=None,
):
    """Move an async export after Windows releases its file handle."""
    source = Path(source)
    destination = Path(destination)
    deadline = time.monotonic() + timeout
    while True:
        check_cancelled(cancel_check)
        try:
            source.replace(destination)
            return destination
        except OSError as error:
            is_file_lock = isinstance(error, PermissionError) or getattr(
                error,
                "winerror",
                None,
            ) in {32, 33}
            if not is_file_lock or time.monotonic() >= deadline:
                raise
            time.sleep(0.1)


def flatten_exported_paths(
    output_directory,
    result_paths,
    cancel_check=None,
):
    """Move files produced in exporter subdirectories to the job root."""
    output_directory = Path(output_directory)
    flattened_paths = []
    for source in result_paths:
        source = Path(source)
        check_cancelled(cancel_check)
        if source.parent == output_directory:
            flattened_paths.append(source)
            continue
        destination = output_directory / source.name
        if destination.exists():
            raise RuntimeError(f"Export filename already exists: {destination.name}")
        flattened_paths.append(
            move_exported_file(
                source,
                destination,
                cancel_check=cancel_check,
            )
        )
    for directory in sorted(
        {path.parent for path in map(Path, result_paths)},
        key=lambda path: len(path.parts),
        reverse=True,
    ):
        while directory != output_directory:
            try:
                directory.rmdir()
            except OSError:
                break
            directory = directory.parent
    return flattened_paths