#!/usr/bin/env python3
"""omniopt_share - Share OmniOpt runs with a unified JSON/ZIP protocol.

Python rewrite of the former ``omniopt_share`` bash script, with a new
on-the-wire format:

  ``manifest.json``     - JSON document describing every file in the share.
  ``bundle.zip``        - ZIP archive with all the actual file contents.

The GUI server (``share_internal.php``) accepts these two as
``multipart/form-data`` parts named ``manifest`` and ``bundle``.  Because
every file is enumerated in the manifest, the server can:

  * verify each file's SHA-256 against the value declared in the manifest
  * reject any file whose ``archive_path`` would escape the share root
  * reject any file larger than a server-configured limit
  * apply its own whitelist independently of the client
  * add support for new file types purely by editing the manifest schema

CLI flags are 100 % backwards-compatible with the bash version.

Behavioural compatibility:
  * ``--update`` sends only files changed since the last successful share
    (cache in ``~/.oo_share_cache``); ``--force`` re-sends everything
  * large ``state_files/`` snapshots are skipped unless
    ``--include_state_files`` is given
  * ``--debug`` / ``--no_color`` / ``--timeout``
  * ``--username`` / ``--password`` / ``--outfile`` / ``--dont_send_singleruns``
  * Run dirs must contain ``results.csv`` (validated server-side)
  * Default ``BASEURL`` is the same as before, overridable via
    ``$HOME/.oo_base_url``
"""

from __future__ import annotations

import argparse
import datetime
import hashlib
import json
import os
import re
import shlex
import sys
import tempfile
import time
import urllib.error
import urllib.request
import zipfile
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple


SCRIPT_DIR = Path(__file__).resolve().parent

MANIFEST_SCHEMA_VERSION = "1.0"

# Server-imposed upper bound on the size of any single shared file
# (1 GiB).  The PHP side enforces the same value.
MAX_FILE_SIZE = 1 << 30

DEFAULT_BASE_URL = "https://imageseg.scads.de/omniax"

# Whitelisted file extensions for the top-level of a run dir.
TOP_LEVEL_EXTENSIONS = (".csv", ".txt", ".log", ".json")

# Extensionless files that have a fixed meaning in the GUI.
TOP_LEVEL_EXTENSIONLESS_FILES = ("outfile", "log", "progressbar", "install_errors", "trial_index_to_params")

# Log files live next to the run dirs, not inside them.  These suffixes map
# them to the names the GUI expects.
LOG_FILE_ALIASES = (
    ("outfile", ""),
    ("log", "_log"),
    ("progressbar", "_progressbar"),
    ("install_errors", "_install_errors"),
    ("trial_index_to_params", "_trial_index_to_param_logs"),
)

# Whitelisted file extensions under state_files/.
STATE_FILE_EXTENSIONS = (".csv", ".txt", ".json")

# Sub-extension patterns used for single_run files (e.g. 0.out, 0.err).
SINGLE_RUN_SUFFIXES = (".out", ".err")

# Files above this size are stored uncompressed in the bundle.  They are
# usually snapshots (json/csv) where deflate costs a lot of CPU but saves
# comparatively little.  Small text files are still compressed.
COMPRESS_MAX_SIZE = 4 * 1024 * 1024

# How long a single HTTP operation may block (seconds).  Large uploads can
# take a while, but should never hang forever.
DEFAULT_POST_TIMEOUT = 300

# Local cache of hashes from the last successful share.  Used to upload only
# changed files on ``--update`` instead of re-zipping/re-sending everything.
CACHE_VERSION = 1
CACHE_MAX_AGE_SECONDS = 30 * 24 * 60 * 60

USERNAME_FILE = Path.home() / ".oo_share_username"
DONT_ASK_FILE = Path.home() / ".oo_share_dont_ask"
BASE_URL_FILE = Path.home() / ".oo_base_url"


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------


def is_valid_username(name: str) -> bool:
    """Reject empty / whitespace-containing usernames."""
    return bool(name) and not any(c.isspace() for c in name)


def sanitize_archive_path(path: str) -> str:
    """Reject paths that try to escape the archive root.

    Accepted forms: ``results.csv`` or ``state_files/run_uuid``.
    Rejected: absolute paths, ``..`` traversal, null bytes, backslashes.
    """
    if not path or path != path.strip():
        raise ValueError(f"Invalid archive_path: {path!r}")
    if "\x00" in path:
        raise ValueError(f"Null byte in archive_path: {path!r}")
    if "\\" in path or path.startswith("/"):
        raise ValueError(f"Absolute or backslash path rejected: {path!r}")
    parts = path.split("/")
    if any(p in ("", ".", "..") for p in parts):
        raise ValueError(f"Unsafe path component in {path!r}")
    return path


# ---------------------------------------------------------------------------
# File discovery
# ---------------------------------------------------------------------------


def _find_log_dir(run_dir: Path) -> Optional[Path]:
    seen = set()
    starts = [run_dir.resolve(), Path.cwd().resolve(), SCRIPT_DIR]
    for start in starts:
        for base in (start, *start.parents):
            if base in seen:
                continue
            seen.add(base)
            candidate = base / "logs"
            if candidate.is_dir():
                return candidate
    return None


def _read_run_uuid(run_dir: Path) -> Optional[str]:
    path = run_dir / "state_files" / "run_uuid"
    try:
        value = path.read_text(encoding="utf-8").strip()
    except OSError:
        return None
    if not value or value in (".", ".."):
        return None
    if not re.fullmatch(r"[A-Za-z0-9_.-]+", value):
        return None
    return value


