#!/usr/bin/env python3
"""Immutable application archives through the platform's primary storage.

Credentials are owner-delivered environment fields, never read from another
workload's Secret. CLI transfer previews by default. Output is redacted.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import secrets
import stat
import tempfile
from datetime import UTC, datetime
from pathlib import Path

ENDPOINT = "https://s3.nl-ams.scw.cloud"
BUCKET = "railiance-platform-pg-backup"
PROFILES = {
    "coulomb-social": "platform-pg/application-archives/coulomb-social/",
    "email-connect": "platform-pg/application-archives/email-connect/",
}
PART_SIZE = 64 * 1024**2
MAX_BYTES = 20 * 1024**3


def private_directory(path):
    info = path.lstat()
    if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.geteuid() or info.st_mode & 0o077:
        raise ValueError("private_owner_directory_required")


def regular_file(path):
    info = path.lstat()
    if not stat.S_ISREG(info.st_mode) or info.st_uid != os.geteuid() or info.st_mode & 0o077:
        raise ValueError("private_owner_file_required")


def digest(path):
    with path.open("rb") as source:
        return hashlib.file_digest(source, "sha256").hexdigest()


def checkpoint(path, receipt):
    fd, name = tempfile.mkstemp(prefix="archive-receipt-", dir=path.parent)
    with os.fdopen(fd, "w") as output:
        json.dump(receipt, output, indent=2)
        output.flush()
        os.fsync(output.fileno())
    Path(name).replace(path)


def key_for(application, captured_at):
    stamp = datetime.fromisoformat(captured_at.replace("Z", "+00:00"))
    if stamp.tzinfo is None:
        raise ValueError("capture_timezone_required")
    stamp = stamp.astimezone(UTC)
    return (
        PROFILES[application]
        + stamp.strftime("%Y/%m/%d/%Y%m%dT%H%M%SZ-")
        + secrets.token_hex(12)
        + ".tar.age"
    )


def validate_key(application, key):
    suffix = key.removeprefix(PROFILES[application])
    if not key.startswith(PROFILES[application]) or not re.fullmatch(
        r"\d{4}/\d{2}/\d{2}/\d{8}T\d{6}Z-[a-f0-9]{24}\.tar\.age", suffix
    ):
        raise ValueError("archive_key_outside_contract")


def missing(error):
    response = getattr(error, "response", {})
    return response.get("ResponseMetadata", {}).get("HTTPStatusCode") == 404


def require_version(value):
    if (
        not isinstance(value, str)
        or not value
        or value == "null"
        or any(c in value for c in "\r\n")
    ):
        raise ValueError("versioned_archive_required")
    return value


def existing(client, key, expected, size):
    try:
        head = client.head_object(Bucket=BUCKET, Key=key)
    except Exception as error:
        if missing(error):
            return None
        raise
    if head.get("ContentLength") != size or head.get("Metadata", {}).get("sha256") != expected:
        raise ValueError("existing_archive_binding_mismatch")
    return require_version(head.get("VersionId"))


def transfer(client, application, source, output, receipt, save=lambda: None):
    regular_file(source)
    private_directory(output.parent)
    size = source.stat().st_size
    if not 0 < size <= MAX_BYTES or output.exists() or output.is_symlink():
        raise ValueError("invalid_archive_source_or_output")
    with source.open("rb") as stream:
        if stream.read(22) != b"age-encryption.org/v1\n":
            raise ValueError("encrypted_archive_required")
    expected = digest(source)
    key = receipt["object_key"]
    validate_key(application, key)
    if receipt.get("ciphertext_sha256", expected) != expected:
        raise ValueError("archive_receipt_binding_mismatch")
    receipt.update(
        application=application,
        destination="s3://" + BUCKET + "/" + key,
        ciphertext_sha256=expected,
        ciphertext_bytes=size,
        remote_retrieval_verified=False,
        stage="existing_object_check",
    )
    save()
    version = existing(client, key, expected, size)
    if receipt.get("version_id") and receipt["version_id"] != version:
        raise ValueError("existing_archive_version_changed")
    upload_id = None
    if version is None:
        try:
            receipt["stage"] = "multipart_upload"
            save()
            upload_id = client.create_multipart_upload(
                Bucket=BUCKET,
                Key=key,
                ContentType="application/octet-stream",
                Metadata={"sha256": expected, "application": application},
            )["UploadId"]
            parts = []
            with source.open("rb") as stream:
                while block := stream.read(PART_SIZE):
                    number = len(parts) + 1
                    result = client.upload_part(
                        Bucket=BUCKET, Key=key, UploadId=upload_id, PartNumber=number, Body=block
                    )
                    parts.append({"PartNumber": number, "ETag": result["ETag"]})
            completed = client.complete_multipart_upload(
                Bucket=BUCKET,
                Key=key,
                UploadId=upload_id,
                MultipartUpload={"Parts": parts},
                IfNoneMatch="*",
            )
            upload_id = None
            version = require_version(completed.get("VersionId"))
            receipt["multipart_completed"] = True
        finally:
            if upload_id:
                try:
                    client.abort_multipart_upload(Bucket=BUCKET, Key=key, UploadId=upload_id)
                    receipt["multipart_aborted"] = True
                except Exception:
                    receipt["multipart_aborted"] = False
                save()
    else:
        receipt["reused_verified_binding"] = True
    receipt.update(stage="version_pinned_readback", version_id=version)
    save()
    download(client, key, version, size, expected, output)
    receipt.update(
        status="offsite_fetched_verified",
        stage="transfer_verified",
        remote_retrieval_verified=True,
        immutable_create=True,
    )
    save()
    return receipt


def download(client, key, version, size, expected, output):
    private_directory(output.parent)
    if output.exists() or output.is_symlink() or not 0 < size <= MAX_BYTES:
        raise ValueError("fresh_bounded_readback_required")
    response = client.get_object(Bucket=BUCKET, Key=key, VersionId=version)
    if response["ContentLength"] != size or response.get("VersionId") != version:
        response["Body"].close()
        raise ValueError("readback_version_or_length_mismatch")
    read_bytes = 0
    with response["Body"] as body, output.open("xb") as target:
        output.chmod(0o600)
        while block := body.read(8 * 1024**2):
            read_bytes += len(block)
            if read_bytes > size:
                raise ValueError("readback_exceeds_expected_size")
            target.write(block)
    if read_bytes != size or digest(output) != expected:
        raise ValueError("readback_digest_mismatch")


def retrieve(client, application, output, receipt):
    """Node-loss recovery needs only the independently retained public receipt."""
    key = receipt.get("object_key", "")
    validate_key(application, key)
    version = require_version(receipt.get("version_id"))
    expected = receipt.get("ciphertext_sha256", "")
    size = receipt.get("ciphertext_bytes", 0)
    if (
        receipt.get("schema") != "platform.application-primary-archive.v1"
        or receipt.get("application") != application
        or receipt.get("status") != "offsite_fetched_verified"
        or receipt.get("remote_retrieval_verified") is not True
        or receipt.get("destination") != "s3://" + BUCKET + "/" + key
        or not re.fullmatch(r"[a-f0-9]{64}", expected)
        or not isinstance(size, int)
        or isinstance(size, bool)
        or not 0 < size <= MAX_BYTES
    ):
        raise ValueError("verified_recovery_receipt_required")
    download(client, key, version, size, expected, output)
    return {
        "status": "retrieved_verified",
        "application": application,
        "object_key": key,
        "version_id": version,
        "ciphertext_sha256": expected,
        "ciphertext_bytes": size,
        "application_restore_verified": False,
        "provider_mutations": False,
    }


def runtime_client():
    import boto3
    from botocore.config import Config

    key = os.environ.get("AWS_ACCESS_KEY_ID", "").strip()
    secret = os.environ.get("AWS_SECRET_ACCESS_KEY", "").strip()
    if not key or not secret:
        raise ValueError("owner_delivered_primary_custody_required")
    return boto3.client(
        "s3",
        endpoint_url=ENDPOINT,
        region_name="nl-ams",
        aws_access_key_id=key,
        aws_secret_access_key=secret,
        config=Config(
            signature_version="s3v4",
            connect_timeout=15,
            read_timeout=180,
            retries={"max_attempts": 3},
            s3={"addressing_style": "path"},
            request_checksum_calculation="when_required",
            response_checksum_validation="when_required",
        ),
    )


def check_runtime():
    from botocore.session import get_session

    model = get_session().get_service_model("s3").operation_model("CompleteMultipartUpload")
    if "IfNoneMatch" not in model.input_shape.members:
        raise ValueError("conditional_multipart_sdk_required")
    if not all(
        os.environ.get(name, "").strip() for name in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"]
    ):
        raise ValueError("owner_delivered_primary_custody_required")
    return {
        "status": "configured",
        "credential_values_printed": False,
        "provider_contacted": False,
        "conditional_multipart_sdk": True,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["check", "transfer", "retrieve"])
    parser.add_argument("--application", choices=PROFILES, required=True)
    parser.add_argument("--source", type=Path)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--receipt", type=Path)
    parser.add_argument("--captured-at")
    parser.add_argument("--execute", action="store_true")
    args = parser.parse_args()
    if args.action == "check":
        print(json.dumps(check_runtime()))
        return 0
    if not args.execute:
        print(
            json.dumps(
                {
                    "status": "reviewed",
                    "provider_contacted": False,
                    "bucket": BUCKET,
                    "prefix": PROFILES[args.application],
                    "immutable_create": True,
                }
            )
        )
        return 0
    if args.action == "retrieve":
        if not args.output or not args.receipt:
            raise ValueError("retrieval_arguments_required")
        regular_file(args.receipt)
        receipt = json.loads(args.receipt.read_text())
        print(json.dumps(retrieve(runtime_client(), args.application, args.output, receipt)))
        return 0
    if not all([args.source, args.output, args.receipt, args.captured_at]):
        raise ValueError("transfer_arguments_required")
    regular_file(args.source)
    private_directory(args.output.parent)
    private_directory(args.receipt.parent)
    if args.output.is_symlink() or args.receipt.is_symlink():
        raise ValueError("linked_transfer_path_refused")
    if args.receipt.exists():
        regular_file(args.receipt)
        receipt = json.loads(args.receipt.read_text())
        if (
            receipt.get("schema") != "platform.application-primary-archive.v1"
            or receipt.get("application") != args.application
            or receipt.get("captured_at") != args.captured_at
        ):
            raise ValueError("transfer_receipt_scope_mismatch")
    else:
        receipt = {
            "schema": "platform.application-primary-archive.v1",
            "status": "pending",
            "application": args.application,
            "captured_at": args.captured_at,
            "object_key": key_for(args.application, args.captured_at),
        }
    validate_key(args.application, receipt["object_key"])
    if args.output.exists():
        regular_file(args.output)
        if (
            receipt.get("status") == "offsite_fetched_verified"
            and receipt.get("remote_retrieval_verified")
            and receipt.get("immutable_create")
            and require_version(receipt.get("version_id"))
            and args.source.stat().st_size == receipt.get("ciphertext_bytes")
            and digest(args.output) == receipt.get("ciphertext_sha256")
            and digest(args.source) == receipt.get("ciphertext_sha256")
        ):
            print(json.dumps(receipt))
            return 0
        raise ValueError("unverified_existing_readback_requires_review")
    try:
        checkpoint(args.receipt, receipt)
        transfer(
            runtime_client(),
            args.application,
            args.source,
            args.output,
            receipt,
            lambda: checkpoint(args.receipt, receipt),
        )
    except Exception as error:
        receipt.update(status="failed", failure_type=type(error).__name__)
        checkpoint(args.receipt, receipt)
        raise
    print(json.dumps(receipt))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as error:
        print(json.dumps({"status": "failed", "failure_type": type(error).__name__}))
        raise SystemExit(1)