"""Authenticate the maintainer's sealed evaluator before publishing private files."""
from __future__ import annotations

import base64
import hashlib
import hmac
import io
import json
import os
from pathlib import Path
import re
import shutil
import stat
import tempfile
import zipfile

from cryptography.hazmat.primitives.ciphers.aead import AESGCM

MAGIC = b"AUTORSE-EVAL-GCM1\n"
MAX_BLOB_BYTES = 64 * 1024**2
MAX_EXPANDED_BYTES = 64 * 1024**2
MAX_MANIFEST_BYTES = 64 * 1024
PRIVATE_FILES = frozenset("""hstu/beauty/expected_user_ids.npy hstu/beauty/catalog.npy hstu/beauty/test_targets.npy
tiger/beauty/expected_user_ids.npy tiger/beauty/catalog.npy tiger/beauty/test_targets.npy
intsr/Amazon/expected_submission.npz intsr/Amazon/test_data_labeled.csv
wukong/movielenslatest/test_data_labeled.csv rankmixer/criteo5m/expected_sample_ids.npy
rankmixer/criteo5m/test_labels.npy onetrans/kuaisar_v2/labels.npz""".split())
FILES = PRIVATE_FILES | {"private-manifest.json"}
ERROR = "Encrypted evaluator verification or initialization failed."


def make_aad(expected_manifest_sha, public_bundle_sha, space_id) -> bytes:
    return json.dumps({"protocol": "autorse-evaluator-aesgcm-v1",
                       "manifest_sha256": expected_manifest_sha,
                       "public_bundle_sha256": public_bundle_sha,
                       "space_id": space_id}, sort_keys=True, separators=(",", ":")).encode("utf-8")


def _require(condition):
    if not condition:
        raise ValueError(ERROR)


def _hash(data):
    return hashlib.sha256(data).hexdigest()


def _sha(value):
    return isinstance(value, str) and re.fullmatch(r"[0-9a-f]{64}", value) is not None


def _unique_object(pairs):
    result = {}
    for key, value in pairs:
        _require(key not in result)
        result[key] = value
    return result


def _unpack(plaintext, expected_manifest_sha, public_bundle_sha):
    _require(len(plaintext) <= MAX_BLOB_BYTES)
    with zipfile.ZipFile(io.BytesIO(plaintext)) as archive:
        infos = archive.infolist()
        _require(len(infos) == len(FILES) and {i.filename for i in infos} == FILES)
        _require(sum(i.file_size for i in infos) <= MAX_EXPANDED_BYTES)
        by_name = {}
        for info in infos:
            name = info.filename
            _require(name not in by_name and not name.startswith("/") and "\\" not in name
                     and all(part not in ("", ".", "..") for part in name.split("/")))
            mode = info.external_attr >> 16
            _require(not info.is_dir() and stat.S_IFMT(mode) in (0, stat.S_IFREG)
                     and not info.flag_bits & 1 and 0 <= info.file_size <= MAX_EXPANDED_BYTES
                     and info.compress_type in (zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED))
            by_name[name] = info

        def read(name, size):
            info = by_name[name]
            _require(info.file_size == size)
            with archive.open(info) as stream:
                value = stream.read(size + 1)
            _require(len(value) == size)
            return value

        manifest_info = by_name["private-manifest.json"]
        _require(manifest_info.file_size <= MAX_MANIFEST_BYTES)
        raw_manifest = read("private-manifest.json", manifest_info.file_size)
        _require(hmac.compare_digest(_hash(raw_manifest), expected_manifest_sha))
        manifest = json.loads(raw_manifest, object_pairs_hook=_unique_object)
        _require(isinstance(manifest, dict) and manifest.get("schema_version") == 1
                 and manifest.get("visibility") == "PRIVATE_EVALUATOR_ONLY"
                 and manifest.get("public_bundle_sha256") == public_bundle_sha)
        rows = manifest.get("files")
        _require(isinstance(rows, list) and len(rows) == len(PRIVATE_FILES))
        result = {"private-manifest.json": raw_manifest}
        for row in rows:
            _require(isinstance(row, dict) and set(row) == {"path", "size", "sha256"})
            name, size, sha = row["path"], row["size"], row["sha256"]
            _require(isinstance(name, str) and name in PRIVATE_FILES and name not in result
                     and type(size) is int and 0 <= size <= MAX_EXPANDED_BYTES and _sha(sha))
            value = read(name, size)
            _require(hmac.compare_digest(_hash(value), sha))
            result[name] = value
        _require(set(result) == FILES)
        return result