def collect_shareable_files(
    run_dir: str,
    send_single_runs: bool = True,
    include_state_files: bool = True,
) -> List[Dict[str, Any]]:
    """Walk ``run_dir`` and return the list of shareable files.

    Each entry is a dict with at least:

        ``local_path``   - absolute Path on disk
        ``archive_path`` - path inside the zip (always safe per
                           :func:`sanitize_archive_path`)
        ``name``         - logical name (filename without extension,
                           matching the legacy ``-F key=@file`` mapping)

    Extension / directory rules match the legacy bash version.
    ``include_state_files=False`` keeps only ``state_files/run_uuid`` (which
    the server needs to match updates) and skips large snapshots such as
    ``pd.json``/``checkpoint.json`` that the GUI does not read.
    """
    root = Path(run_dir)
    if not root.is_dir():
        return []

    out: List[Dict[str, Any]] = []
    seen_archive_paths = set()

    def add_entry(
        name: str, archive_path: str, local_path: Path
    ) -> None:
        archive_path = sanitize_archive_path(archive_path)
        if archive_path in seen_archive_paths:
            return
        seen_archive_paths.add(archive_path)
        out.append(
            {
                "name": name,
                "archive_path": archive_path,
                "local_path": local_path,
            }
        )

    # Top-level files
    for entry in sorted(root.iterdir()):
        if entry.is_file() and entry.suffix.lower() in TOP_LEVEL_EXTENSIONS:
            add_entry(entry.stem, entry.name, entry)

    for entry_name in TOP_LEVEL_EXTENSIONLESS_FILES:
        entry = root / entry_name
        if entry.is_file():
            add_entry(entry_name, entry_name, entry)

    # git_version (no specific extension requirement)
    git_version = root / "git_version"
    if git_version.is_file():
        add_entry("git_version", "git_version", git_version)

    # state_files/ - only csv/txt/json.  Large checkpoint/model snapshots are
    # skipped by default; the GUI only needs run_uuid (added separately).
    sf = root / "state_files"
    if sf.is_dir() and include_state_files:
        for entry in sorted(sf.iterdir()):
            if entry.is_file() and entry.suffix.lower() in STATE_FILE_EXTENSIONS:
                out.append(
                    {
                        "name": entry.stem,
                        "archive_path": sanitize_archive_path(
                            f"state_files/{entry.name}"
                        ),
                        "local_path": entry,
                    }
                )

    # state_files/run_uuid (no extension — always include for update matching)
    run_uuid_file = root / "state_files" / "run_uuid"
    if run_uuid_file.is_file():
        add_entry("run_uuid", "state_files/run_uuid", run_uuid_file)

    # Main log files are kept outside the run dir (logs/<run_uuid>...).
    log_dir = _find_log_dir(root)
    run_uuid_value = _read_run_uuid(root)
    if log_dir is not None and run_uuid_value is not None:
        for archive_name, suffix in LOG_FILE_ALIASES:
            local_path = log_dir / f"{run_uuid_value}{suffix}"
            if local_path.is_file():
                add_entry(archive_name, archive_name, local_path)

    # single_runs/<digit>/<file.{out,err}>  (optional)
    if send_single_runs:
        sr = root / "single_runs"
        if sr.is_dir():
            for run_folder in sorted(sr.iterdir()):
                if not run_folder.is_dir() or not run_folder.name.isdigit():
                    continue
                for entry in sorted(run_folder.iterdir()):
                    if (
                        entry.is_file()
                        and entry.suffix.lower() in SINGLE_RUN_SUFFIXES
                    ):
                        add_entry(
                            f"single_run_file_{run_folder.name}_{entry.name}",
                            f"single_runs/{run_folder.name}/{entry.name}",
                            entry,
                        )

    return out


# ---------------------------------------------------------------------------
# Single-run log filtering and large-upload diagnostics
# ---------------------------------------------------------------------------

LARGE_SHARE_WARNING_BYTES = 50 * 1024 * 1024
MAX_FAILED_LOG_BYTES = 50 * 1024 * 1024
LOG_INSPECT_BYTES = 256 * 1024
OMITTED_LOGS_FILENAME = "omitted_logs.json"


def _read_text_excerpt(path: Path, max_bytes: int = LOG_INSPECT_BYTES) -> str:
    """Read a bounded head+tail excerpt without loading a huge log fully."""
    try:
        size = path.stat().st_size
        with path.open("rb") as fh:
            if size <= 2 * max_bytes:
                data = fh.read()
            else:
                head = fh.read(max_bytes)
                fh.seek(-max_bytes, os.SEEK_END)
                tail = fh.read(max_bytes)
                data = head + b"\n... [middle omitted] ...\n" + tail
    except OSError:
        return ""
    return data.decode("utf-8", errors="replace")


def _single_run_log_indicates_failure(path: Path) -> bool:
    """Best-effort: should this single-run log be shared by default?

    Successful submitit logs contain an explicit completion marker.  Logs
    with a non-zero EXIT_CODE, OOM/time-limit markers or a traceback are
    kept.  Logs without any marker are kept too: they are usually still
    running or were interrupted, which is exactly when users need them.
    """
    try:
        if path.stat().st_size == 0:
            return False
    except OSError:
        return False

    text = _read_text_excerpt(path)
    match = re.search(r"EXIT_CODE:\s*(-?\d+)", text)
    if match:
        return int(match.group(1)) != 0
    if (
        "Job completed successfully" in text
        or "Exiting after successful completion" in text
    ):
        return False
    if re.search(
        r"(oom_kill event|CUDA out of memory|Traceback \(most recent call last\)|"
        r"DUE TO TIME LIMIT|Time limit exceeded|CANCELLED)",
        text,
        re.IGNORECASE,
    ):
        return True
    return True


