#!/bin/python3
"""
Parallel rsync launcher.

Requires rsync on PATH (e.g. via MSYS/Cygwin on Windows).
Runs multiple rsync invocations in parallel using either top-level entries
or balanced recursive file chunks.
"""

# Usage examples:
#
#   Basic sync (recursive-files mode, auto worker count):
#     python pyrsync.py /src/dir user@host:/dest/
#
#   Auto-tune workers based on file size and local vs remote destination:
#     python pyrsync.py /src/dir user@host:/dest/ --auto-tune
#
#   Limit to a specific number of parallel workers:
#     python pyrsync.py /src/dir /local/dest/ -w 8
#
#   Use top-level mode (one rsync process per immediate source entry):
#     python pyrsync.py /src/dir /local/dest/ --mode top-level
#
#   Dry run (show what would be transferred without doing it):
#     python pyrsync.py /src/dir user@host:/dest/ --dry-run
#
#   Verbose output (shows command templates, task list preview, and elapsed time):
#     python pyrsync.py /src/dir /local/dest/ -v
#
#   Show rsync output for all tasks:
#     python pyrsync.py /src/dir /local/dest/ --show-rsync-output
#
#   Show rsync output only when a task fails:
#     python pyrsync.py /src/dir user@host:/dest/ --show-rsync-output-on-failure
#
#   Full-featured example (auto-tune, verbose, show failures):
#     python pyrsync.py /src/dir user@host:/dest/ --auto-tune -v --show-rsync-output-on-failure
#
#   Override rsync flags entirely:
#     python pyrsync.py /src/dir /local/dest/ -r "-az --checksum"


import argparse
import os
import re
import shlex
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import List, Sequence, Tuple

# DEFAULT_RSYNC_ARGS = ["-av", "--partial", "--compress", "--info=progress2"]
DEFAULT_RSYNC_ARGS = ["-a", "--partial", "--info=stats2"]


def build_file_list(src: str) -> List[Path]:
    p = Path(src)
    if p.is_dir():
        # sync each top-level entry separately (files and immediate subdirs)
        return [x for x in p.iterdir()]
    # allow glob patterns
    if any(ch in src for ch in ["*", "?"]):
        return list(Path(".").glob(src))
    # single file
    return [p]


def build_recursive_file_list(src_dir: Path) -> List[Path]:
    """Return all files under src_dir as relative paths for --files-from."""
    return [p.relative_to(src_dir) for p in src_dir.rglob("*") if p.is_file()]


def get_sized_files(src_dir: Path, rel_files: Sequence[Path]) -> List[Tuple[int, Path]]:
    """Return (size_bytes, rel_path) tuples with best-effort size lookup."""
    sized_files: List[Tuple[int, Path]] = []
    for rel_path in rel_files:
        abs_path = src_dir / rel_path
        try:
            size = abs_path.stat().st_size
        except OSError:
            size = 0
        sized_files.append((size, rel_path))
    return sized_files


def partition_by_size(sized_files: Sequence[Tuple[int, Path]], workers: int) -> List[List[Path]]:
    """Greedy size-balanced partitioning to reduce worker skew."""
    buckets: List[List[Path]] = [[] for _ in range(workers)]
    bucket_sizes = [0] * workers

    # Largest files first gives better load balancing across workers.
    sorted_files = sorted(sized_files, key=lambda x: x[0], reverse=True)
    for size, rel_path in sorted_files:
        idx = min(range(workers), key=lambda i: bucket_sizes[i])
        buckets[idx].append(rel_path)
        bucket_sizes[idx] += size

    return [b for b in buckets if b]


def detect_remote_destination(dest: str) -> bool:
    """Heuristic detection of rsync remote destinations."""
    if dest.startswith("rsync://"):
        return True
    if "::" in dest:
        return True
    if re.match(r"^[a-zA-Z]:[\\/]", dest):
        return False
    if re.match(r"^\\\\", dest):
        return False
    return bool(re.match(r"^[^:/\\]+@?[^:/\\]*:[^/].*", dest))


def format_bytes(num_bytes: float) -> str:
    units = ["B", "KB", "MB", "GB", "TB"]
    value = float(num_bytes)
    for unit in units:
        if value < 1024.0 or unit == units[-1]:
            return f"{value:.1f}{unit}"
        value /= 1024.0
    return f"{num_bytes:.1f}B"


