#
# Proot-Distro - manage proot containers.
#
# Created by Sylirre <sylirre@termux.dev> for Termux project.
# Development assisted by Claude Code (https://claude.ai/code).
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

# Architecture: The push pipeline.
#
#   1. Load the cached manifest + image_config produced by `build`.
#      Re-canonicalize the config and verify its digest matches
#      manifest.config.digest — guards against cache corruption.
#   2. Exchange PD_DOCKER_AUTH for a Bearer token with `pull,push` scope.
#   3. For each layer: HEAD-probe (skip if present), else
#      POST /v2/<repo>/blobs/uploads/ + PUT <location>?digest=<digest>
#      streaming from the cached blob through a progress reader.
#   4. Upload the image config blob.
#   5. PUT the manifest under the tag.

import hashlib
import os
import sys
import urllib.error
import urllib.parse
import urllib.request

from proot_distro.message import (
    C, is_quiet, log_info,
)
from proot_distro.constants import PROGRAM_NAME
from proot_distro.progress import (
    clear_bar, fmt_size,
)
from proot_distro.helpers.docker.cache import (
    blob_present, load_manifest_cache, manifest_layers, open_required_layer,
)
from proot_distro.helpers.docker.media import (
    OCI_MANIFEST_MEDIA, canonical_json,
)
from proot_distro.helpers.download import retry_http
from proot_distro.helpers.docker.refs import parse_image_ref
from proot_distro.helpers.docker.transport import (
    auth_note,
    get_auth_token,
    opener,
    push_denied_msg,
    _ua,
)


def _resolve_upload_url(base: str, location: str) -> str:
    """Resolve the Location header from POST /v2/<repo>/blobs/uploads/."""
    if not location:
        raise RuntimeError(
            "Registry did not return an upload Location header."
        )
    if location.startswith(("http://", "https://")):
        return location
    if location.startswith("/"):
        return base + location
    return base.rstrip("/") + "/" + location


def _blob_exists(
    repo: str, digest: str, token: str, base: str, insecure: bool = False,
) -> bool:
    """Return True iff blob *digest* already exists on the registry."""
    url = f"{base}/v2/{repo}/blobs/{digest}"
    headers = {**_ua()}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    req = urllib.request.Request(url, method="HEAD", headers=headers)

    def _attempt():
        with opener(insecure).open(req) as resp:
            return 200 <= resp.status < 300

    try:
        return retry_http(_attempt, what=f"Checking {digest[:19]}")
    except urllib.error.HTTPError as exc:
        if exc.code == 404:
            return False
        raise