def _split_single_run_logs(
    files: List[Dict[str, Any]],
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]:
    """Split files into (non-single-run, failed single-run, successful single-run)."""
    regular: List[Dict[str, Any]] = []
    by_folder: Dict[str, List[Dict[str, Any]]] = {}
    for f in files:
        archive_path = str(f["archive_path"])
        if not archive_path.startswith("single_runs/"):
            regular.append(f)
            continue
        parts = archive_path.split("/")
        folder = parts[1] if len(parts) > 1 else ""
        by_folder.setdefault(folder, []).append(f)

    failed: List[Dict[str, Any]] = []
    successful: List[Dict[str, Any]] = []
    for entries in by_folder.values():
        out_entries = [
            e for e in entries
            if Path(e["local_path"]).suffix.lower() == ".out"
        ]
        check_entries = out_entries or entries
        if any(
            _single_run_log_indicates_failure(Path(e["local_path"]))
            for e in check_entries
        ):
            failed.extend(entries)
        else:
            successful.extend(entries)
    return regular, failed, successful


def _limit_failed_logs(
    failed_logs: List[Dict[str, Any]],
    max_bytes: int,
    seed: str,
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], int]:
    """Deterministically sample failed-log groups until ``max_bytes`` is reached.

    Returns ``(selected, omitted, selected_size)``.  Grouping by run folder
    keeps a job's ``.out``/``.err`` together.  The order is pseudo-random but
    stable for the same run, so repeated live updates do not churn through a
    different subset each time.
    """
    groups: Dict[str, List[Dict[str, Any]]] = {}
    for f in failed_logs:
        archive_path = str(f["archive_path"])
        parts = archive_path.split("/")
        folder = "/".join(parts[:2])
        groups.setdefault(folder, []).append(f)

    ordered = sorted(
        groups.items(),
        key=lambda item: hashlib.sha256(
            f"{seed}\0{item[0]}".encode("utf-8")
        ).hexdigest(),
    )
    selected: List[Dict[str, Any]] = []
    omitted: List[Dict[str, Any]] = []
    selected_size = 0
    for _, entries in ordered:
        group_size = sum(_file_size(e) for e in entries)
        if group_size > max_bytes or selected_size + group_size > max_bytes:
            omitted.extend(entries)
            continue
        selected.extend(entries)
        selected_size += group_size
    return selected, omitted, selected_size


def _write_omitted_logs_metadata(
    run_dir: str,
    user_id: str,
    omitted_folders: Dict[str, str],
    single_run_logs_skipped: bool,
) -> Optional[Dict[str, Any]]:
    """Write a small GUI-readable notice about logs that were not shared."""
    run_abs = str(Path(run_dir).resolve())
    base_cmd = f"omniopt_share {shlex.quote(run_abs)} --update"
    if user_id:
        base_cmd += f" --username={shlex.quote(user_id)}"
    data = {
        "single_run_logs_skipped": bool(single_run_logs_skipped),
        "omitted": [
            {"folder": folder, "reason": reason}
            for folder, reason in sorted(omitted_folders.items())
        ],
        "commands": {
            "send_all_failed_logs": base_cmd + " --send_all_failed_logs",
            "send_all_single_runs": base_cmd + " --send_all_single_runs",
        },
    }

    path = Path(run_dir) / OMITTED_LOGS_FILENAME
    serialized = json.dumps(data, indent=2, sort_keys=True) + "\n"
    try:
        if path.is_file() and path.read_text(encoding="utf-8") == serialized:
            # Avoid touching mtime on every live-share tick.
            pass
        else:
            tmp = path.with_name(path.name + ".tmp")
            tmp.write_text(serialized, encoding="utf-8")
            os.replace(tmp, path)
    except OSError as e:
        print(f"Warning: could not write {path}: {e}", file=sys.stderr)
        return None
    return {
        "name": "omitted_logs",
        "archive_path": OMITTED_LOGS_FILENAME,
        "local_path": path,
    }


def _format_size(num_bytes: int) -> str:
    if num_bytes >= 1024 ** 3:
        return f"{num_bytes / (1024 ** 3):.2f} GiB"
    if num_bytes >= 1024 ** 2:
        return f"{num_bytes / (1024 ** 2):.1f} MiB"
    if num_bytes >= 1024:
        return f"{num_bytes / 1024:.1f} KiB"
    return f"{num_bytes} B"


def _file_size(entry: Dict[str, Any]) -> int:
    try:
        return Path(entry["local_path"]).stat().st_size
    except OSError:
        return 0


def _share_category(archive_path: str) -> str:
    if archive_path.startswith("single_runs/"):
        return "single-run logs"
    if archive_path.startswith("state_files/"):
        return "state files"
    name = Path(archive_path).name
    if archive_path.endswith(".log") or name in {
        "outfile", "log", "progressbar", "verbose_log.txt",
        "oo_errors.txt", "evaluation_errors.log",
    }:
        return "logs"
    if archive_path.endswith((".csv", ".txt")):
        return "results/tables"
    return "other"


