clip_manager.py

from __future__ import annotations

import argparse
import glob
import logging
import os
import shutil
import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable

# Enable OpenEXR support in OpenCV — needed for EXR I/O throughout the pipeline.
# Must be set before any cv2.imread/imwrite calls on .exr files.
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

import cv2
import numpy as np

from backend.frame_io import EXR_WRITE_FLAGS, read_image_frame
from device_utils import resolve_device

if TYPE_CHECKING:
    from gvm_core import GVMProcessor
from BiRefNetModule.wrapper import BiRefNetHandler, usage_to_weights_file

logger = logging.getLogger(__name__)


@dataclass
class InferenceSettings:
    """Settings for CorridorKey inference, extracted from interactive prompts.

    Can be constructed directly for non-interactive use (Nuke, Houdini, batch scripts).
    """

    input_is_linear: bool = False
    despill_strength: float = 0.5  # 0.0–1.0
    auto_despeckle: bool = True
    despeckle_size: int = 400
    refiner_scale: float = 1.0


# Core Paths
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CLIPS_DIR = os.path.join(BASE_DIR, "ClipsForInference")
OUTPUT_DIR = os.path.join(BASE_DIR, "Output")

# Network Mapping
# Windows Drive -> Linux Mount Point
WIN_DRIVE_ROOT = "V:\\"
LINUX_MOUNT_ROOT = "/mnt/ssd-storage"


# --- Helpers ---
def is_image_file(filename: str) -> bool:
    return filename.lower().endswith((".png", ".jpg", ".jpeg", ".exr", ".tif", ".tiff", ".bmp"))


def is_video_file(filename: str) -> bool:
    return filename.lower().endswith((".mp4", ".mov", ".avi", ".mkv"))


def map_path(win_path: str) -> str:
    r"""
    Converts a Windows path (example: V:\Projects\Shot1) to the local Linux path.
    """
    # Normalize slashes
    win_path = win_path.strip()

    # Check if it starts with the drive letter
    if win_path.upper().startswith(WIN_DRIVE_ROOT.upper()):
        # Remove drive letter
        rel_path = win_path[len(WIN_DRIVE_ROOT) :]
        # Combine and flip slashes
        linux_path = os.path.join(LINUX_MOUNT_ROOT, rel_path).replace("\\", "/")
        return linux_path

    # If not on V:, maybe it's already a linux path or invalid?
    return win_path


# --- Classes ---
class ClipAsset:
    def __init__(self, path: str, asset_type: str) -> None:
        self.path = path
        self.type = asset_type  # 'sequence' or 'video'
        self.frame_count = 0
        self._calculate_length()

    def _calculate_length(self) -> None:
        if self.type == "sequence":
            files = sorted([f for f in os.listdir(self.path) if is_image_file(f)])
            self.frame_count = len(files)
        elif self.type == "video":
            cap = cv2.VideoCapture(self.path)
            if not cap.isOpened():
                self.frame_count = 0
            else:
                self.frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            cap.release()