class _ProgressReader:
    """File wrapper that draws an upload progress bar as read() runs."""

    def __init__(self, fh, total: int, label: str):
        self._fh = fh
        self.total = total
        self.sent = 0
        self._label = label
        self._tty = sys.stderr.isatty() and not is_quiet()
        self._last_shown = 0

    def read(self, size=-1):
        data = self._fh.read(size)
        self.sent += len(data)
        if self._tty and (self.sent - self._last_shown >= 262144
                          or len(data) == 0):
            self._last_shown = self.sent
            pfx = f"{C['BLUE']}[{C['GREEN']}*{C['BLUE']}] {C['CYAN']}"
            if self.total:
                pct = min(self.sent * 100 // self.total, 100)
                bar = "#" * (pct // 5) + "-" * (20 - pct // 5)
                line = (
                    f"\r{pfx}{self._label}: [{bar}] {pct:3d}%  "
                    f"{fmt_size(self.sent)} / "
                    f"{fmt_size(self.total)}\033[K{C['RST']}"
                )
            else:
                line = (
                    f"\r{pfx}{self._label}: "
                    f"{fmt_size(self.sent)} uploaded...\033[K{C['RST']}"
                )
            sys.stderr.write(line)
            sys.stderr.flush()
        return data


def _upload_blob_bytes(
    repo: str, digest: str, data: bytes, token: str,
    base: str, insecure: bool = False,
) -> None:
    """Upload a small in-memory blob (POST + monolithic PUT).

    Per the OCI Distribution spec the PUT body uses
    'application/octet-stream'; the manifest's mediaType field describes
    the *content* of the blob and is set separately when the manifest
    itself is uploaded.
    """
    headers = {**_ua()}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    def _attempt():
        # A retry re-opens the upload session: the POST yields a fresh upload
        # Location, then the monolithic (content-addressed) PUT re-sends the
        # same bytes — safe to repeat.
        post_req = urllib.request.Request(
            f"{base}/v2/{repo}/blobs/uploads/",
            data=b"",
            method="POST",
            headers={**headers, "Content-Length": "0"},
        )
        with opener(insecure).open(post_req) as resp:
            location = resp.headers.get("Location", "")
        put_url = _resolve_upload_url(base, location)
        sep = "&" if "?" in put_url else "?"
        put_url = f"{put_url}{sep}digest={urllib.parse.quote(digest, safe='')}"
        put_req = urllib.request.Request(
            put_url,
            data=data,
            method="PUT",
            headers={
                **headers,
                "Content-Type": "application/octet-stream",
                "Content-Length": str(len(data)),
            },
        )
        with opener(insecure).open(put_req) as resp:
            if not 200 <= resp.status < 300:
                raise RuntimeError(
                    f"Blob upload failed for {digest}: HTTP {resp.status}"
                )

    retry_http(_attempt, what=f"Uploading {digest[:19]}")


def _upload_blob_fd(
    repo: str, digest: str, src_fd: int, token: str,
    base: str, insecure: bool = False, label: str = "",
) -> None:
    """Upload the blob behind *src_fd* (streamed, POST + monolithic PUT).

    A descriptor rather than a path, so the bytes published are the ones
    the digest check was made on — see cache.open_required_layer. Each
    attempt rewinds it, since a retry re-streams from the start.
    """
    headers = {**_ua()}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    size = os.fstat(src_fd).st_size

    def _attempt():
        # A retry re-opens the upload session and re-streams the file from the
        # start, so each attempt is a complete, self-contained upload.
        post_req = urllib.request.Request(
            f"{base}/v2/{repo}/blobs/uploads/",
            data=b"",
            method="POST",
            headers={**headers, "Content-Length": "0"},
        )
        with opener(insecure).open(post_req) as resp:
            location = resp.headers.get("Location", "")
        put_url = _resolve_upload_url(base, location)
        sep = "&" if "?" in put_url else "?"
        put_url = f"{put_url}{sep}digest={urllib.parse.quote(digest, safe='')}"
        try:
            os.lseek(src_fd, 0, os.SEEK_SET)
            with open(src_fd, "rb", closefd=False) as fh:
                reader = _ProgressReader(fh, size, label or digest[:19])
                put_req = urllib.request.Request(
                    put_url,
                    data=reader,
                    method="PUT",
                    headers={
                        **headers,
                        "Content-Type": "application/octet-stream",
                        "Content-Length": str(size),
                    },
                )
                with opener(insecure).open(put_req) as resp:
                    if not 200 <= resp.status < 300:
                        raise RuntimeError(
                            f"Blob upload failed for {digest}: HTTP {resp.status}"
                        )
        finally:
            clear_bar()

    retry_http(_attempt, what=f"Uploading {label or digest[:19]}")


def _put_manifest(
    repo: str, reference: str, body: bytes, media_type: str,
    token: str, base: str, insecure: bool = False,
) -> str:
    """PUT a manifest at <reference> (tag or digest). Returns the registry
    digest from the Docker-Content-Digest header, if provided."""
    url = f"{base}/v2/{repo}/manifests/{reference}"
    headers = {
        **_ua(),
        "Content-Type": media_type,
        "Content-Length": str(len(body)),
    }
    if token:
        headers["Authorization"] = f"Bearer {token}"
    req = urllib.request.Request(url, data=body, method="PUT", headers=headers)

    def _attempt():
        with opener(insecure).open(req) as resp:
            if not 200 <= resp.status < 300:
                raise RuntimeError(
                    f"Manifest upload failed: HTTP {resp.status}"
                )
            return resp.headers.get("Docker-Content-Digest", "")

    return retry_http(_attempt, what=f"Uploading manifest {reference}")


def _strip_private_keys(d: dict) -> dict:
    """Return a shallow copy of *d* without keys starting with '_'.

    `_get_manifest` stuffs the response Content-Type into `_ct` for
    internal use; that key must not be serialised back to the registry.
    """
    return {k: v for k, v in d.items() if not k.startswith("_")}


def push_image(image_ref: str, arch: str, insecure: bool = False) -> dict:
    """Push a built image (resolved from the manifest cache) to its registry.

    The image must have been produced by `proot-distro build` under
    exactly this *image_ref* and *arch* — `build` stores the manifest
    in MANIFEST_CACHE_DIR and the layer + config blobs in
    LAYER_CACHE_DIR using the same digests we transmit here.

    Registry traffic uses verified HTTPS unless *insecure* is set. With
    *insecure* a custom registry is reached over HTTPS with certificate
    verification disabled, falling back to plain HTTP when the registry only
    speaks HTTP (Docker Hub stays verified-HTTPS regardless). When enforcing
    HTTPS, an untrusted certificate or an HTTP-only registry surfaces a
    RuntimeError pointing the user at ``--allow-insecure`` — the same handling
    as the install command.
    """
    manifest, repo, image_config = load_manifest_cache(image_ref, arch)
    if manifest is None:
        raise RuntimeError(
            f"No cached manifest for '{image_ref}' ({arch}). Build image "
            f"first with: {PROGRAM_NAME} build -t {image_ref}"
        )

    # The cache entry is a file on disk, and on Termux that directory is
    # guest-writable; a manifest that is not the shape the loop below
    # subscripts is a traceback rather than a message without this.
    layers = manifest_layers(manifest, image_ref)
    if not layers:
        raise RuntimeError(
            f"Cached manifest for '{image_ref}' has no filesystem layers."
        )

    # Asked of the layer cache's own descriptor, not of a composed
    # path: `oci_layers` is a name a guest can leave behind as a
    # symlink, and os.path.isfile() follows one -- so a host directory's
    # files answered for the blobs this push is about to upload.
    missing = [
        layer["digest"] for layer in layers
        if not blob_present(layer["digest"])
    ]
    if missing:
        raise RuntimeError(
            f"Cannot push '{image_ref}': {len(missing)} layer blob(s) are "
            f"missing from the local cache (first missing: {missing[0]}). "
            f"Rebuild the image to repopulate the cache."
        )

    registry, _, tag = parse_image_ref(image_ref)

    # Re-canonicalize the image config and verify the digest. The manifest
    # carries config.digest, which the registry verifies against the bytes
    # we PUT. Round-tripping the dict through json.dump+json.load preserves
    # all keys, so the canonical form is reproducible here.
    config_bytes = canonical_json(image_config)
    cached_config = manifest.get("config", {})
    expected_cfg_digest = (
        cached_config.get("digest", "")
        if isinstance(cached_config, dict) else ""
    )
    if not isinstance(expected_cfg_digest, str):
        expected_cfg_digest = ""
    actual_cfg_digest = "sha256:" + hashlib.sha256(config_bytes).hexdigest()
    if expected_cfg_digest != actual_cfg_digest:
        raise RuntimeError(
            f"Image config digest mismatch (cached manifest expects "
            f"{expected_cfg_digest}, regenerated bytes hash to "
            f"{actual_cfg_digest}). The local cache appears corrupted; "
            f"rebuild the image."
        )

    log_info(f"Authenticating with registry{auth_note()}...")
    try:
        # Resolve the scheme/base for this registry. Under --allow-insecure a
        # bad certificate is tolerated and an HTTP-only registry falls back to
        # http://; otherwise a cert/plaintext failure raises a RuntimeError
        # pointing at --allow-insecure (handled inside get_auth_token).
        token, base = get_auth_token(
            repo, registry, actions="pull,push", insecure=insecure,
        )
    except urllib.error.HTTPError as exc:
        if exc.code in (401, 403):
            raise RuntimeError(push_denied_msg(image_ref, exc.code)) from exc
        raise

    n_layers = len(layers)
    bytes_uploaded = 0

    for i, layer in enumerate(layers):
        digest = layer["digest"]
        short_id = digest.split(":")[-1][:12]

        try:
            if _blob_exists(repo, digest, token, base, insecure):
                log_info(f"{short_id}: Layer {i + 1}/{n_layers} already "
                         f"exists on registry, skipping upload.")
                continue

            # The registry hashes what we PUT and rejects a mismatch with
            # an opaque 400; checking here names the real problem, and
            # keeps a cache blob that was swapped for someone else's
            # content from being published under this image's digest.
            # The upload then streams the descriptor the check was made
            # on, so the name cannot be re-pointed in between.
            src_fd = open_required_layer(digest)
            try:
                size = os.fstat(src_fd).st_size
                log_info(f"{short_id}: Uploading layer {i + 1}/{n_layers} "
                         f"({fmt_size(size)})...")
                _upload_blob_fd(
                    repo, digest, src_fd, token, base, insecure,
                    label=short_id,
                )
            finally:
                os.close(src_fd)
            bytes_uploaded += size
        except urllib.error.HTTPError as exc:
            if exc.code in (401, 403):
                raise RuntimeError(
                    push_denied_msg(image_ref, exc.code)
                ) from exc
            raise

    cfg_short = expected_cfg_digest.split(":")[-1][:12]
    try:
        if _blob_exists(repo, expected_cfg_digest, token, base, insecure):
            log_info(f"{cfg_short}: Image config already exists on "
                     f"registry, skipping upload.")
        else:
            log_info(f"{cfg_short}: Uploading image config "
                     f"({fmt_size(len(config_bytes))})...")
            _upload_blob_bytes(
                repo, expected_cfg_digest, config_bytes, token, base, insecure,
            )
            bytes_uploaded += len(config_bytes)
    except urllib.error.HTTPError as exc:
        if exc.code in (401, 403):
            raise RuntimeError(push_denied_msg(image_ref, exc.code)) from exc
        raise

    manifest_media = manifest.get("mediaType") or OCI_MANIFEST_MEDIA
    manifest_bytes = canonical_json(_strip_private_keys(manifest))
    log_info(f"Uploading manifest for tag '{tag}' "
             f"({fmt_size(len(manifest_bytes))})...")
    try:
        registry_digest = _put_manifest(
            repo, tag, manifest_bytes, manifest_media, token, base, insecure,
        )
    except urllib.error.HTTPError as exc:
        if exc.code in (401, 403):
            raise RuntimeError(push_denied_msg(image_ref, exc.code)) from exc
        raise
    bytes_uploaded += len(manifest_bytes)

    return {
        "manifest_digest": registry_digest,
        "bytes_uploaded": bytes_uploaded,
        "registry": registry or "docker.io",
        "repo": repo,
        "tag": tag,
    }