def _print_size_breakdown(
    all_files: List[Dict[str, Any]],
    selected_files: List[Dict[str, Any]],
    successful_logs: List[Dict[str, Any]],
    send_all_single_runs: bool,
    omitted_failed_logs: Optional[List[Dict[str, Any]]] = None,
) -> None:
    all_size = sum(_file_size(f) for f in all_files)
    selected_size = sum(_file_size(f) for f in selected_files)
    if all_size < LARGE_SHARE_WARNING_BYTES and selected_size < LARGE_SHARE_WARNING_BYTES:
        return

    print(
        f"Share size check: {_format_size(all_size)} in {len(all_files)} file(s); "
        f"selected: {_format_size(selected_size)} in {len(selected_files)} file(s).",
        file=sys.stderr,
    )

    categories: Dict[str, List[int]] = {}
    for f in all_files:
        cat = _share_category(str(f["archive_path"]))
        categories.setdefault(cat, [0, 0])
        categories[cat][0] += _file_size(f)
        categories[cat][1] += 1
    for cat, (size, count) in sorted(
        categories.items(), key=lambda item: item[1][0], reverse=True
    ):
        extra = ""
        omitted_here = 0
        if cat == "single-run logs":
            if successful_logs and not send_all_single_runs:
                omitted_here += len(successful_logs)
            if omitted_failed_logs:
                omitted_here += len(omitted_failed_logs)
        if omitted_here:
            extra = f", {omitted_here} log file(s) omitted"
        print(
            f"  {_format_size(size):>10}  {cat} ({count} file(s){extra})",
            file=sys.stderr,
        )

    largest = sorted(all_files, key=_file_size, reverse=True)[:10]
    if largest:
        print("  Largest files:", file=sys.stderr)
        omitted_paths = {str(f["archive_path"]): "successful"
                         for f in successful_logs}
        for f in omitted_failed_logs or []:
            omitted_paths[str(f["archive_path"])] = "log cap"
        for f in largest:
            reason = omitted_paths.get(str(f["archive_path"]))
            marker = f" [omitted: {reason}]" if reason else ""
            print(
                f"    {_format_size(_file_size(f)):>10}  {f['archive_path']}{marker}",
                file=sys.stderr,
            )

    if successful_logs and not send_all_single_runs:
        print(
            "Only failed/interrupted single-run logs are sent by default. "
            "Use --send_all_single_runs to include successful logs, or "
            "--dont_send_singleruns to skip logs entirely.",
            file=sys.stderr,
        )
    if omitted_failed_logs:
        print(
            "Failed/interrupted logs were randomly sampled to stay under "
            f"{_format_size(MAX_FAILED_LOG_BYTES)}. Use --send_all_failed_logs "
            "to send all of them.",
            file=sys.stderr,
        )


# ---------------------------------------------------------------------------
# Manifest assembly & validation
# ---------------------------------------------------------------------------


def _sha256_of(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def build_manifest(
    *,
    run_dir: str,
    user_id: str,
    experiment_name: str,
    update: bool,
    update_uuid: Optional[str],
    password: Optional[str],
    send_single_runs: bool = True,
    include_state_files: bool = True,
    files: Optional[List[Dict[str, Any]]] = None,
) -> Dict[str, Any]:
    """Build a complete share manifest for ``run_dir``.

    ``files`` may be a pre-filtered subset (used for incremental updates).
    """
    if not is_valid_username(user_id):
        raise ValueError(f"Invalid username: {user_id!r}")

    if files is None:
        files = collect_shareable_files(
            run_dir,
            send_single_runs=send_single_runs,
            include_state_files=include_state_files,
        )
    entries: List[Dict[str, Any]] = []
    for f in files:
        local_path: Path = f["local_path"]
        size = local_path.stat().st_size
        if size > MAX_FILE_SIZE:
            raise ValueError(
                f"{local_path} is {size} bytes, larger than max {MAX_FILE_SIZE}"
            )
        entries.append(
            {
                "name": f["name"],
                "archive_path": f["archive_path"],
                "size": size,
                "sha256": _sha256_of(local_path),
                "content_type": _guess_content_type(local_path.name),
            }
        )

    # Always read the run_uuid if present — the server uses it to detect
    # whether this is a re-share of the same run (and updates in place).
    if update_uuid is None:
        run_uuid_path = Path(run_dir) / "state_files" / "run_uuid"
        if run_uuid_path.is_file():
            update_uuid = run_uuid_path.read_text(encoding="utf-8").strip()

    return {
        "schema_version": MANIFEST_SCHEMA_VERSION,
        "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "client_version": "0.97",
        "user_id": user_id,
        "experiment_name": experiment_name,
        "update": bool(update),
        "update_uuid": update_uuid,
        "password": password,
        "files": entries,
    }


def _guess_content_type(name: str) -> str:
    ext = Path(name).suffix.lower()
    return {
        ".csv": "text/csv",
        ".json": "application/json",
        ".txt": "text/plain",
        ".log": "text/plain",
        ".out": "text/plain",
        ".err": "text/plain",
    }.get(ext, "application/octet-stream")


def verify_manifest(manifest: Dict[str, Any]) -> str:
    """Return ``""`` on success, error message otherwise.

    The PHP side runs an equivalent check.  This function exists so the
    Python client can dry-validate its own manifest before uploading.
    """
    required = (
        "schema_version",
        "user_id",
        "experiment_name",
        "update",
        "update_uuid",
        "password",
        "files",
    )
    for key in required:
        if key not in manifest:
            return f"Missing manifest key: {key!r}"

    if manifest["schema_version"] != MANIFEST_SCHEMA_VERSION:
        return (
            f"Unsupported manifest schema_version "
            f"{manifest['schema_version']!r}, expected "
            f"{MANIFEST_SCHEMA_VERSION!r}"
        )
    if not is_valid_username(str(manifest["user_id"])):
        return f"Invalid user_id: {manifest['user_id']!r}"
    # `update` must be a real bool, not just truthy.
    if not isinstance(manifest["update"], bool):
        return (
            f"update must be bool, got "
            f"{type(manifest['update']).__name__}: {manifest['update']!r}"
        )

    files = manifest["files"]
    if not isinstance(files, list):
        return "files must be a list"
    if len(files) == 0:
        return "Manifest has no files (nothing to share)"
    for f in files:
        try:
            sanitize_archive_path(str(f["archive_path"]))
        except (KeyError, ValueError) as e:
            return f"Bad archive_path: {e}"
        try:
            size = int(f["size"])
        except (KeyError, ValueError, TypeError):
            return f"Bad size: {f.get('size')!r}"
        if isinstance(f["size"], bool) or not isinstance(f["size"], int):
            return f"Size must be int, got {type(f['size']).__name__}: {f['size']!r}"
        if size < 0 or size > MAX_FILE_SIZE:
            return f"Size out of range: {size}"
        sha = f.get("sha256")
        if not isinstance(sha, str) or len(sha) != 64:
            return f"Bad sha256: {sha!r}"
    return ""


# ---------------------------------------------------------------------------
# Incremental-update cache
# ---------------------------------------------------------------------------


def _cache_dir() -> Path:
    return Path(
        os.environ.get("OO_SHARE_CACHE_DIR", str(Path.home() / ".oo_share_cache"))
    )


def _cache_file_path(
    run_dir: str,
    user_id: str,
    experiment_name: str,
    base_url: str,
    update_uuid: Optional[str],
) -> Path:
    identity = update_uuid or str(Path(run_dir).resolve())
    raw = "\0".join([base_url, user_id, experiment_name, identity])
    key = hashlib.sha256(raw.encode("utf-8")).hexdigest()
    return _cache_dir() / f"{key}.json"


def _load_cache(path: Path) -> Dict[str, Dict[str, Any]]:
    if not path.is_file():
        return {}
    try:
        if time.time() - path.stat().st_mtime > CACHE_MAX_AGE_SECONDS:
            # Server-side shares expire; an old cache must not suppress the
            # re-upload of an otherwise unchanged run.
            return {}
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}
    if not isinstance(data, dict) or data.get("version") != CACHE_VERSION:
        return {}
    files = data.get("files")
    return files if isinstance(files, dict) else {}