class ClipEntry:
    def __init__(self, name: str, root_path: str) -> None:
        self.name = name
        self.root_path = root_path
        self.input_asset: ClipAsset | None = None
        self.alpha_asset: ClipAsset | None = None

    def find_assets(self) -> None:
        # 1. Look for Input
        input_dir = os.path.join(self.root_path, "Input")

        # Check for directory first
        if os.path.isdir(input_dir):
            if not os.listdir(input_dir):
                raise ValueError(f"Clip '{self.name}': 'Input' directory is empty.")
            self.input_asset = ClipAsset(input_dir, "sequence")
        else:
            # Check for video file (Case-Insensitive)
            candidates = glob.glob(os.path.join(self.root_path, "[Ii]nput.*"))
            candidates = [c for c in candidates if is_video_file(c)]

            if candidates:
                self.input_asset = ClipAsset(candidates[0], "video")
            else:
                # Fallback: Look for ANY video file in the directory
                all_files = glob.glob(os.path.join(self.root_path, "*"))
                video_files = [f for f in all_files if is_video_file(f)]

                if video_files:
                    logger.info(f"Clip '{self.name}': Using '{os.path.basename(video_files[0])}' as Input.")
                    self.input_asset = ClipAsset(video_files[0], "video")
                else:
                    raise ValueError(f"Clip '{self.name}': No 'Input' directory or video file found.")

        if self.input_asset.frame_count == 0:
            raise ValueError(f"Clip '{self.name}': Input asset has 0 frames or could not be read.")

        # 2. Look for Alpha
        # Check for 'AlphaHint' or 'alphahint' directory
        alpha_dir_upper = os.path.join(self.root_path, "AlphaHint")
        alpha_dir_lower = os.path.join(self.root_path, "alphahint")

        target_alpha_dir = None
        if os.path.isdir(alpha_dir_upper):
            target_alpha_dir = alpha_dir_upper
        elif os.path.isdir(alpha_dir_lower):
            target_alpha_dir = alpha_dir_lower

        if target_alpha_dir:
            if not os.listdir(target_alpha_dir):
                logger.warning(f"Clip '{self.name}': AlphaHint directory exists but is empty. Marking for generation.")
                self.alpha_asset = None
            else:
                # Check for image sequence first
                self.alpha_asset = ClipAsset(target_alpha_dir, "sequence")
                if self.alpha_asset.frame_count == 0:
                    # Fallback: check for video file inside the AlphaHint directory
                    video_candidates = [f for f in os.listdir(target_alpha_dir) if is_video_file(f)]
                    if video_candidates:
                        self.alpha_asset = ClipAsset(os.path.join(target_alpha_dir, video_candidates[0]), "video")
                    else:
                        logger.warning(
                            f"Clip '{self.name}': AlphaHint directory has no valid image or video files."
                            " Marking for generation."
                        )
                        self.alpha_asset = None
        else:
            # Check for video file (Case-Insensitive)
            # Match AlphaHint.* or alphahint.*
            candidates = glob.glob(os.path.join(self.root_path, "[Aa]lpha[Hh]int.*"))
            candidates = [c for c in candidates if is_video_file(c)]

            if candidates:
                self.alpha_asset = ClipAsset(candidates[0], "video")
            else:
                self.alpha_asset = None  # Missing, needs generation

    def validate_pair(self) -> None:
        if self.input_asset and self.alpha_asset:
            if self.input_asset.frame_count != self.alpha_asset.frame_count:
                raise ValueError(
                    f"Clip '{self.name}': Frame count mismatch! "
                    f"Input: {self.input_asset.frame_count}, Alpha: {self.alpha_asset.frame_count}"
                )


# --- Logic ---


def get_gvm_processor(device: str = "cpu") -> GVMProcessor:
    try:
        from gvm_core import GVMProcessor

        return GVMProcessor(device=device)
    except ImportError:
        raise ImportError(
            "Could not import gvm_core. Please ensure 'gvm_core' is in the project root and requirements are installed."
        ) from None
    except Exception as e:
        raise RuntimeError(f"Failed to initialize GVM Processor: {e}") from e


def generate_alphas(
    clips,
    device=None,
    *,
    on_clip_start: Callable[[str, int], None] | None = None,
):
    clips_to_process = [c for c in clips if c.alpha_asset is None]

    if not clips_to_process:
        logger.info("All clips have valid Alpha assets. No generation needed.")
        return

    logger.info(f"Found {len(clips_to_process)} clips missing Alpha.")

    if device is None:
        device = resolve_device()

    try:
        processor = get_gvm_processor(device=device)
    except ImportError as e:
        logger.error(f"GVM Import Error: {e}")
        logger.error("Skipping GVM generation. Please install GVM requirements if you wish to use this feature.")
        return
    except Exception as e:
        logger.error(f"GVM Initialization Error: {e}")
        return

    for clip in clips_to_process:
        logger.info(f"Generating Alpha for: {clip.name}")
        if on_clip_start:
            on_clip_start(clip.name, len(clips_to_process))

        alpha_output_dir = os.path.join(clip.root_path, "AlphaHint")
        if os.path.exists(alpha_output_dir):
            shutil.rmtree(alpha_output_dir)
        os.makedirs(alpha_output_dir, exist_ok=True)

        try:
            processor.process_sequence(
                input_path=clip.input_asset.path,
                output_dir=None,
                num_frames_per_batch=1,
                decode_chunk_size=1,
                denoise_steps=1,
                mode="matte",
                write_video=False,
                direct_output_dir=alpha_output_dir,
            )

            # Post-Process: Naming Convention
            generated_files = sorted([f for f in os.listdir(alpha_output_dir) if f.endswith(".png")])

            if not generated_files:
                logger.error(f"GVM finished but no PNGs found in {alpha_output_dir}")
                continue

            if clip.input_asset.type == "sequence":
                in_files = sorted([f for f in os.listdir(clip.input_asset.path) if is_image_file(f)])
                stems = [os.path.splitext(f)[0] for f in in_files]
            else:
                base_name = os.path.splitext(os.path.basename(clip.input_asset.path))[0]
                stems = [base_name] * len(generated_files)

            for i, gvm_file in enumerate(generated_files):
                if i >= len(stems):
                    break

                stem = stems[i]
                new_name = f"{stem}_alphaHint_{i:04d}.png"

                old_path = os.path.join(alpha_output_dir, gvm_file)
                new_path = os.path.join(alpha_output_dir, new_name)

                if old_path != new_path:
                    os.rename(old_path, new_path)

            logger.info(f"Saved {len(generated_files)} alpha frames to {alpha_output_dir}")

        except Exception as e:
            logger.error(f"Error generating alpha for {clip.name}: {e}")
            import traceback

            traceback.print_exc()


