Merge astrophotography session folders into one combined folder (move + renumber lights, fan-out calibration frames)

merge_sessions.py

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Dinko Korunic <dinko.korunic@gmail.com>
# SPDX-License-Identifier: BSD-3-Clause
"""
Merge multiple astrophotography session folders into one combined folder.

For each session folder (e.g. Messier_101_2026-05-22_22-27-24), the script:
  * picks the single bias*.fits file and MOVES it to the destination as the
    first of 5 sequential names (bias0007.fits -> bias0007..bias0011.fits;
    bias0003.fits -> bias0003..bias0007.fits), then copies the moved file to
    the remaining 4 names
  * does the same for dark*.fits
  * does the same for flat*.fits (un-numbered flat.fits -> flat0001..0005)
  * MOVES every Light*.fits across all sources into the destination with
    sequential renumbering (Light0001.fits, Light0002.fits, ...)

The destination is never clobbered: if any target file already exists, the
script aborts BEFORE moving anything, so the source set stays intact.

Usage:
    ./merge_sessions.py SOURCE_FOLDER [SOURCE_FOLDER ...] [-d DEST] [-n]

If --dest is omitted, the destination name is derived from the common
"<name>_YYYY-MM" prefix of the source folders (e.g. Messier_101_2026-05) and
placed alongside them.
"""

from __future__ import annotations

import argparse
import re
import shutil
import sys
from pathlib import Path

SESSION_RE = re.compile(r"^(?P<base>.+)_(?P<ym>\d{4}-\d{2})-\d{2}_\d{2}-\d{2}-\d{2}$")
# (output_prefix, seed_glob) — capture index varies per session.
CALIBRATION_SOURCES = (
    ("bias", "bias*.fits"),
    ("dark", "dark*.fits"),
    ("flat", "flat*.fits"),
)
CALIBRATION_COPIES = 5
CALIBRATION_DEFAULT_START = 1  # for seeds without digits (e.g. flat.fits)
SEED_INDEX_RE = re.compile(r"(\d+)$")
LIGHT_GLOB = "Light*.fits"
LIGHT_WIDTH = 4  # Light0001.fits style


def derive_dest_name(folders: list[Path]) -> str | None:
    """Return e.g. 'Messier_101_2026-05' if all sources share a base+YYYY-MM."""
    keys = set()
    for f in folders:
        m = SESSION_RE.match(f.name)
        if not m:
            return None
        keys.add((m["base"], m["ym"]))
    if len(keys) != 1:
        return None
    base, ym = keys.pop()
    return f"{base}_{ym}"


def _is_inside(child: Path, parent: Path) -> bool:
    """True if child is parent or lives anywhere underneath it (resolved paths)."""
    try:
        child.relative_to(parent)
        return True
    except ValueError:
        return False


def find_seed(sources: list[Path], pattern: str) -> Path | None:
    """Return the single seed file matching pattern in the first source that has it."""
    for s in sources:
        matches = sorted(p for p in s.glob(pattern) if p.is_file())
        if not matches:
            continue
        if len(matches) > 1:
            names = ", ".join(m.name for m in matches)
            print(f"  warn: multiple matches for {pattern} in {s.name}: {names}; using {matches[0].name}", file=sys.stderr)
        return matches[0]
    return None


def parse_seed_index(seed_name: str, default: int) -> int:
    """Extract the trailing integer from a filename stem; fall back to default."""
    m = SEED_INDEX_RE.search(Path(seed_name).stem)
    return int(m.group(1)) if m else default


MOVE, COPY = "move", "copy"