def _save_cache(path: Path, files: Dict[str, Dict[str, Any]]) -> None:
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_name(path.name + ".tmp")
        tmp.write_text(
            json.dumps({"version": CACHE_VERSION, "files": files}, sort_keys=True),
            encoding="utf-8",
        )
        os.replace(tmp, path)
    except OSError as e:
        print(f"Warning: could not write share cache {path}: {e}", file=sys.stderr)


def _file_state(path: Path) -> Dict[str, int]:
    st = path.stat()
    return {
        "size": st.st_size,
        "mtime_ns": getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000)),
    }


def _split_changed_files(
    files: List[Dict[str, Any]],
    cache: Dict[str, Dict[str, Any]],
) -> Tuple[List[Dict[str, Any]], Dict[str, Dict[str, int]]]:
    """Return (files to upload, pre-upload state of every selected file)."""
    selected: List[Dict[str, Any]] = []
    pre_states: Dict[str, Dict[str, int]] = {}
    for f in files:
        local = Path(f["local_path"])
        try:
            state = _file_state(local)
        except OSError:
            continue
        old = cache.get(f["archive_path"])
        if (
            isinstance(old, dict)
            and old.get("size") == state["size"]
            and old.get("mtime_ns") == state["mtime_ns"]
        ):
            continue
        selected.append(f)
        pre_states[f["archive_path"]] = state
    return selected, pre_states


def _update_cache(
    cache_path: Path,
    old_cache: Dict[str, Dict[str, Any]],
    manifest: Dict[str, Any],
    pre_states: Dict[str, Dict[str, int]],
    source_dir: str,
    *,
    replace: bool,
    local_paths: Optional[Dict[str, Path]] = None,
) -> Dict[str, Dict[str, Any]]:
    new_cache: Dict[str, Dict[str, Any]] = {} if replace else dict(old_cache)
    for entry in manifest["files"]:
        archive_path = entry["archive_path"]
        if local_paths is not None and archive_path in local_paths:
            local = Path(local_paths[archive_path])
        else:
            local = Path(source_dir) / archive_path
        try:
            state = _file_state(local)
        except OSError:
            new_cache.pop(archive_path, None)
            continue
        before = pre_states.get(archive_path)
        if before and (
            before["size"] != state["size"]
            or before["mtime_ns"] != state["mtime_ns"]
        ):
            # The file changed while it was being zipped.  Do not remember it
            # as uploaded, so the next update sends the newer version.
            new_cache.pop(archive_path, None)
            continue
        new_cache[archive_path] = {
            "size": entry["size"],
            "mtime_ns": state["mtime_ns"],
            "sha256": entry["sha256"],
        }
    _save_cache(cache_path, new_cache)
    return new_cache


# ---------------------------------------------------------------------------
# Bundle writing
# ---------------------------------------------------------------------------