def _safe_chain(directory):
    for path in (directory, *directory.parents):
        _require(not path.is_symlink() and (not path.exists() or path.is_dir()))


def _existing_matches(root, payload):
    entries = list(root.rglob("*"))
    _require(not any(p.is_symlink() for p in entries))
    if not entries:
        return False
    _require({p.relative_to(root).as_posix() for p in entries if p.is_file()} == FILES)
    for path in (root, *entries):
        metadata = path.stat()
        if path.is_dir():
            _require(stat.S_IMODE(metadata.st_mode) == 0o700)
        else:
            name = path.relative_to(root).as_posix()
            _require(stat.S_ISREG(metadata.st_mode) and metadata.st_nlink == 1
                     and stat.S_IMODE(metadata.st_mode) == 0o600
                     and metadata.st_size == len(payload[name]))
            with path.open("rb") as stream:
                digest = hashlib.file_digest(stream, "sha256").hexdigest()
            _require(hmac.compare_digest(digest, _hash(payload[name])))
    return True


def _write_file(path, data):
    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), 0o600)
    with os.fdopen(descriptor, "wb") as stream:
        os.fchmod(stream.fileno(), 0o600)
        stream.write(data)


def provision_encrypted_data(*, blob_path: Path, private_root: Path, key_b64: str,
                             expected_cipher_sha: str, expected_manifest_sha: str,
                             public_bundle_sha: str,
                             space_id: str = "Ray0311/AutoRSE-Evaluator") -> None:
    stage = None
    try:
        _require(_sha(expected_cipher_sha) and _sha(expected_manifest_sha) and _sha(public_bundle_sha)
                 and isinstance(space_id, str) and 0 < len(space_id) <= 256 and isinstance(key_b64, str))
        key = base64.b64decode(key_b64, validate=True)
        _require(len(key) == 32)
        blob_path, root = Path(blob_path), Path(private_root).absolute()
        _require(not blob_path.is_symlink() and blob_path.is_file()
                 and len(MAGIC) + 12 + 16 <= blob_path.stat().st_size <= MAX_BLOB_BYTES)
        with blob_path.open("rb") as stream:
            blob = stream.read(MAX_BLOB_BYTES + 1)
        _require(len(blob) <= MAX_BLOB_BYTES and blob.startswith(MAGIC)
                 and hmac.compare_digest(_hash(blob), expected_cipher_sha))
        offset = len(MAGIC)
        plaintext = AESGCM(key).decrypt(blob[offset:offset + 12], blob[offset + 12:],
                                         make_aad(expected_manifest_sha, public_bundle_sha, space_id))
        payload = _unpack(plaintext, expected_manifest_sha, public_bundle_sha)
        _safe_chain(root)
        if root.exists() and _existing_matches(root, payload):
            return
        missing = []
        parent = root.parent
        while not parent.exists():
            missing.append(parent)
            parent = parent.parent
        for parent in reversed(missing):
            parent.mkdir(mode=0o700)
            parent.chmod(0o700)
        _safe_chain(root)
        stage = Path(tempfile.mkdtemp(prefix="." + root.name + "-", dir=root.parent))
        stage.chmod(0o700)
        for name, data in payload.items():
            path = stage / name
            path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
            _write_file(path, data)
        for directory in stage.rglob("*"):
            if directory.is_dir():
                directory.chmod(0o700)
        _safe_chain(root)
        os.replace(stage, root)
        stage = None
    except Exception:
        raise ValueError(ERROR) from None
    finally:
        if stage is not None:
            shutil.rmtree(stage, ignore_errors=True)