def build_plan(sources: list[Path], dest: Path) -> list[tuple[Path, Path, str]]:
    """Return an ordered list of (src, dst, kind) operations.

    Order matters: each calibration MOVE precedes its dependent COPYs.
    """
    ops: list[tuple[Path, Path, str]] = []

    for prefix, pattern in CALIBRATION_SOURCES:
        seed = find_seed(sources, pattern)
        if seed is None:
            print(f"  warn: no match for {pattern} in any source, skipping {prefix}", file=sys.stderr)
            continue
        start = parse_seed_index(seed.name, default=CALIBRATION_DEFAULT_START)
        end = start + CALIBRATION_COPIES - 1
        print(f"  {prefix}: seed = {seed} (emitting {prefix}{start:04d}..{end:04d})")
        first_target = dest / f"{prefix}{start:04d}.fits"
        ops.append((seed, first_target, MOVE))
        print(f"    -> {first_target.name}")
        for i in range(start + 1, start + CALIBRATION_COPIES):
            target = dest / f"{prefix}{i:04d}.fits"
            ops.append((first_target, target, COPY))
            print(f"    -> {target.name}")

    lights: list[Path] = []
    for s in sorted(sources, key=lambda p: p.name):
        lights.extend(sorted(s.glob(LIGHT_GLOB)))

    if not lights:
        print("  warn: no Light*.fits files found in any source", file=sys.stderr)
    else:
        width = max(LIGHT_WIDTH, len(str(len(lights))))
        print(f"  lights: {len(lights)} files -> Light{1:0{width}d}.fits..Light{len(lights):0{width}d}.fits")
        for i, src in enumerate(lights, start=1):
            ops.append((src, dest / f"Light{i:0{width}d}.fits", MOVE))

    return ops


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("sources", nargs="+", type=Path, help="Session folders to merge")
    ap.add_argument("-d", "--dest", type=Path, help="Destination folder (auto-derived if omitted)")
    ap.add_argument("-n", "--dry-run", action="store_true", help="Print actions without moving any files")
    ap.add_argument("-v", "--verbose", action="store_true", help="Print every planned operation (src -> dst)")
    args = ap.parse_args()

    # Dedup: a doubled folder would crash mid-move with no rollback.
    seen: set[Path] = set()
    sources: list[Path] = []
    for s in args.sources:
        if not s.is_dir():
            print(f"error: not a directory: {s}", file=sys.stderr)
            return 2
        rs = s.resolve()
        if rs in seen:
            print(f"  warn: ignoring duplicate source {rs}", file=sys.stderr)
            continue
        seen.add(rs)
        sources.append(rs)

    if args.dest:
        dest = args.dest.resolve()
    else:
        name = derive_dest_name(sources)
        if not name:
            print("error: cannot derive destination name from sources; pass --dest", file=sys.stderr)
            return 2
        dest = sources[0].parent / name

    # Reject overlap in either direction, not just exact equality.
    for s in sources:
        if dest == s or _is_inside(dest, s) or _is_inside(s, dest):
            print(f"error: destination {dest} overlaps source {s}", file=sys.stderr)
            return 2

    print(f"Destination: {dest}")
    ops = build_plan(sources, dest)

    if args.verbose:
        for src, dst, kind in ops:
            print(f"    {kind} {src} -> {dst}")

    # Pre-flight: any existing dst aborts before we move anything.
    collisions = [dst for _, dst, _ in ops if dst.exists()]
    if collisions:
        print(f"error: {len(collisions)} destination file(s) already exist, aborting:", file=sys.stderr)
        for p in collisions[:10]:
            print(f"  {p}", file=sys.stderr)
        if len(collisions) > 10:
            print(f"  ... and {len(collisions) - 10} more", file=sys.stderr)
        return 1

    n_moves = sum(1 for _, _, k in ops if k == MOVE)
    n_copies = sum(1 for _, _, k in ops if k == COPY)
    if args.dry_run:
        print(f"Dry-run: {n_moves} moves, {n_copies} copies")
        return 0

    dest.mkdir(parents=True, exist_ok=True)
    for i, (src, dst, kind) in enumerate(ops, start=1):
        # Re-check: shutil.move/copy2 silently overwrite; closes TOCTOU window.
        if dst.exists():
            print(f"error: op {i}/{len(ops)}: {dst} appeared after pre-flight, aborting", file=sys.stderr)
            print(f"       {i - 1} operations completed; manual recovery required", file=sys.stderr)
            return 1
        try:
            if kind == MOVE:
                shutil.move(src, dst)
            else:
                shutil.copy2(src, dst)
        except OSError as e:
            print(f"error: op {i}/{len(ops)} failed: {kind} {src} -> {dst}: {e}", file=sys.stderr)
            print(f"       {i - 1} operations completed; manual recovery required", file=sys.stderr)
            return 1

    print(f"Done: {n_moves} moves, {n_copies} copies.")
    return 0


if __name__ == "__main__":
    sys.exit(main())