def write_bundle_zip(
    manifest: Dict[str, Any],
    zip_path: str | os.PathLike,
    *,
    source_dir: Optional[str | os.PathLike] = None,
    local_paths: Optional[Dict[str, Path]] = None,
) -> Path:
    """Write only the bundle ZIP to ``zip_path``.

    This is used by the CLI so it can keep the bundle in a single temporary
    file instead of creating an ``oo_share_*`` directory with
    ``manifest.json`` and ``bundle.zip`` inside it.
    """
    err = verify_manifest(manifest)
    if err:
        raise ValueError(f"Refusing to write invalid manifest: {err}")

    zip_path = Path(zip_path)
    zip_path.parent.mkdir(parents=True, exist_ok=True)

    base = Path(source_dir) if source_dir is not None else None
    # Hash and size are computed from the exact bytes that are written into
    # the zip.  This keeps manifest and bundle consistent even if a file is
    # modified (or truncated) between manifest creation and bundling, e.g.
    # while the run is still live and ``worker_usage.csv`` is being updated.
    with zipfile.ZipFile(zip_path, "w") as zf:
        for f in manifest["files"]:
            arc = sanitize_archive_path(f["archive_path"])
            if local_paths is not None and arc in local_paths:
                local = Path(local_paths[arc])
            else:
                local = (base / arc) if base is not None else Path(arc)
            if not local.is_file():
                raise FileNotFoundError(f"Bundle missing source file: {local}")

            info = zipfile.ZipInfo(arc)
            info.external_attr = 0o600 << 16
            info.compress_type = (
                zipfile.ZIP_DEFLATED
                if local.stat().st_size <= COMPRESS_MAX_SIZE
                else zipfile.ZIP_STORED
            )

            digest = hashlib.sha256()
            size = 0
            with local.open("rb") as src, zf.open(info, "w") as dst:
                for chunk in iter(lambda: src.read(1 << 20), b""):
                    size += len(chunk)
                    if size > MAX_FILE_SIZE:
                        raise ValueError(
                            f"{local} is larger than max {MAX_FILE_SIZE}"
                        )
                    digest.update(chunk)
                    dst.write(chunk)

            f["size"] = size
            f["sha256"] = digest.hexdigest()

    # Re-validate after the authoritative size/hash values have been filled
    # in, so callers can rely on the manifest matching the written bundle.
    err = verify_manifest(manifest)
    if err:
        raise ValueError(f"Refusing to write invalid manifest: {err}")

    return zip_path


def write_bundle(
    manifest: Dict[str, Any],
    out_dir: str | os.PathLike,
    *,
    source_dir: Optional[str | os.PathLike] = None,
    local_paths: Optional[Dict[str, Path]] = None,
) -> Tuple[Path, Path]:
    """Write ``manifest.json`` and ``bundle.zip`` to ``out_dir``.

    Returns ``(manifest_path, zip_path)``.  Kept for callers/tests that
    expect a directory layout; the CLI itself uses :func:`write_bundle_zip`
    and a single temporary file.
    """
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = out_dir / "manifest.json"
    zip_path = out_dir / "bundle.zip"

    write_bundle_zip(
        manifest, zip_path, source_dir=source_dir, local_paths=local_paths
    )
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True))

    return manifest_path, zip_path


# ---------------------------------------------------------------------------
# CLI / main
# ---------------------------------------------------------------------------


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="omniopt_share",
        description="Share OmniOpt runs.",
        add_help=False,
    )
    p.add_argument("--help", action="store_true")
    p.add_argument("--update", action="store_true",
                   help="Update a previously-shared run.")
    p.add_argument("--debug", action="store_true")
    p.add_argument("--no_color", action="store_true")
    p.add_argument("--force", action="store_true",
                   help="Ignore the incremental cache and re-send everything.")
    p.add_argument("--include_state_files", action="store_true",
                   help="Also share state_files/*.json/csv/txt (large "
                        "checkpoints/model snapshots; normally skipped).")
    p.add_argument("--timeout", type=float, default=DEFAULT_POST_TIMEOUT,
                   help=f"HTTP timeout in seconds (default: {DEFAULT_POST_TIMEOUT}).")
    p.add_argument("--username", default=None)
    p.add_argument("--password", default=None)
    p.add_argument(
        "--dont_send_singleruns", "--no_single_run_logs",
        dest="dont_send_singleruns", action="store_true",
        help="Do not share any single-run logs (alias: --no_single_run_logs).",
    )
    p.add_argument("--send_all_single_runs", action="store_true",
                   help="Share all single-run logs, including successful ones "
                        "(default: only failed/interrupted logs).")
    p.add_argument("--send_all_failed_logs", action="store_true",
                   help="Do not randomly sample failed logs when they exceed "
                        f"{_format_size(MAX_FAILED_LOG_BYTES)}.")
    p.add_argument("--outfile", default=None)
    p.add_argument("run_dirs", nargs="*", default=[],
                   help="One or more run directories to share.")
    return p


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    args = _build_parser().parse_args(list(argv))
    # Convenience: invert --dont_send_singleruns into a positive bool.
    args.send_single_runs = not args.dont_send_singleruns
    return args


def _print_help() -> None:
    print(
        "Usage: omniopt_share [OPTIONS] RUN_DIR\n"
        "\n"
        "Share your hyperparameter optimization results with others.\n"
        "\n"
        "Example:\n"
        "  omniopt_share runs/my_experiment/0\n"
        "\n"
        "Options:\n"
        "  --update                  Update a job that had this run-UUID.\n"
        "                            Only changed files are uploaded.\n"
        "  --debug                   Enable debug output\n"
        "  --no_color                Disable colored output\n"
        "  --force                   Ignore cache and re-send all files\n"
        "  --include_state_files     Also send state_files snapshots\n"
        "  --timeout=SECONDS         HTTP timeout (default: "
        f"{DEFAULT_POST_TIMEOUT:g})\n"
        "  --username=NAME           Override username\n"
        "  --password=PWD            Set password\n"
        "  --dont_send_singleruns    Skip single-run outputs completely\n"
        "                            (alias: --no_single_run_logs)\n"
        "  --send_all_single_runs    Include successful logs too\n"
        "                            (default: only failed logs)\n"
        "  --send_all_failed_logs    Do not sample failed logs down to the\n"
        f"                            {_format_size(MAX_FAILED_LOG_BYTES)} cap\n"
        "  --outfile=PATH            Path to the SLURM .out file\n"
        "  --help                    This help\n"
    )


def _read_username_from_cache() -> Optional[str]:
    if DONT_ASK_FILE.exists() and USERNAME_FILE.exists():
        cached = USERNAME_FILE.read_text(encoding="utf-8").strip()
        if is_valid_username(cached):
            return cached
    return None


