#!/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`` / ``--force`` / ``--debug`` / ``--no_color``
  * ``--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 shutil
import sys
import tempfile
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")

# 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")

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 collect_shareable_files(
    run_dir: str, send_single_runs: 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.
    """
    root = Path(run_dir)
    if not root.is_dir():
        return []

    out: List[Dict[str, Any]] = []

    # Top-level files
    for entry in sorted(root.iterdir()):
        if entry.is_file() and entry.suffix.lower() in TOP_LEVEL_EXTENSIONS:
            name = entry.stem
            out.append(
                {
                    "name": name,
                    "archive_path": sanitize_archive_path(entry.name),
                    "local_path": entry,
                }
            )

    # git_version (no specific extension requirement)
    git_version = root / "git_version"
    if git_version.is_file():
        out.append(
            {
                "name": "git_version",
                "archive_path": "git_version",
                "local_path": git_version,
            }
        )

    # state_files/ - only csv/txt/json
    sf = root / "state_files"
    if sf.is_dir():
        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,
                    }
                )

    # 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
                    ):
                        out.append(
                            {
                                "name": f"single_run_file_{run_folder.name}_{entry.name}",
                                "archive_path": sanitize_archive_path(
                                    f"single_runs/{run_folder.name}/{entry.name}"
                                ),
                                "local_path": entry,
                            }
                        )

    return out


# ---------------------------------------------------------------------------
# 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,
) -> Dict[str, Any]:
    """Build a complete share manifest for ``run_dir``."""
    if not is_valid_username(user_id):
        raise ValueError(f"Invalid username: {user_id!r}")

    files = collect_shareable_files(run_dir, send_single_runs=send_single_runs)
    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),
            }
        )

    # If update is requested and the run dir has a run_uuid, copy it
    if update_uuid is None and update:
        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 ""


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


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

    Returns ``(manifest_path, zip_path)``.

    If ``source_dir`` is provided, every file's local path is resolved
    as ``source_dir / archive_path``.  Otherwise, ``archive_path`` is
    treated as the local path (relative or absolute).
    """
    err = verify_manifest(manifest)
    if err:
        raise ValueError(f"Refusing to write invalid manifest: {err}")

    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"

    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True))

    base = Path(source_dir) if source_dir is not None else None
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for f in manifest["files"]:
            arc = sanitize_archive_path(f["archive_path"])
            local = (base / arc) if base is not None else Path(arc)
            if not local.is_file():
                raise FileNotFoundError(f"Bundle missing source file: {local}")
            zf.write(local, arcname=arc)

    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 cache and re-send everything.")
    p.add_argument("--username", default=None)
    p.add_argument("--password", default=None)
    p.add_argument("--dont_send_singleruns", action="store_true")
    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"
        "  --debug                   Enable debug output\n"
        "  --no_color                Disable colored output\n"
        "  --force                   Ignore cache\n"
        "  --username=NAME           Override username\n"
        "  --password=PWD            Set password\n"
        "  --dont_send_singleruns    Skip single-run outputs\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()
    body = _build_multipart(boundary, manifest, bundle_zip)
    req = urllib.request.Request(
        url,
        data=body,
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=60) 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 as e:
        return 7, f"Connection error: {e}"


def _build_multipart(
    boundary: str, manifest: Dict[str, Any], bundle_zip: Path
) -> bytes:
    """Build a ``multipart/form-data`` body with two parts: manifest + bundle."""
    lines: List[bytes] = []
    crlf = b"\r\n"

    # Part 1: manifest
    lines.append(f"--{boundary}".encode())
    lines.append(b'Content-Disposition: form-data; name="manifest"')
    lines.append(b"Content-Type: application/json")
    lines.append(b"")
    lines.append(json.dumps(manifest).encode())

    # Part 2: bundle
    lines.append(f"--{boundary}".encode())
    lines.append(
        f'Content-Disposition: form-data; name="bundle"; filename="bundle.zip"'.encode()
    )
    lines.append(b"Content-Type: application/zip")
    lines.append(b"")
    lines.append(bundle_zip.read_bytes())

    # Close
    lines.append(f"--{boundary}--".encode())
    lines.append(b"")
    return crlf.join(lines)


def _resolve_base_url() -> str:
    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 args.dont_send_singleruns

    manifest = build_manifest(
        run_dir=run_dir,
        user_id=user_id,
        experiment_name=experiment_name,
        update=args.update,
        update_uuid=None,
        password=args.password,
        send_single_runs=send_single_runs,
    )

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

    out_dir = Path(tempfile.mkdtemp(prefix="oo_share_"))
    try:
        _, zip_path = write_bundle(manifest, out_dir, source_dir=run_dir)
    except ValueError as e:
        print(f"Refusing to share: {e}", file=sys.stderr)
        return 1

    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

    print(body)
    return 0


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)
