batch_utils

def add_reels_to_batch(
    batch: flame.PyBatch,
    reels: Optional[List[str]] = None,
    shelf_reels: Optional[List[str]] = None,
):
    """ Add reels and shelf reels to batch.
    """
    if reels:
        existing_reel_names = [
            reel.name.get_value()
            for reel in batch.reels
        ]
        for new_reel in reels:
            if new_reel not in existing_reel_names:
                batch.create_reel(new_reel)

    if shelf_reels:
        existing_shelf_reel_names = [
            reel.name.get_value()
            for reel in batch.shelf_reels
        ]
        for new_sr in shelf_reels:
            if new_sr not in existing_shelf_reel_names:
                batch.create_shelf_reel(new_sr)

def clear_node_metadata(node: flame.PyNode):
    """ Remove AYON instance data from a node's note attribute."""
    node.note = ""

def create_batch(
    name,
    frame_start: int,
    frame_duration: int,
    handle_start: int = 0,
    handle_end: int = 0,
) -> flame.PyBatch:
    """ Create Batch Group in active project's Desktop
    """
    frame_start -= handle_start
    frame_duration += handle_start + handle_end

    return flame.batch.create_batch_group(
        name,
        start_frame=frame_start,
        duration=frame_duration,
    )

def get_batch_from_workspace(
    name: str,
    workspace: Optional[flame.PyWorkspace] = None
) -> Optional[flame.PyBatch]:
    """ Get batch group from name and workspace.
    """
    if workspace is None:
        project = flame.project.current_project
        workspace = project.current_workspace

    desktop = workspace.desktop
    for batchgroup in desktop.batch_groups:
        if batchgroup.name.get_value() == name:
            return batchgroup

    return None

def load_batch_from_consolidated_json(
    filepath: str,
    name: Optional[str] = None,
    temporary_folder: Optional[str] = None,
) -> Optional[flame.PyBatch]:
    """ Load a batch from a consolidated json file.
    """
    with open(filepath, "r", encoding="utf-8") as file_:
        data = json.load(file_)

    tmp = tempfile.TemporaryDirectory() if not temporary_folder else None
    tmp_dir = pathlib.Path(tmp.name if tmp else temporary_folder)

    try:
        batch_file = None
        for relative_file, content in data.items():
            file_path = tmp_dir / relative_file
            file_path.parent.mkdir(parents=True, exist_ok=True)
            if content.startswith("__b64__:"):
                file_path.write_bytes(base64.b64decode(content[8:]))
            else:
                file_path.write_text(content, encoding="utf-8")

            if relative_file.endswith(".batch"):
                batch_file = relative_file

        if batch_file is None:
            raise ValueError(
                f"No valid batch found in consolidated json: {filepath}"
            )

        flame.batch.load_setup(str(tmp_dir / batch_file))

        # Restore the batch group name from the provided name
        # or use the .batch filename stem otherwise.
        batch_name = name or pathlib.Path(batch_file).stem
        flame.batch.name = batch_name

        return flame.batch

    finally:
        if tmp is not None:
            tmp.cleanup()

def read_node_metadata(node: flame.PyNode) -> Optional[Dict[str, Any]]:
    """ Read AYON instance data from a node's note attribute.
    """
    try:
        raw = node.note.get_value()
    except Exception:
        return None

    if not raw:
        return None

    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        return None

    return data if data.get(AYON_NOTE_MARKER) else None

def save_batch_as_consolidated_json(
    batch: flame.PyBatch,
    filepath: str,
    temporary_folder: Optional[str] = None,  # where flame run native export
) -> str:
    """ Export batch as a consolidated json file.
    """
    tmp = tempfile.TemporaryDirectory() if temporary_folder is None else None
    tmp_dir = pathlib.Path(tmp.name if tmp else temporary_folder)

    try:
        batch_name = batch.name.get_value()
        bgroup_file = tmp_dir / f"{batch_name}.batch"
        batch.save_setup(str(bgroup_file))

        if not tmp_dir.is_dir():
            raise RuntimeError(
                f"Unable to save batchgroup to folder: {tmp_dir}."
            )

        # Concatenate all intermediary files as 1 single consolidated JSON.
        # Binary files (e.g. .mx presets are not utf-8) are base64-encoded
        # using a "__b64__:" prefix so we can deserialize them.
        json_output = {}
        for file_path in tmp_dir.rglob("*"):
            if file_path.is_file():
                relative_path = file_path.relative_to(tmp_dir)
                try:
                    content = file_path.read_text(encoding="utf-8")
                except (UnicodeDecodeError, ValueError):
                    content = "__b64__:" + base64.b64encode(
                        file_path.read_bytes()
                    ).decode("ascii")
                json_output[str(relative_path)] = content

        with open(filepath, "w") as file_handler:
            json.dump(json_output, file_handler, indent=4)

    finally:
        # Delete temporary directory if created.
        if tmp is not None:
            tmp.cleanup()

    return filepath

def update_batch(
    batch: flame.PyBatch,
    name: Optional[str] = None,
    frame_start: Optional[int] = None,
    frame_duration: Optional[int] = None,
    handle_start: int = 0,
    handle_end: int = 0,
) -> flame.PyBatch:
    """ Update provided batch with new values.
    """
    if name:
        batch.name = name
    if frame_start is not None:
        batch.start_frame = frame_start - handle_start
    if frame_duration is not None:
        frame_duration += handle_start + handle_end
        batch.duration = frame_duration

    return batch

def write_node_metadata(node: flame.PyNode, data: Dict[str, Any]):
    """ Write AYON instance data into a node's note attribute."""
    payload = dict(data)
    payload[AYON_NOTE_MARKER] = True
    node.note = json.dumps(payload)
    node.note_collapsed = True