def _resolve_username(args: argparse.Namespace) -> Optional[str]:
    if args.username:
        if not is_valid_username(args.username):
            print(f"Invalid username: {args.username!r}", file=sys.stderr)
            return None
        return args.username
    return _read_username_from_cache()


def _post_manifest(base_url: str, manifest: Dict[str, Any],
                   bundle_zip: Path, args: argparse.Namespace) -> Tuple[int, str]:
    """Upload ``manifest.json`` + ``bundle.zip`` as multipart POST.

    Returns ``(exit_code, body)``.  Falls back to a JSON POST when no
    bundle is present.
    """
    url = (
        f"{base_url}/share_internal.php"
        f"?user_id={manifest['user_id']}"
        f"&experiment_name={manifest['experiment_name']}"
    )
    if manifest["update"]:
        url += "&update=1"
    if manifest["update_uuid"]:
        url += f"&update_uuid={manifest['update_uuid']}"
    if manifest["password"]:
        url += f"&password={manifest['password']}"

    boundary = "----omniopt-share-" + os.urandom(8).hex()
    content_length, body = _build_multipart_stream(boundary, manifest, bundle_zip)
    req = urllib.request.Request(
        url,
        data=body,
        headers={
            "Content-Type": f"multipart/form-data; boundary={boundary}",
            "Content-Length": str(content_length),
        },
        method="POST",
    )
    timeout = float(getattr(args, "timeout", DEFAULT_POST_TIMEOUT) or DEFAULT_POST_TIMEOUT)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode("utf-8", errors="replace")
    except (urllib.error.URLError, TimeoutError, OSError) as e:
        return 7, f"Connection error: {e}"


def _build_multipart_stream(
    boundary: str, manifest: Dict[str, Any], bundle_zip: Path
) -> Tuple[int, Iterable[bytes]]:
    """Return ``(content_length, chunks)`` for a streaming multipart body.

    The zip is streamed from disk in 1 MiB chunks.  This keeps memory usage
    constant even for very large run directories (the old implementation read
    the whole ``bundle.zip`` into a single ``bytes`` object, then copied it
    again while building the request body).
    """
    prefix = (
        f"--{boundary}\r\n"
        'Content-Disposition: form-data; name="manifest"\r\n'
        "Content-Type: application/json\r\n"
        "\r\n"
    ).encode()
    manifest_bytes = json.dumps(manifest).encode()
    middle = (
        f"\r\n--{boundary}\r\n"
        'Content-Disposition: form-data; name="bundle"; filename="bundle.zip"\r\n'
        "Content-Type: application/zip\r\n"
        "\r\n"
    ).encode()
    suffix = f"\r\n--{boundary}--\r\n".encode()
    zip_size = bundle_zip.stat().st_size
    content_length = (
        len(prefix) + len(manifest_bytes) + len(middle) + zip_size + len(suffix)
    )

    def chunks() -> Iterable[bytes]:
        yield prefix
        yield manifest_bytes
        yield middle
        with bundle_zip.open("rb") as fh:
            for chunk in iter(lambda: fh.read(1 << 20), b""):
                yield chunk
        yield suffix

    return content_length, chunks()


def _resolve_base_url() -> str:
    env_url = os.environ.get("OMNIOPT_BASE_URL")
    if env_url:
        return env_url
    if BASE_URL_FILE.exists():
        return BASE_URL_FILE.read_text(encoding="utf-8").strip() or DEFAULT_BASE_URL
    return DEFAULT_BASE_URL