def get_birefnet_usage_options():
    return list(usage_to_weights_file.keys())


def run_birefnet(
    clips,
    device=None,
    usage="General",
    dilate_radius=0,
    *,
    on_clip_start: Callable[[str, int], None] | None = None,
    on_frame_complete: Callable[[int, int], None] | None = None,
):
    clips_to_process = [c for c in clips if c.alpha_asset is None]

    if not clips_to_process:
        logger.info("All clips have valid Alpha assets. No BiRefNet generation needed.")
        return

    if device is None:
        device = resolve_device()

    logger.info(f"Found {len(clips_to_process)} clips missing Alpha.")

    logger.info(f"Initializing BiRefNet ({usage}) on {device}...")
    # Initialize the handler once
    try:
        handler = BiRefNetHandler(device=device, usage=usage)
    except ImportError as e:
        logger.error(f"BiRefNet Import Error: {e}")
        return
    except Exception as e:
        logger.error(f"BiRefNet Initialization Error: {e}")
        return

    try:
        for clip in clips_to_process:
            logger.info(f"Generating BiRefNet Alpha for: {clip.name}")
            if on_clip_start:
                on_clip_start(clip.name, clip.input_asset.frame_count)

            alpha_output_dir = os.path.join(clip.root_path, "AlphaHint")
            os.makedirs(alpha_output_dir, exist_ok=True)

            try:
                handler.process(
                    input_path=clip.input_asset.path,
                    alpha_output_dir=alpha_output_dir,
                    dilate_radius=dilate_radius,
                    on_frame_complete=on_frame_complete,
                )
                logger.info(f"BiRefNet complete for {clip.name}")
            except Exception as e:
                logger.error(f"BiRefNet failed for {clip.name}: {e}")
                import traceback

                traceback.print_exc()

    finally:
        handler.cleanup()