def auto_tune_workers(
    total_units: int,
    avg_size_bytes: float,
    is_remote: bool,
    cpu_count: int,
    configured_workers: int,
) -> int:
    """Pick worker count from workload size profile and destination type."""
    cpu = max(1, cpu_count)

    if avg_size_bytes < 1 * 1024 * 1024:
        base = min(32, cpu * 4)
    elif avg_size_bytes < 32 * 1024 * 1024:
        base = min(16, cpu * 2)
    elif avg_size_bytes < 256 * 1024 * 1024:
        base = min(8, cpu)
    else:
        base = min(4, max(2, cpu // 2))

    if is_remote:
        base = max(2, min(base, 8))
    else:
        base = min(base + 2, 24)

    tuned = max(1, min(base, total_units))
    return min(tuned, configured_workers)


def maybe_print_task_output(label: str, out: str, err: str, show_output: bool, show_on_fail_only: bool, code: int) -> None:
    if not show_output:
        return
    if show_on_fail_only and code == 0:
        return

    if out.strip():
        print(f"--- stdout: {label} ---")
        print(out.rstrip())
    if err.strip():
        print(f"--- stderr: {label} ---", file=sys.stderr)
        print(err.rstrip(), file=sys.stderr)


def run_rsync(item: Path, dest: str, rsync_args: List[str], dry_run: bool) -> Tuple[Path, int, str, str]:
    cmd = ["rsync", *rsync_args]
    if dry_run:
        cmd.append("--dry-run")
    # preserve relative path by sending the item name, not its parent
    cmd += [str(item), dest]
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True)
        return item, proc.returncode, proc.stdout, proc.stderr
    except FileNotFoundError:
        return item, 127, "", "rsync not found on PATH"


def run_rsync_files_from(
    src_dir: Path,
    rel_files: Sequence[Path],
    dest: str,
    rsync_args: List[str],
    dry_run: bool,
) -> Tuple[str, int, str, str]:
    cmd = ["rsync", *rsync_args]
    if dry_run:
        cmd.append("--dry-run")

    cmd += ["--from0", "--files-from=-", f"{src_dir}{os.sep}", dest]

    payload = "\0".join(str(p) for p in rel_files) + "\0"
    try:
        proc = subprocess.run(cmd, input=payload, capture_output=True, text=True)
        label = f"chunk({len(rel_files)} files)"
        return label, proc.returncode, proc.stdout, proc.stderr
    except FileNotFoundError:
        return "chunk(0 files)", 127, "", "rsync not found on PATH"


def main():
    ap = argparse.ArgumentParser(description="Transfer files using rsync in parallel.")
    ap.add_argument("src", help="Source file, directory or glob pattern")
    ap.add_argument("dest", help="Rsync destination (local path or user@host:/path)")
    ap.add_argument(
        "-w",
        "--workers",
        type=int,
        default=min(8, (os.cpu_count() or 4)),
        help="Parallel rsync workers (default: min(8, CPU count))",
    )
    ap.add_argument(
        "--mode",
        choices=["top-level", "recursive-files"],
        default="recursive-files",
        help=(
            "top-level: one rsync per immediate src entry; "
            "recursive-files: split full file tree into balanced rsync chunks"
        ),
    )
    ap.add_argument("-r", "--rsync-args", default=" ".join(DEFAULT_RSYNC_ARGS),
                    help="Additional rsync args (quoted). Default: -a --partial --info=stats2")
    ap.add_argument(
        "--auto-tune",
        action="store_true",
        help=(
            "Auto-tune worker count from average file size and destination type. "
            "Uses --workers as an upper bound."
        ),
    )
    ap.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Show detailed progress, including task start/completion and rsync command lines.",
    )
    ap.add_argument(
        "--show-rsync-output",
        action="store_true",
        help="Print rsync stdout/stderr for each task (can be noisy).",
    )
    ap.add_argument(
        "--show-rsync-output-on-failure",
        action="store_true",
        help="Print rsync stdout/stderr only for failed tasks.",
    )
    ap.add_argument("--dry-run", action="store_true", help="Don't transfer, just show what would happen")
    args = ap.parse_args()

    if args.workers < 1:
        print("--workers must be >= 1", file=sys.stderr)
        sys.exit(2)

    rsync_args = shlex.split(args.rsync_args)
    results = []
    started_at = time.perf_counter()

    if args.show_rsync_output and args.show_rsync_output_on_failure:
        print(
            "Both --show-rsync-output and --show-rsync-output-on-failure set; "
            "showing output for all tasks.",
            file=sys.stderr,
        )

    show_output = args.show_rsync_output or args.show_rsync_output_on_failure
    show_on_fail_only = args.show_rsync_output_on_failure and not args.show_rsync_output

    src_path = Path(args.src)
    is_remote_dest = detect_remote_destination(args.dest)
    cpu_count = os.cpu_count() or 4
    if args.mode == "recursive-files" and src_path.is_dir():
        rel_files = build_recursive_file_list(src_path)
        if not rel_files:
            print("No files found for:", args.src, file=sys.stderr)
            sys.exit(2)

        sized_files = get_sized_files(src_path, rel_files)
        avg_size = sum(size for size, _ in sized_files) / max(1, len(sized_files))
        if args.auto_tune:
            workers = auto_tune_workers(len(rel_files), avg_size, is_remote_dest, cpu_count, args.workers)
            dest_kind = "remote" if is_remote_dest else "local"
            print(
                f"Auto-tuned workers={workers} ({dest_kind} dest, avg file {format_bytes(avg_size)})"
            )
        else:
            workers = min(args.workers, len(rel_files))

        chunks = partition_by_size(sized_files, workers)
        print(f"Starting {len(chunks)} rsync chunk task(s) with {workers} worker(s)...")

        if args.verbose:
            cmd_preview = ["rsync", *rsync_args]
            if args.dry_run:
                cmd_preview.append("--dry-run")
            cmd_preview += ["--from0", "--files-from=-", f"{src_path}{os.sep}", args.dest]
            print("Command template:", " ".join(shlex.quote(part) for part in cmd_preview))

        with ThreadPoolExecutor(max_workers=workers) as ex:
            futures = {
                ex.submit(run_rsync_files_from, src_path, chunk, args.dest, rsync_args, args.dry_run): chunk
                for chunk in chunks
            }
            for fut in as_completed(futures):
                label, code, out, err = fut.result()
                results.append((label, code, out, err))
                status = "OK" if code == 0 else f"ERR({code})"
                print(f"[{status}] {label}")
                maybe_print_task_output(label, out, err, show_output, show_on_fail_only, code)
    else:
        items = build_file_list(args.src)
        if not items:
            print("No files found for:", args.src, file=sys.stderr)
            sys.exit(2)

        size_samples = []
        for item in items:
            try:
                if item.is_file():
                    size_samples.append(item.stat().st_size)
                else:
                    size_samples.append(64 * 1024 * 1024)
            except OSError:
                size_samples.append(0)

        avg_size = sum(size_samples) / max(1, len(size_samples))
        if args.auto_tune:
            workers = auto_tune_workers(len(items), avg_size, is_remote_dest, cpu_count, args.workers)
            dest_kind = "remote" if is_remote_dest else "local"
            print(
                f"Auto-tuned workers={workers} ({dest_kind} dest, avg unit {format_bytes(avg_size)})"
            )
        else:
            workers = min(args.workers, len(items))

        print(f"Starting {len(items)} rsync task(s) with {workers} worker(s)...")

        if args.verbose:
            print("Task list preview:")
            preview_limit = 10
            for idx, item in enumerate(items[:preview_limit], start=1):
                print(f"  {idx}. {item}")
            if len(items) > preview_limit:
                print(f"  ... and {len(items) - preview_limit} more")

        with ThreadPoolExecutor(max_workers=workers) as ex:
            futures = {}
            for item in items:
                if args.verbose:
                    cmd_preview = ["rsync", *rsync_args]
                    if args.dry_run:
                        cmd_preview.append("--dry-run")
                    cmd_preview += [str(item), args.dest]
                    print("Scheduling:", " ".join(shlex.quote(part) for part in cmd_preview))
                futures[ex.submit(run_rsync, item, args.dest, rsync_args, args.dry_run)] = item

            for fut in as_completed(futures):
                item, code, out, err = fut.result()
                results.append((str(item), code, out, err))
                status = "OK" if code == 0 else f"ERR({code})"
                print(f"[{status}] {item}")
                maybe_print_task_output(str(item), out, err, show_output, show_on_fail_only, code)

    # summary
    failed = [r for r in results if r[1] != 0]
    elapsed = time.perf_counter() - started_at
    print(f"\nCompleted: {len(results)} tasks, {len(failed)} failed in {elapsed:.2f}s.")
    if failed:
        print("Failures detail:")
        for label, code, out, err in failed:
            print(f"--- {label} (exit {code}) ---")
            if out:
                print(out.strip())
            if err:
                print(err.strip())

    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()