def _share_one_with(
    args: argparse.Namespace,
    user_id: str,
    run_dir: str,
    base_url: str,
    *,
    post_manifest: Callable[..., Tuple[int, str]],
) -> int:
    """Variant of :func:`_share_one` that takes an injected HTTP
    transport for testing."""
    if not os.path.isdir(run_dir):
        print(f"{run_dir} does not exist", file=sys.stderr)
        return 1

    experiment_name = os.path.basename(
        os.path.dirname(os.path.abspath(run_dir))
    )
    send_single_runs = not getattr(args, "dont_send_singleruns", False)
    send_all_single_runs = bool(getattr(args, "send_all_single_runs", False))
    send_all_failed_logs = bool(getattr(args, "send_all_failed_logs", False))
    include_state_files = bool(getattr(args, "include_state_files", False))

    run_uuid: Optional[str] = None
    run_uuid_path = Path(run_dir) / "state_files" / "run_uuid"
    if run_uuid_path.is_file():
        try:
            run_uuid = run_uuid_path.read_text(encoding="utf-8").strip() or None
        except OSError:
            run_uuid = None

    all_files = collect_shareable_files(
        run_dir,
        send_single_runs=send_single_runs,
        include_state_files=include_state_files,
    )
    if not all_files:
        print(f"Could not find any files in {run_dir}", file=sys.stderr)
        return 1

    successful_logs: List[Dict[str, Any]] = []
    omitted_failed_logs: List[Dict[str, Any]] = []
    files_for_selection = all_files
    if send_single_runs and not send_all_single_runs:
        regular_files, failed_logs, successful_logs = _split_single_run_logs(all_files)
        if not send_all_failed_logs:
            failed_size = sum(_file_size(f) for f in failed_logs)
            if failed_size > MAX_FAILED_LOG_BYTES:
                failed_logs, omitted_failed_logs, kept_size = _limit_failed_logs(
                    failed_logs,
                    MAX_FAILED_LOG_BYTES,
                    run_uuid or str(Path(run_dir).resolve()),
                )
                print(
                    f"Failed logs total {_format_size(failed_size)}; selected "
                    f"{_format_size(kept_size)} (cap "
                    f"{_format_size(MAX_FAILED_LOG_BYTES)}). "
                    "Use --send_all_failed_logs to send all of them.",
                    file=sys.stderr,
                )
        files_for_selection = regular_files + failed_logs

    omitted_folders: Dict[str, str] = {}
    for f in successful_logs:
        parts = str(f["archive_path"]).split("/")
        omitted_folders["/".join(parts[:2])] = "successful"
    for f in omitted_failed_logs:
        parts = str(f["archive_path"]).split("/")
        omitted_folders["/".join(parts[:2])] = "log cap"

    metadata_entry = _write_omitted_logs_metadata(
        run_dir,
        user_id,
        omitted_folders,
        single_run_logs_skipped=not send_single_runs,
    )
    all_files = [
        f for f in all_files if f["archive_path"] != OMITTED_LOGS_FILENAME
    ]
    files_for_selection = [
        f for f in files_for_selection
        if f["archive_path"] != OMITTED_LOGS_FILENAME
    ]
    if metadata_entry is not None:
        all_files.append(metadata_entry)
        files_for_selection.append(metadata_entry)

    if not files_for_selection:
        print(f"Could not find any files in {run_dir}", file=sys.stderr)
        return 1

    cache_path = _cache_file_path(
        run_dir, user_id, experiment_name, base_url, run_uuid
    )
    cache = _load_cache(cache_path)
    incremental = (
        bool(getattr(args, "update", False))
        and not bool(getattr(args, "force", False))
        and bool(cache)
    )

    if incremental:
        files_to_send, pre_states = _split_changed_files(files_for_selection, cache)
        if not files_to_send:
            _print_size_breakdown(
                all_files,
                files_to_send,
                successful_logs,
                send_all_single_runs,
                omitted_failed_logs,
            )
            print(
                f"Nothing changed in {run_dir}; skipping upload.",
                file=sys.stderr,
            )
            return 0
    else:
        files_to_send = files_for_selection
        pre_states = {}
        for f in files_for_selection:
            try:
                pre_states[f["archive_path"]] = _file_state(Path(f["local_path"]))
            except OSError:
                continue

    _print_size_breakdown(
        all_files,
        files_to_send,
        successful_logs,
        send_all_single_runs,
        omitted_failed_logs,
    )

    manifest = build_manifest(
        run_dir=run_dir,
        user_id=user_id,
        experiment_name=experiment_name,
        update=bool(getattr(args, "update", False)),
        update_uuid=run_uuid,
        password=getattr(args, "password", None),
        send_single_runs=send_single_runs,
        include_state_files=include_state_files,
        files=files_to_send,
    )

    if not manifest["files"]:
        print(f"Could not find any files in {run_dir}", file=sys.stderr)
        return 1

    # Use a single temporary ZIP file instead of a temporary directory.  The
    # old ``mkdtemp(prefix="oo_share_")`` created an ``oo_share_*`` folder on
    # every share (and left it behind if the process was interrupted).
    zip_fd, zip_name = tempfile.mkstemp(prefix="omniopt_share_", suffix=".zip")
    os.close(zip_fd)
    zip_path = Path(zip_name)
    try:
        local_paths = {
            f["archive_path"]: Path(f["local_path"]) for f in files_to_send
        }
        try:
            write_bundle_zip(
                manifest, zip_path, source_dir=run_dir, local_paths=local_paths
            )
        except ValueError as e:
            print(f"Refusing to share: {e}", file=sys.stderr)
            return 1

        total_size = sum(int(f["size"]) for f in manifest["files"])
        if getattr(args, "debug", False) or total_size >= (1 << 20):
            print(
                f"Uploading {len(manifest['files'])} file(s), "
                f"{total_size / (1 << 20):.1f} MiB"
                + (" (incremental update)" if incremental else "")
                + "...",
                file=sys.stderr,
            )

        status, body = post_manifest(base_url, manifest, zip_path, args)

        if status != 200 or "Error" in body:
            print(body, file=sys.stderr)
            if status == 7:
                print(
                    "Exit code 7 means the server was not reachable. "
                    "Are you online? Is the server properly started?",
                    file=sys.stderr,
                )
            return status if status != 200 else 1

        _update_cache(
            cache_path,
            cache,
            manifest,
            pre_states,
            run_dir,
            replace=not incremental,
            local_paths=local_paths,
        )
        print(body)
        return 0
    finally:
        try:
            zip_path.unlink()
        except FileNotFoundError:
            pass


def _share_one(
    args: argparse.Namespace,
    user_id: str,
    run_dir: str,
    base_url: str,
) -> int:
    """Share a single run_dir; returns 0 on success, non-zero on error."""
    return _share_one_with(
        args, user_id, run_dir, base_url,
        post_manifest=_post_manifest,
    )


def main(
    argv: Optional[Sequence[str]] = None,
    *,
    post_manifest: Optional[Callable[..., Tuple[int, str]]] = None,
) -> int:
    args = parse_args(argv if argv is not None else sys.argv[1:])
    if args.help:
        _print_help()
        return 0
    if not args.run_dirs:
        print("Please run with at least one folder as argument", file=sys.stderr)
        return 1

    user_id = _resolve_username(args)
    if not user_id:
        print("No username (use --username or run interactively once)", file=sys.stderr)
        return 1

    base_url = _resolve_base_url()

    last_rc = 0
    for run_dir in args.run_dirs:
        if post_manifest is not None:
            rc = _share_one_with(args, user_id, run_dir, base_url,
                                post_manifest=post_manifest)
        else:
            rc = _share_one(args, user_id, run_dir, base_url)
        if rc != 0:
            last_rc = rc
    return last_rc


# Legacy alias kept for tests / external callers that still use the
# single-dir API.
_share_one_run_dir = _share_one


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(0)