def run_videomama(
    clips: list[ClipEntry],
    chunk_size: int = 50,
    device: str | None = None,
    *,
    on_clip_start: Callable[[str, int], None] | None = None,
    on_frame_complete: Callable[[int, int], None] | None = None,
) -> None:
    """
    Runs VideoMaMa on clips that have VideoMamaMaskHint but NO AlphaHint.
    """
    # Process if:
    # 1. Has VideoMamaMaskHint (File or Folder, Case-Insensitive)
    # 2. AND (Alpha is Missing OR Alpha is a Video File we want to upgrade)

    clips_to_process = []
    clip_mask_paths = {}  # Store the resolved mask path for each clip

    for c in clips:
        # Search for 'videomamamaskhint' asset (Strict: videomamamaskhint.ext or VideoMamaMaskHint/)
        candidates = []
        for f in os.listdir(c.root_path):
            stem, _ = os.path.splitext(f)
            if stem.lower() == "videomamamaskhint":
                candidates.append(f)

        mask_asset_path = None
        has_mask = False

        # Priority: Directory > Video File
        # Check directories first
        for cand in candidates:
            path = os.path.join(c.root_path, cand)
            if os.path.isdir(path) and len(os.listdir(path)) > 0:
                has_mask = True
                mask_asset_path = path
                break

        # If no directory, check files
        if not has_mask:
            for cand in candidates:
                path = os.path.join(c.root_path, cand)
                if os.path.isfile(path) and is_video_file(path):
                    has_mask = True
                    mask_asset_path = path
                    break

        if not has_mask:
            continue

        # Store for later
        clip_mask_paths[c.name] = mask_asset_path

        if c.alpha_asset is None:
            clips_to_process.append(c)
        elif c.alpha_asset.type == "video":
            clips_to_process.append(c)

    if not clips_to_process:
        logger.info("No candidates for VideoMaMa (looking for VideoMamaMaskHint + [NoAlpha OR VideoAlpha]).")
        return

    logger.info(f"Found {len(clips_to_process)} clips for VideoMaMa processing.")

    # Import locally — sys.path mutation is needed because VideoMaMaInferenceModule
    # uses intra-package imports that assume its directory is on the path.
    try:
        sys.path.append(os.path.join(BASE_DIR, "VideoMaMaInferenceModule"))
        from VideoMaMaInferenceModule.inference import load_videomama_model
        from VideoMaMaInferenceModule.inference import run_inference as run_videomama_frames
    except ImportError as e:
        logger.error(f"Failed to import VideoMaMa: {e}")
        return

    if device is None:
        device = resolve_device()

    logger.info("Loading VideoMaMa Pipeline...")
    pipeline = load_videomama_model(device=device)

    for clip in clips_to_process:
        logger.info(f"Running VideoMaMa on: {clip.name}")
        if on_clip_start:
            on_clip_start(clip.name, len(clips_to_process))

        # Retrieve resolved path
        mask_hint_path = clip_mask_paths[clip.name]
        logger.info(f"  Using VideoMamaMaskHint: {os.path.basename(mask_hint_path)}")

        alpha_output_dir = os.path.join(clip.root_path, "AlphaHint")

        # Load Inputs
        # 1. Input Frames (RGB)
        input_frames = []
        if clip.input_asset.type == "video":
            cap = cv2.VideoCapture(clip.input_asset.path)
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                input_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            cap.release()
        else:
            files = sorted([f for f in os.listdir(clip.input_asset.path) if is_image_file(f)])
            for f in files:
                fpath = os.path.join(clip.input_asset.path, f)
                # Handle EXR (Float 0-1) vs Standard (Int 0-255)
                if f.lower().endswith(".exr"):
                    img = cv2.imread(fpath, cv2.IMREAD_UNCHANGED)
                    if img is not None:
                        # Normalize Float 0-1
                        img = np.clip(img, 0.0, 1.0)
                        # Linear -> sRGB (Gamma 2.2 Approximation) for VideoMaMa
                        img = img ** (1.0 / 2.2)
                        # 0-255 uint8
                        img = (img * 255.0).astype(np.uint8)
                        # Ensure 3 channels
                        if img.ndim == 2:
                            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
                        elif img.shape[2] == 4:
                            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
                else:
                    img = cv2.imread(fpath)

                if img is not None:
                    input_frames.append(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

        # 2. Mask Frames
        mask_frames = []

        # Check if VideoMamaMaskHint is a directory or a file (video)
        if os.path.isdir(mask_hint_path):
            # Directory of Images
            mask_files = sorted([f for f in os.listdir(mask_hint_path) if is_image_file(f)])
            for f in mask_files:
                fpath = os.path.join(mask_hint_path, f)
                m = None

                # Handle EXR Masks
                if f.lower().endswith(".exr"):
                    m = cv2.imread(fpath, cv2.IMREAD_UNCHANGED)
                    if m is not None:
                        if m.ndim == 3:
                            m = m[:, :, 0]
                        m = np.clip(m, 0.0, 1.0)
                        m = (m * 255.0).astype(np.uint8)
                else:
                    # Standard Masks
                    m = cv2.imread(fpath, cv2.IMREAD_GRAYSCALE)

                if m is not None:
                    # Force Binary Thresholding
                    _, m = cv2.threshold(m, 10, 255, cv2.THRESH_BINARY)
                    mask_frames.append(m)

        elif os.path.isfile(mask_hint_path):
            # Handle Video File
            cap = cv2.VideoCapture(mask_hint_path)
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                # Convert to Grayscale
                m = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                # Force Binary Thresholding
                _, m = cv2.threshold(m, 10, 255, cv2.THRESH_BINARY)
                mask_frames.append(m)
            cap.release()

        # Validate Lengths
        num_frames = min(len(input_frames), len(mask_frames))
        input_frames = input_frames[:num_frames]
        mask_frames = mask_frames[:num_frames]

        if num_frames == 0:
            logger.error(f"Skipping {clip.name}: No valid frame pairs found.")
            continue

        # Run Inference
        try:
            # Prepare Output Directory First
            # Logic: If it exists as a FILE (legacy/error), delete it.
            if os.path.exists(alpha_output_dir) and not os.path.isdir(alpha_output_dir):
                logger.warning(f"Removing file '{alpha_output_dir}' to create directory.")
                os.remove(alpha_output_dir)

            # If there was a Video Alpha Asset (e.g. AlphaHint.mp4), rename it to backup so it doesn't conflict
            if clip.alpha_asset and clip.alpha_asset.type == "video":
                old_path = clip.alpha_asset.path
                if os.path.exists(old_path):
                    dir_name = os.path.dirname(old_path)
                    base, ext = os.path.splitext(os.path.basename(old_path))
                    backup_path = os.path.join(dir_name, f"{base}_backup{ext}")
                    logger.info(
                        f"Backing up existing Alpha Video: "
                        f"{os.path.basename(old_path)} -> {os.path.basename(backup_path)}"
                    )
                    os.rename(old_path, backup_path)
                    # Clear it from memory so we rely on the new one
                    clip.alpha_asset = None

            os.makedirs(alpha_output_dir, exist_ok=True)

            # Name setup
            if clip.input_asset.type == "sequence":
                in_names = sorted(
                    [os.path.splitext(f)[0] for f in os.listdir(clip.input_asset.path) if is_image_file(f)]
                )
            else:
                stem = os.path.splitext(os.path.basename(clip.input_asset.path))[0]
                in_names = [f"{stem}_{i:05d}" for i in range(num_frames)]

            total_saved = 0

            # Iterate generator
            for chunk_frames in run_videomama_frames(pipeline, input_frames, mask_frames, chunk_size=chunk_size):
                for frame in chunk_frames:
                    if total_saved >= len(in_names):
                        break

                    name = in_names[total_saved]
                    out_path = os.path.join(alpha_output_dir, f"{name}.png")

                    # Convert to BGR and Save
                    cv2.imwrite(out_path, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
                    total_saved += 1

                    if on_frame_complete:
                        on_frame_complete(total_saved, num_frames)

                logger.info(f"  Saved {total_saved}/{num_frames} frames...")

            logger.info(f"VideoMaMa Complete: Saved {total_saved} frames to AlphaHint.")

            # Update clip state in memory (dummy) - re-scan will pick it up properly
            clip.alpha_asset = ClipAsset(alpha_output_dir, "sequence")

        except Exception as e:
            logger.error(f"VideoMaMa failed for {clip.name}: {e}")
            import traceback

            traceback.print_exc()


def run_inference(
    clips,
    device=None,
    backend=None,
    max_frames=None,
    skip_existing=False,
    settings: InferenceSettings | None = None,
    *,
    on_clip_start: Callable[[str, int], None] | None = None,
    on_frame_complete: Callable[[int, int], None] | None = None,
):
    ready_clips = [c for c in clips if c.input_asset and c.alpha_asset]

    if not ready_clips:
        logger.info("No clips found with both Input and Alpha assets. Run generate_coarse_alpha first?")
        return

    logger.info(f"Found {len(ready_clips)} clips ready for inference.")

    # Backward compat for callers that don't pass settings
    if settings is None:
        settings = InferenceSettings()

    # Ensure Output Directory exists
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR, exist_ok=True)

    import numpy as np

    if device is None:
        device = resolve_device()
    from CorridorKeyModule.backend import create_engine

    engine = create_engine(backend=backend, device=device)

    for clip in ready_clips:
        logger.info(f"Running Inference on: {clip.name}")

        # Setup Outputs in ClipFolder/Output/...
        clip_out_root = os.path.join(clip.root_path, "Output")
        fg_dir = os.path.join(clip_out_root, "FG")
        matte_dir = os.path.join(clip_out_root, "Matte")
        comp_dir = os.path.join(clip_out_root, "Comp")
        proc_dir = os.path.join(clip_out_root, "Processed")

        for d in [fg_dir, matte_dir, comp_dir, proc_dir]:
            os.makedirs(d, exist_ok=True)

        num_frames = min(clip.input_asset.frame_count, clip.alpha_asset.frame_count)
        if max_frames is not None:
            num_frames = min(num_frames, max_frames)
        logger.info(
            f"  Input frames: {clip.input_asset.frame_count},"
            f" Alpha frames: {clip.alpha_asset.frame_count} -> Processing {num_frames} frames"
        )

        if num_frames == 0:
            logger.warning(f"Clip '{clip.name}': 0 frames to process, skipping.")
            continue

        input_cap = None
        alpha_cap = None
        input_files = []
        alpha_files = []

        if clip.input_asset.type == "video":
            input_cap = cv2.VideoCapture(clip.input_asset.path)
        else:
            input_files = sorted([f for f in os.listdir(clip.input_asset.path) if is_image_file(f)])

        if clip.alpha_asset.type == "video":
            alpha_cap = cv2.VideoCapture(clip.alpha_asset.path)
        else:
            alpha_files = sorted([f for f in os.listdir(clip.alpha_asset.path) if is_image_file(f)])

        if on_clip_start:
            on_clip_start(clip.name, num_frames)

        skipped_count = 0

        for i in range(num_frames):
            # Pre-compute output stem for skip-existing check (mirrors how input_stem
            # is set later: video -> zero-padded index, sequence -> file stem)
            if clip.input_asset.type == "video":
                expected_stem = f"{i:05d}"
            else:
                expected_stem = os.path.splitext(input_files[i])[0]

            if skip_existing and os.path.exists(os.path.join(comp_dir, f"{expected_stem}.png")):
                logger.debug("Frame %d already rendered, skipping.", i)
                skipped_count += 1
                if on_frame_complete:
                    on_frame_complete(i, num_frames)
                continue

            # 1. Read Input
            img_srgb = None
            input_stem = f"{i:05d}"

            # Use the settings-defined gamma
            input_is_linear = settings.input_is_linear

            if input_cap:
                ret, frame = input_cap.read()
                if not ret:
                    break
                img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img_srgb = img_rgb.astype(np.float32) / 255.0
                input_stem = f"{i:05d}"
            else:
                fpath = os.path.join(clip.input_asset.path, input_files[i])
                input_stem = os.path.splitext(input_files[i])[0]

                is_exr = fpath.lower().endswith(".exr")
                if is_exr:
                    img_srgb = read_image_frame(fpath, gamma_correct_exr=not input_is_linear)
                    if img_srgb is None:
                        continue
                else:
                    img_bgr = cv2.imread(fpath)
                    if img_bgr is None:
                        continue
                    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
                    img_srgb = img_rgb.astype(np.float32) / 255.0

            # 2. Read Alpha (Mask)
            mask_linear = None
            if alpha_cap:
                ret, frame = alpha_cap.read()
                if not ret:
                    break
                mask_linear = frame[:, :, 2].astype(np.float32) / 255.0
            else:
                fpath = os.path.join(clip.alpha_asset.path, alpha_files[i])
                mask_in = cv2.imread(fpath, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_UNCHANGED)

                if mask_in is None:
                    continue

                if mask_in.ndim == 3:
                    if mask_in.shape[2] == 3:
                        mask_linear = mask_in[:, :, 0]
                    else:
                        mask_linear = mask_in
                else:
                    mask_linear = mask_in

                if mask_linear.dtype == np.uint8:
                    mask_linear = mask_linear.astype(np.float32) / 255.0
                elif mask_linear.dtype == np.uint16:
                    mask_linear = mask_linear.astype(np.float32) / 65535.0
                else:
                    mask_linear = mask_linear.astype(np.float32)

            if mask_linear.shape[:2] != img_srgb.shape[:2]:
                mask_linear = cv2.resize(
                    mask_linear, (img_srgb.shape[1], img_srgb.shape[0]), interpolation=cv2.INTER_LINEAR
                )

            # 3. Process
            USE_STRAIGHT_MODEL = True
            res = engine.process_frame(
                img_srgb,
                mask_linear,
                input_is_linear=input_is_linear,
                fg_is_straight=USE_STRAIGHT_MODEL,
                despill_strength=settings.despill_strength,
                auto_despeckle=settings.auto_despeckle,
                despeckle_size=settings.despeckle_size,
                refiner_scale=settings.refiner_scale,
            )

            pred_fg = res["fg"]  # sRGB
            pred_alpha = res["alpha"]  # Linear

            # 4. Save (EXR half-float, PXR24 compression — see backend/frame_io.py)

            # Save FG
            # pred_fg is RGB 0-1 float. Convert to BGR for OpenCV
            fg_bgr = cv2.cvtColor(pred_fg, cv2.COLOR_RGB2BGR)
            cv2.imwrite(os.path.join(fg_dir, f"{input_stem}.exr"), fg_bgr, EXR_WRITE_FLAGS)

            # Save Matte
            if pred_alpha.ndim == 3:
                pred_alpha = pred_alpha[:, :, 0]
            # Matte is single channel linear float
            cv2.imwrite(os.path.join(matte_dir, f"{input_stem}.exr"), pred_alpha, EXR_WRITE_FLAGS)

            # 5. Generate Reference Comp
            comp_srgb = res["comp"]
            # Save Comp (PNG 8-bit)
            comp_bgr = cv2.cvtColor((np.clip(comp_srgb, 0.0, 1.0) * 255.0).astype(np.uint8), cv2.COLOR_RGB2BGR)
            cv2.imwrite(os.path.join(comp_dir, f"{input_stem}.png"), comp_bgr)

            # 6. Save Processed (RGBA EXR)
            if "processed" in res:
                # Result is RGBA
                proc_rgba = res["processed"]
                # Convert to BGRA for OpenCV
                proc_bgra = cv2.cvtColor(proc_rgba, cv2.COLOR_RGBA2BGRA)
                cv2.imwrite(os.path.join(proc_dir, f"{input_stem}.exr"), proc_bgra, EXR_WRITE_FLAGS)

            if on_frame_complete:
                on_frame_complete(i, num_frames)

        if input_cap:
            input_cap.release()
        if alpha_cap:
            alpha_cap.release()

        if skip_existing and skipped_count > 0:
            logger.info(
                "  Skipped %d of %d frames (outputs already exist).",
                skipped_count,
                num_frames,
            )

        # 7. Stitch comp frames into MP4 (if input was video)
        if clip.input_asset and clip.input_asset.type == "video":
            try:
                from backend.ffmpeg_tools import find_ffmpeg, probe_video, stitch_video

                if find_ffmpeg():
                    # Get source fps
                    try:
                        video_info = probe_video(clip.input_asset.path)
                        fps = video_info.get("fps", 24.0)
                    except Exception:
                        fps = 24.0

                    comp_video_path = os.path.join(clip_out_root, f"{clip.name}_comp.mp4")

                    # Detect frame pattern from saved files
                    comp_files = sorted(f for f in os.listdir(comp_dir) if f.endswith(".png"))
                    if comp_files:
                        # Frames are named {input_stem}.png — e.g. 00000.png
                        # Build ffmpeg pattern from first file
                        first = comp_files[0]
                        stem = os.path.splitext(first)[0]
                        if stem.isdigit():
                            pattern = f"%0{len(stem)}d.png"
                        else:
                            pattern = "frame_%06d.png"

                        logger.info(f"Stitching comp video: {comp_dir} -> {comp_video_path} @ {fps} fps")
                        stitch_video(comp_dir, comp_video_path, fps=fps, pattern=pattern)
                    else:
                        logger.warning(f"No comp frames found in {comp_dir}, skipping video stitch.")
                else:
                    logger.info("ffmpeg not found — skipping comp video stitch.")
            except Exception as e:
                logger.warning(f"Comp video stitch failed (non-fatal): {e}")

        logger.info(f"Clip {clip.name} Complete.")


def organize_target(target_dir: str) -> None:
    """
    Organizes a specific folder.
    1. If loose video -> Rename to Input.ext (if safe).
    2. If sequence -> Move to Input/.
    3. Ensure AlphaHint and VideoMamaMaskHint folders exist.
    """
    logger.info(f"Organizing Target: {target_dir}")

    if not os.path.exists(target_dir):
        logger.error(f"Target directory not found: {target_dir}")
        return

    # Check for loose video
    # Strategy: Find largest video file that ISN'T named Input.*
    candidates = [f for f in os.listdir(target_dir) if is_video_file(f)]
    candidates = [f for f in candidates if not os.path.splitext(f)[0].lower() == "input"]

    if candidates and not os.path.exists(os.path.join(target_dir, "Input")):
        # If multiple, pick largest (heuristic for 'Main Plate')
        candidates.sort(key=lambda f: os.path.getsize(os.path.join(target_dir, f)), reverse=True)
        main_clip = candidates[0]
        ext = os.path.splitext(main_clip)[1]

        try:
            shutil.move(os.path.join(target_dir, main_clip), os.path.join(target_dir, f"Input{ext}"))
            logger.info(f"Renamed '{main_clip}' to 'Input{ext}'")
        except Exception as e:
            logger.error(f"Failed to rename '{main_clip}': {e}")

    # Check for Image Sequence (Flat)
    # Only if Input folder doesn't exist and Input video doesn't exist
    has_input_dir = os.path.isdir(os.path.join(target_dir, "Input"))
    has_input_video = any(
        is_video_file(f) and os.path.basename(f).lower().startswith("input") for f in os.listdir(target_dir)
    )

    if not has_input_dir and not has_input_video:
        all_files = sorted(glob.glob(os.path.join(target_dir, "*")))
        image_files = [f for f in all_files if is_image_file(f)]

        if len(image_files) > 0:
            try:
                input_subdir = os.path.join(target_dir, "Input")
                os.makedirs(input_subdir)
                for img in image_files:
                    shutil.move(img, os.path.join(input_subdir, os.path.basename(img)))
                logger.info(
                    f"Organized: Moved {len(image_files)} images in '{os.path.basename(target_dir)}' to 'Input/'"
                )
            except Exception as e:
                logger.error(f"Failed to organize sequence in '{target_dir}': {e}")

    # Create Hints
    for hint in ["AlphaHint", "VideoMamaMaskHint"]:
        hint_path = os.path.join(target_dir, hint)
        if not os.path.exists(hint_path):
            os.makedirs(hint_path)


def organize_clips(clips_dir: str) -> None:
    """
    Legacy wrapper for backward compatibility with 'ClipsForInference' folder.
    Organizes all subfolders in the given directory using the new logic.
    """
    if not os.path.exists(clips_dir):
        logger.warning(f"Clips directory not found: {clips_dir}")
        return

    logger.info(f"Organizing Clips Directory: {clips_dir}")

    # Check for loose videos in root
    loose_videos = [f for f in os.listdir(clips_dir) if is_video_file(f) and os.path.isfile(os.path.join(clips_dir, f))]

    # Organize loose videos first
    for v in loose_videos:
        clip_name = os.path.splitext(v)[0]
        ext = os.path.splitext(v)[1]
        target_folder = os.path.join(clips_dir, clip_name)

        if os.path.exists(target_folder):
            logger.warning(f"Skipping loose video '{v}': Target folder '{clip_name}' already exists.")
            continue

        try:
            os.makedirs(target_folder)
            target_file = os.path.join(target_folder, f"Input{ext}")
            shutil.move(os.path.join(clips_dir, v), target_file)
            logger.info(f"Organized: Moved '{v}' to '{clip_name}/Input{ext}'")

            # Also initialize hints immediately
            for hint in ["AlphaHint", "VideoMamaMaskHint"]:
                os.makedirs(os.path.join(target_folder, hint), exist_ok=True)
        except Exception as e:
            logger.error(f"Failed to organize video '{v}': {e}")

    # Now iterate all subdirectories and run organize_target
    for entry in os.listdir(clips_dir):
        full_path = os.path.join(clips_dir, entry)
        if os.path.isdir(full_path) and entry not in ["IgnoredClips", "Output"]:
            organize_target(full_path)


def scan_clips() -> list[ClipEntry]:
    if not os.path.exists(CLIPS_DIR):
        os.makedirs(CLIPS_DIR, exist_ok=True)
        return []

    # Auto-organize first
    organize_clips(CLIPS_DIR)

    clip_dirs = [d for d in os.listdir(CLIPS_DIR) if os.path.isdir(os.path.join(CLIPS_DIR, d))]

    valid_clips = []
    invalid_clips = []

    for d in clip_dirs:
        if d.startswith(".") or d.startswith("_") or d == "IgnoredClips":
            continue

        full_path = os.path.join(CLIPS_DIR, d)

        try:
            entry = ClipEntry(d, full_path)
            entry.find_assets()
            entry.validate_pair()
            valid_clips.append(entry)
        except ValueError as ve:
            invalid_clips.append((d, str(ve)))
        except Exception as e:
            invalid_clips.append((d, f"Unexpected error: {e}"))

    if invalid_clips:
        logger.warning("INVALID OR SKIPPED CLIPS:")
        for name, reason in invalid_clips:
            logger.warning("  - %s: %s", name, reason)
    else:
        logger.info("All clip folders appear valid.")

    return valid_clips


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="CorridorKey Clip Manager")
    parser.add_argument("--action", choices=["generate_alphas", "run_inference", "list", "wizard"], required=True)
    parser.add_argument("--win_path", help=r"Windows Path (example: V:\...) for Wizard Mode", default=None)
    parser.add_argument(
        "--device",
        choices=["auto", "cuda", "mps", "cpu"],
        default="auto",
        help="Compute device (default: auto-detect CUDA > MPS > CPU)",
    )
    parser.add_argument(
        "--backend",
        choices=["auto", "torch", "mlx"],
        default="auto",
        help="Inference backend (default: auto-detect MLX on Apple Silicon, else Torch)",
    )
    parser.add_argument(
        "--max-frames",
        type=int,
        default=None,
        help="Limit number of frames to process per clip (e.g. 1 for first frame only)",
    )

    args = parser.parse_args()

    device = resolve_device(args.device)
    logger.info(f"Using device: {device}")

    if args.action == "list":
        scan_clips()
    elif args.action == "generate_alphas":
        clips = scan_clips()
        generate_alphas(clips, device=device)
    elif args.action == "run_inference":
        clips = scan_clips()
        run_inference(clips, device=device, backend=args.backend, max_frames=args.max_frames)
    elif args.action == "wizard":
        if not args.win_path:
            print("Error: --win_path required for wizard.")
        else:
            raise NotImplementedError("interactive_wizard is not yet implemented")