VFXProxy/parser.py

"""
parser.py — Filename parsing and Shot data model.

Filename format:
    SHOW_EPISODE_SEQUENCE_SHOT_ELEMENT_TASK_VERSION.FRAME.exr

Example:
    UNG_101_012_010_BG01_TT_v001.01001.exr
"""
from __future__ import annotations

import re
import struct
from dataclasses import dataclass, field
from pathlib import Path


# ──────────────────────────────────────────────────────────────────────────────
# Regex patterns
# ──────────────────────────────────────────────────────────────────────────────

# Matches the basename (no frame, no extension):
#   SHOW_EPISODE_SEQUENCE_SHOT_ELEMENT_TASK_VERSION
BASENAME_RE = re.compile(
    r"^(?P<show>[^_]+)_"
    r"(?P<episode>[^_]+)_"
    r"(?P<sequence>[^_]+)_"
    r"(?P<shot>[^_]+)_"
    r"(?P<element>[^_]+)_"
    r"(?P<task>[^_]+)_"
    r"(?P<version>v\d+)$",
    re.IGNORECASE,
)

# Matches a full EXR sequence filename: {basename}.{4-7 digit frame}.exr
SEQ_FILE_RE = re.compile(r"^(.+)\.(\d{4,7})\.exr$", re.IGNORECASE)


# ──────────────────────────────────────────────────────────────────────────────
# Data model
# ──────────────────────────────────────────────────────────────────────────────

@dataclass
class Shot:
    show: str
    episode: str
    sequence: str
    shot: str
    element: str
    task: str
    version: str
    basename: str
    first_frame: int
    last_frame: int
    frame_count: int
    missing_frames: list[int] = field(default_factory=list)
    width: int = 0
    height: int = 0
    input_sequence: str = ""   # Nuke-style path, e.g. /path/basename.%05d.exr
    output_folder: str = ""
    output_movie: str = ""
    output_exr: str = ""
    status: str = "Ready"

    @property
    def display_frames(self) -> str:
        s = f"{self.first_frame}\u2013{self.last_frame}  ({self.frame_count}f)"
        if self.missing_frames:
            s += f"  [{len(self.missing_frames)} missing]"
        return s

    @property
    def display_resolution(self) -> str:
        return f"{self.width}\u00d7{self.height}" if self.width and self.height else "\u2014"


# ──────────────────────────────────────────────────────────────────────────────
# Parsing helpers
# ──────────────────────────────────────────────────────────────────────────────

def parse_basename(basename: str) -> dict | None:
    """Return a dict of named fields for *basename*, or None if it does not match."""
    m = BASENAME_RE.match(basename)
    return m.groupdict() if m else None


def build_shot(
    basename: str,
    source_dir: str,
    frames: list[int],
    output_root: str,
    parsed: dict,
    frame_padding: int = 4,
) -> Shot:
    """Construct a :class:`Shot` from parsed filename parts and frame list."""
    frames_sorted = sorted(frames)
    first = frames_sorted[0]
    last = frames_sorted[-1]
    missing = sorted(set(range(first, last + 1)) - set(frames_sorted))

    show = parsed["show"]
    episode = parsed["episode"]

    # Nuke requires forward slashes even on Windows
    input_dir = Path(source_dir).as_posix()
    input_seq = f"{input_dir}/{basename}.%0{frame_padding}d.exr"

    out_folder = Path(output_root) / show / f"{show}_{episode}" / basename
    out_folder_str = out_folder.as_posix()

    # Attempt to read resolution from the first frame of the sequence
    first_exr = Path(source_dir) / f"{basename}.{str(first).zfill(frame_padding)}.exr"
    width, height = read_exr_resolution(str(first_exr))

    return Shot(
        show=show,
        episode=episode,
        sequence=parsed["sequence"],
        shot=parsed["shot"],
        element=parsed["element"],
        task=parsed["task"],
        version=parsed["version"],
        basename=basename,
        first_frame=first,
        last_frame=last,
        frame_count=len(frames_sorted),
        missing_frames=missing,
        width=width,
        height=height,
        input_sequence=input_seq,
        output_folder=out_folder_str,
        output_movie=f"{out_folder_str}/{basename}.mov",
        output_exr=f"{out_folder_str}/{basename}.%0{frame_padding}d.exr",
    )


# ──────────────────────────────────────────────────────────────────────────────
# EXR header reader (no external dependencies)
# ──────────────────────────────────────────────────────────────────────────────

_EXR_MAGIC = b"\x76\x2f\x31\x01"


def read_exr_resolution(filepath: str) -> tuple[int, int]:
    """Read (width, height) from an OpenEXR file header without external libraries.

    Returns (0, 0) if the file cannot be read or is not a valid EXR.
    """
    try:
        with open(filepath, "rb") as fh:
            if fh.read(4) != _EXR_MAGIC:
                return 0, 0
            fh.read(4)  # version + feature flags
            while True:
                name = _read_cstring(fh)
                if not name:
                    break  # end of header
                type_str = _read_cstring(fh)
                size = struct.unpack("<I", fh.read(4))[0]
                data = fh.read(size)
                if name == "dataWindow" and type_str == "box2i":
                    xmin, ymin, xmax, ymax = struct.unpack("<iiii", data)
                    return xmax - xmin + 1, ymax - ymin + 1
    except Exception:
        pass
    return 0, 0


def _read_cstring(fh) -> str:
    """Read a null-terminated ASCII string from a binary file handle."""
    buf = bytearray()
    while True:
        c = fh.read(1)
        if not c or c == b"\x00":
            break
        buf.extend(c)
    return buf.decode("latin-1")