#!/usr/bin/env python3
"""Select tenant HDD files, submit NAS copies and download verified archives."""

import argparse
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import ssl
import stat
import sys
import uuid
from urllib import error, parse, request


MAX_JSON_BYTES = 1024 * 1024
IDENTIFIER = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]{0,127}\Z")


class ClientError(Exception):
    """A safe error that never includes authorization headers or response bodies."""


class NoRedirect(request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise ClientError("redirect_refused")


def identifier(value):
    if not isinstance(value, str) or not IDENTIFIER.fullmatch(value):
        raise ClientError("invalid_identifier")
    return value


def base_url(value):
    if not isinstance(value, str) or any(ord(c) < 33 or ord(c) > 126 for c in value):
        raise ClientError("https_api_base_required")
    try:
        parsed = parse.urlsplit(value)
        port = parsed.port
    except ValueError as exc:
        raise ClientError("https_api_base_required") from exc
    if (parsed.scheme != "https" or not parsed.hostname or parsed.username is not None
            or parsed.password is not None or parsed.query or parsed.fragment
            or port == 0 or parsed.hostname.endswith(".invalid")
            or parsed.path.rstrip("/") != "/v1/storage"):
        raise ClientError("https_api_base_required")
    return value.rstrip("/")


def selection(path):
    with Path(path).open("rb") as stream:
        raw = stream.read(MAX_JSON_BYTES + 1)
    if len(raw) > MAX_JSON_BYTES:
        raise ClientError("selection_too_large")
    try:
        rows = json.loads(raw)
    except (ValueError, UnicodeError) as exc:
        raise ClientError("selection_json_required") from exc
    if not isinstance(rows, list) or not rows or len(rows) > 1000:
        raise ClientError("selection_list_required")
    seen = set()
    for row in rows:
        if not isinstance(row, dict) or set(row) != {"file_id", "version"}:
            raise ClientError("file_id_and_version_required")
        identifier(row["file_id"])
        identifier(row["version"])
        if row["file_id"] in seen:
            raise ClientError("duplicate_file_selection")
        seen.add(row["file_id"])
    return rows


def request_uuid(value):
    try:
        normalized = str(uuid.UUID(value))
    except (ValueError, TypeError, AttributeError) as exc:
        raise ClientError("request_uuid_required") from exc
    if normalized != value:
        raise ClientError("canonical_request_uuid_required")
    return normalized


class ArchiveClient:
    def __init__(self, origin, token, *, timeout=30, opener=None):
        self.origin = base_url(origin)
        if (not isinstance(token, str) or not 1 <= len(token) <= 4096
                or any(ord(c) < 33 or ord(c) > 126 for c in token)):
            raise ClientError("archive_api_key_required")
        if type(timeout) is not int or not 1 <= timeout <= 300:
            raise ClientError("invalid_client_timeout")
        self.token = token
        self.timeout = timeout
        self.opener = opener or request.build_opener(
            request.ProxyHandler({}), NoRedirect(), request.HTTPSHandler(context=ssl.create_default_context()))

    def call(self, method, suffix, *, body=None, request_id=None):
        headers = {"Authorization": "Bearer " + self.token, "Accept": "application/json"}
        data = None
        if body is not None:
            headers["Content-Type"] = "application/json"
            headers["Idempotency-Key"] = request_uuid(request_id)
            data = json.dumps(body, separators=(",", ":")).encode("utf-8")
        message = request.Request(self.origin + suffix, data=data, headers=headers, method=method)
        try:
            # One request only. On timeout, callers query or reuse the same intent.
            with self.opener.open(message, timeout=self.timeout) as response:
                if response.status not in (200, 202):
                    raise ClientError("unexpected_http_status")
                if response.headers.get_content_type() != "application/json":
                    raise ClientError("json_response_required")
                raw = response.read(MAX_JSON_BYTES + 1)
        except error.HTTPError as exc:
            raise ClientError("http_" + str(exc.code)) from None
        except (error.URLError, TimeoutError, OSError) as exc:
            raise ClientError("transport_unconfirmed_do_not_resubmit") from None
        if len(raw) > MAX_JSON_BYTES:
            raise ClientError("response_too_large")
        try:
            result = json.loads(raw)
        except (ValueError, UnicodeError) as exc:
            raise ClientError("invalid_json_response") from exc
        if not isinstance(result, dict):
            raise ClientError("object_response_required")
        return result

    def submit(self, destination, files, request_id):
        return self.call("POST", "/archive-jobs", request_id=request_id, body={
            "destination_id": identifier(destination), "mode": "copy", "files": files})

    def download(self, object_id, output, *, resume=False):
        identifier(object_id)
        target = Path(output).absolute()
        if target.exists() or target.is_symlink() or not target.parent.is_dir():
            raise ClientError("new_output_file_required")
        partial = target.with_name(target.name + ".part")
        metadata = target.with_name(target.name + ".part.json")
        flags = os.O_RDWR | os.O_NOFOLLOW | (0 if resume else os.O_CREAT | os.O_EXCL)
        with os.fdopen(os.open(partial, flags, 0o600), "r+b") as stream:
            fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB)
            info = os.fstat(stream.fileno())
            if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or info.st_uid != os.geteuid() or info.st_mode & 0o077:
                raise ClientError("private_partial_file_required")
            saved = None
            if resume:
                with os.fdopen(os.open(metadata, os.O_RDONLY | os.O_NOFOLLOW), "rb") as details:
                    mode = os.fstat(details.fileno())
                    if mode.st_size > 4096 or mode.st_mode & 0o077 or mode.st_nlink != 1 or mode.st_uid != os.geteuid():
                        raise ClientError("private_download_metadata_required")
                    saved = json.load(details)
                if saved.get("api_base") != self.origin or saved.get("object_id") != object_id:
                    raise ClientError("download_binding_changed")
            offset = info.st_size
            if saved is not None:
                if type(saved.get("bytes")) is not int or not 0 <= offset <= saved["bytes"]:
                    raise ClientError("download_length_invalid")
                # Revalidate the origin even when all bytes arrived before a disconnect.
                if offset and offset == saved["bytes"]:
                    offset -= 1
            headers = {"Authorization": "Bearer " + self.token}
            if offset:
                headers["Range"] = f"bytes={offset}-"
            message = request.Request(self.origin + "/objects/" + object_id, headers=headers)
            try:
                with self.opener.open(message, timeout=self.timeout) as response:
                    if response.status != (206 if offset else 200):
                        raise ClientError("download_range_unconfirmed")
                    digest = response.headers.get("X-Content-SHA256", "")
                    if not re.fullmatch(r"[a-f0-9]{64}", digest):
                        raise ClientError("download_digest_missing")
                    try:
                        remaining = int(response.headers["Content-Length"])
                    except (TypeError, ValueError):
                        raise ClientError("download_length_missing") from None
                    if remaining < 0:
                        raise ClientError("download_length_invalid")
                    total = offset + remaining
                    if offset and response.headers.get("Content-Range") != f"bytes {offset}-{total-1}/{total}":
                        raise ClientError("download_range_unconfirmed")
                    current = {"api_base": self.origin, "object_id": object_id, "bytes": total, "sha256": digest}
                    if saved is not None and saved != current:
                        raise ClientError("download_version_changed")
                    if saved is None:
                        with os.fdopen(os.open(metadata, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600), "w") as details:
                            json.dump(current, details)
                            details.flush(); os.fsync(details.fileno())
                    stream.seek(offset)
                    while remaining:
                        chunk = response.read(min(1024 * 1024, remaining))
                        if not chunk:
                            raise ClientError("download_incomplete_resume_same_object")
                        stream.write(chunk)
                        remaining -= len(chunk)
                    stream.flush(); os.fsync(stream.fileno())
            except error.HTTPError as exc:
                raise ClientError("http_" + str(exc.code)) from None
            except (error.URLError, TimeoutError, OSError):
                raise ClientError("download_unconfirmed_resume_same_object") from None
            stream.seek(0)
            actual = hashlib.sha256()
            while chunk := stream.read(1024 * 1024):
                actual.update(chunk)
            if actual.hexdigest() != digest:
                raise ClientError("download_readback_mismatch")
            # Publish without overwriting an existing output, retaining failed partials.
            os.link(partial, target, follow_symlinks=False)
            parent = os.open(target.parent, os.O_RDONLY | os.O_DIRECTORY)
            try:
                os.fsync(parent)
            finally:
                os.close(parent)
        partial.unlink()
        metadata.unlink()
        return {"object_id": object_id, "bytes": total, "sha256": digest, "output": str(target)}


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--api-base", default=os.environ.get("FILE_ARCHIVE_API_BASE"))
    parser.add_argument("--timeout", type=int, default=30)
    commands = parser.add_subparsers(dest="command", required=True)
    files = commands.add_parser("files")
    files.add_argument("--cursor")
    files.add_argument("--directory")
    commands.add_parser("destinations")
    commands.add_parser("quota")
    commands.add_parser("jobs")
    download = commands.add_parser("download")
    download.add_argument("object_id")
    download.add_argument("--output", required=True)
    download.add_argument("--resume", action="store_true")
    submit = commands.add_parser("submit")
    submit.add_argument("--destination", required=True)
    submit.add_argument("--selection", required=True)
    submit.add_argument("--request-id", required=True)
    status = commands.add_parser("status")
    status.add_argument("job_id")
    lookup = commands.add_parser("lookup")
    lookup.add_argument("request_id")
    args = parser.parse_args(argv)
    client = ArchiveClient(args.api_base, os.environ.get("FILE_ARCHIVE_API_KEY"), timeout=args.timeout)
    if args.command == "download":
        result = client.download(args.object_id, args.output, resume=args.resume)
    elif args.command == "jobs":
        result = client.call("GET", "/archive-jobs")
    elif args.command == "submit":
        result = client.submit(args.destination, selection(args.selection), args.request_id)
    elif args.command == "status":
        result = client.call("GET", "/archive-jobs/" + identifier(args.job_id))
    elif args.command == "lookup":
        result = client.call("GET", "/archive-requests/" + request_uuid(args.request_id))
    elif args.command == "files":
        fields = {name: identifier(value) for name, value in (("cursor", args.cursor), ("directory_id", args.directory)) if value}
        query = "?" + parse.urlencode(fields) if fields else ""
        result = client.call("GET", "/files" + query)
    else:
        result = client.call("GET", "/" + args.command)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    try:
        main()
    except (ClientError, OSError) as exc:
        code = str(exc) if isinstance(exc, ClientError) else "local_file_unavailable"
        print(json.dumps({"error": code}), file=sys.stderr)
        raise SystemExit(2)