#!/usr/bin/env python3
# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""Upload tarballs to S3.

Uploads all .tar.gz files from the input directory to the tarballs/
subdirectory of the workflow output root in S3.

Example with ``--run-id 12345 --platform linux --release-type dev``:

    /tmp/tarballs/therock-dist-linux-gfx94X-dcgpu-7.10.0.tar.gz
      -> s3://therock-dev-artifacts/12345-linux/tarballs/therock-dist-linux-gfx94X-dcgpu-7.10.0.tar.gz

Usage:
    python build_tools/github_actions/upload_tarballs.py \\
        --input-tarballs-dir /tmp/tarballs \\
        --run-id 12345 --platform linux --release-type dev

    python build_tools/github_actions/upload_tarballs.py \\
        --input-tarballs-dir /tmp/tarballs \\
        --run-id 12345 --platform linux --release-type dev --dry-run

    # Local testing (no S3 credentials needed):
    python build_tools/github_actions/upload_tarballs.py \\
        --input-tarballs-dir /tmp/tarballs \\
        --run-id 12345 --platform linux \\
        --output-dir /tmp/upload-test
"""

import argparse
import json
import logging
import platform as platform_module
import sys
from pathlib import Path
from urllib.parse import quote

_BUILD_TOOLS_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_BUILD_TOOLS_DIR))

from _therock_utils.storage_backend import create_storage_backend
from _therock_utils.workflow_outputs import WorkflowOutputRoot
from github_actions_api import gha_set_output

logger = logging.getLogger(__name__)


def _tarball_url(output_root: WorkflowOutputRoot, name: str) -> str:
    return output_root.tarball(quote(name, safe="")).https_url


def _is_test_tarball(name: str) -> bool:
    return "-tests-" in name


def _select_multiarch_tarball_url(
    *,
    tarball_files: list[Path],
    output_root: WorkflowOutputRoot,
    platform: str,
) -> str | None:
    """Gets the default "multiarch" tarball URL from a list of tarball files.

    Note: this should look simpler once we drop the "multiarch" part of the
    file name, at which point the file will just be therock-dist-{platform}.
    """

    # Skip over "test" tarballs, only look at "base" tarballs.
    non_test_tarball_files = [f for f in tarball_files if not _is_test_tarball(f.name)]

    for f in non_test_tarball_files:
        name = f.name
        if name.startswith(f"therock-dist-{platform}-multiarch-"):
            return _tarball_url(output_root, name)

    return None


def run(
    input_tarballs_dir: Path,
    run_id: str,
    platform: str,
    release_type: str,
    output_dir: Path | None = None,
    dry_run: bool = False,
) -> int:
    tarballs_dir = input_tarballs_dir.resolve()
    if not tarballs_dir.is_dir():
        raise FileNotFoundError(f"Tarballs directory not found: {tarballs_dir}")

    tarball_files = sorted(tarballs_dir.glob("*.tar.gz"))
    if not tarball_files:
        raise FileNotFoundError(f"No .tar.gz files found in {tarballs_dir}")

    logger.info("Found %d tarballs in %s:", len(tarball_files), tarballs_dir)
    for f in tarball_files:
        size_mb = f.stat().st_size / (1024 * 1024)
        logger.info("  %s (%.1f MB)", f.name, size_mb)

    output_root = WorkflowOutputRoot.from_workflow_run(
        run_id=run_id,
        platform=platform,
        release_type=release_type or None,
    )

    dest = output_root.tarballs()
    logger.info("Destination: %s", dest.s3_uri)

    backend = create_storage_backend(
        staging_dir=output_dir,
        dry_run=dry_run,
    )
    count = backend.upload_directory(
        tarballs_dir,
        dest,
        include=["*.tar.gz"],
    )

    logger.info("Uploaded %d files", count)

    multiarch_tarball_url = _select_multiarch_tarball_url(
        tarball_files=tarball_files,
        output_root=output_root,
        platform=platform,
    )

    if not multiarch_tarball_url:
        raise ValueError(
            "No multiarch tarball URL was produced; check tarball naming and upload logic"
        )

    gha_set_output({"tarball_urls": json.dumps({"multiarch": multiarch_tarball_url})})

    return 0


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description="Upload tarballs to S3")
    parser.add_argument(
        "--input-tarballs-dir",
        type=Path,
        required=True,
        help="Directory containing .tar.gz tarballs to upload",
    )
    parser.add_argument("--run-id", required=True, help="Workflow run ID")
    parser.add_argument(
        "--platform",
        default=platform_module.system().lower(),
        choices=["linux", "windows"],
        help="Platform (default: current system)",
    )
    parser.add_argument(
        "--release-type",
        default="ci",
        help='Release type: "ci", "dev", "nightly", or "prerelease"',
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Output to local directory instead of S3 (for testing)",
    )
    parser.add_argument(
        "--dry-run", action="store_true", help="Print plan without uploading"
    )
    args = parser.parse_args(argv)

    return run(
        input_tarballs_dir=args.input_tarballs_dir,
        run_id=args.run_id,
        platform=args.platform,
        release_type=args.release_type,
        output_dir=args.output_dir,
        dry_run=args.dry_run,
    )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    sys.exit(main(sys.argv[1:]))