"""Immutable staged-to-committed object workflow shared by local and cloud backends."""

from __future__ import annotations

import hashlib
import os
import stat
import tempfile
from pathlib import Path
from typing import Dict, Mapping, Optional, Tuple

from arr_ingestion.contracts import ROLE_CONTRACTS, SHA256_RE, IngestionError
from arr_storage.contracts import (
    OBJECT_METADATA_SCHEMA,
    BackendError,
    BackendObject,
    ObjectAddress,
    ObjectBackend,
    ObjectKeyPolicy,
    StoredObject,
    role_limit,
    valid_delivery_filename,
)


_REQUIRED_METADATA = frozenset(
    {
        "arr-object-schema",
        "arr-state",
        "arr-role",
        "arr-job-id",
        "arr-attempt-no",
        "arr-sha256",
        "arr-byte-size",
        "arr-mime-type",
    }
)


class ManagedObjectStore:
    """ARR-owned immutable object store and ingestion ArtifactStore implementation."""

    def __init__(self, backend: ObjectBackend, policy: Optional[ObjectKeyPolicy] = None) -> None:
        self._backend = backend
        self._policy = policy or ObjectKeyPolicy()

    @property
    def key_policy(self) -> ObjectKeyPolicy:
        return self._policy

    def stage_file(
        self,
        *,
        job_id: str,
        attempt_no: int,
        role: str,
        source: Path,
        original_filename: str,
        expected_sha256: Optional[str] = None,
        expected_byte_size: Optional[int] = None,
    ) -> StoredObject:
        address = ObjectAddress(job_id, attempt_no, "staged", role)
        object_key = self._policy.build(address)
        if not valid_delivery_filename(role, original_filename):
            raise IngestionError(
                "ARTIFACT_REFERENCE_INVALID", "artifact filename is invalid"
            )
        if expected_sha256 is not None and not SHA256_RE.fullmatch(expected_sha256):
            raise IngestionError(
                "ARTIFACT_REFERENCE_INVALID", "artifact hash is invalid"
            )
        if (
            expected_byte_size is not None
            and (
                not isinstance(expected_byte_size, int)
                or isinstance(expected_byte_size, bool)
                or expected_byte_size < 0
            )
        ):
            raise IngestionError(
                "ARTIFACT_REFERENCE_INVALID", "artifact size is invalid"
            )
        _file_kind, _extension, mime_type = ROLE_CONTRACTS[role]
        with tempfile.TemporaryDirectory(prefix="arr-object-stage-") as temporary:
            snapshot = Path(temporary) / "upload.bin"
            sha256, byte_size = self._snapshot(source, snapshot, role_limit(role))
            if expected_sha256 is not None and sha256 != expected_sha256:
                raise IngestionError(
                    "ARTIFACT_HASH_MISMATCH", "artifact identity does not match"
                )
            if expected_byte_size is not None and byte_size != expected_byte_size:
                raise IngestionError(
                    "ARTIFACT_HASH_MISMATCH", "artifact identity does not match"
                )
            metadata = self._metadata(address, sha256, byte_size, mime_type)
            try:
                self._backend.put_file(
                    object_key,
                    str(snapshot),
                    mime_type,
                    metadata,
                    if_absent=True,
                )
            except BackendError as error:
                if error.kind != "conflict":
                    self._raise_backend(error)
                existing = self.inspect(object_key, original_filename)
                if (existing.sha256, existing.byte_size, existing.mime_type) != (
                    sha256,
                    byte_size,
                    mime_type,
                ):
                    raise IngestionError(
                        "ARTIFACT_CONFLICT", "object key already contains different bytes"
                    ) from None
                return existing
        return self.inspect(object_key, original_filename)

    def commit(self, staged: StoredObject) -> StoredObject:
        self._validate_descriptor(staged, required_state="staged")
        committed_address = ObjectAddress(
            staged.job_id,
            staged.attempt_no,
            "committed",
            staged.role,
        )
        committed_key = self._policy.build(committed_address)
        try:
            verified_staged = self.inspect(staged.object_key, staged.original_filename)
        except IngestionError as error:
            if error.code != "ARTIFACT_NOT_FOUND":
                raise
            committed = self.inspect(committed_key, staged.original_filename)
            self._require_same_identity(staged, committed)
            return committed
        self._require_same_identity(staged, verified_staged)
        metadata = self._metadata(
            committed_address,
            staged.sha256,
            staged.byte_size,
            staged.mime_type,
        )
        try:
            self._backend.copy_object(
                staged.object_key,
                committed_key,
                metadata,
                if_absent=True,
            )
        except BackendError as error:
            if error.kind != "conflict":
                self._raise_backend(error)
            committed = self.inspect(committed_key, staged.original_filename)
            self._require_same_identity(staged, committed)
        else:
            committed = self.inspect(committed_key, staged.original_filename)
            self._require_same_identity(staged, committed)
        try:
            self._backend.delete_object(staged.object_key)
        except BackendError as error:
            self._raise_backend(error)
        return committed

    def upload_committed(
        self,
        *,
        job_id: str,
        attempt_no: int,
        role: str,
        source: Path,
        original_filename: str,
        expected_sha256: Optional[str] = None,
        expected_byte_size: Optional[int] = None,
    ) -> StoredObject:
        staged = self.stage_file(
            job_id=job_id,
            attempt_no=attempt_no,
            role=role,
            source=source,
            original_filename=original_filename,
            expected_sha256=expected_sha256,
            expected_byte_size=expected_byte_size,
        )
        return self.commit(staged)

    def inspect(self, object_key: str, original_filename: str) -> StoredObject:
        address = self._policy.parse(object_key)
        if not valid_delivery_filename(address.role, original_filename):
            raise IngestionError(
                "ARTIFACT_REFERENCE_INVALID", "artifact filename is invalid"
            )
        try:
            head = self._backend.head(object_key)
        except BackendError as error:
            self._raise_backend(error)
        descriptor = self._descriptor_from_head(address, original_filename, head)
        actual_sha, actual_size = self._hash_object(object_key, role_limit(address.role))
        if (actual_sha, actual_size) != (descriptor.sha256, descriptor.byte_size):
            raise IngestionError(
                "ARTIFACT_HASH_MISMATCH", "stored object identity does not match metadata"
            )
        return descriptor

    def inspect_committed(self, object_key: str, original_filename: str) -> StoredObject:
        descriptor = self.inspect(object_key, original_filename)
        if descriptor.state != "committed":
            raise IngestionError(
                "ARTIFACT_NOT_COMMITTED", "only committed objects may be read"
            )
        return descriptor

    def materialize(self, object_key: str, destination: Path, max_bytes: int) -> None:
        address = self._policy.parse(object_key)
        if address.state != "committed":
            raise IngestionError(
                "ARTIFACT_NOT_COMMITTED", "only committed objects may be read"
            )
        if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes < 0:
            raise IngestionError("ARTIFACT_REFERENCE_INVALID", "artifact size limit is invalid")
        try:
            head = self._backend.head(object_key)
        except BackendError as error:
            self._raise_backend(error)
        descriptor = self._descriptor_from_head(address, "placeholder" + ROLE_CONTRACTS[address.role][1], head)
        if descriptor.byte_size > max_bytes:
            raise IngestionError(
                "ARTIFACT_TOO_LARGE", "delivery artifact exceeds its size limit"
            )
        destination.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
        output_descriptor = -1
        destination_created = False
        copied = 0
        digest = hashlib.sha256()
        try:
            output_descriptor = os.open(
                destination,
                os.O_WRONLY | os.O_CREAT | os.O_EXCL,
                0o600,
            )
            destination_created = True
            with self._open_reader(object_key) as reader, os.fdopen(output_descriptor, "wb") as target:
                output_descriptor = -1
                while True:
                    chunk = self._read_chunk(reader)
                    if not chunk:
                        break
                    if not isinstance(chunk, bytes):
                        raise IngestionError(
                            "OBJECT_STORE_UNAVAILABLE", "object storage read did not complete"
                        )
                    copied += len(chunk)
                    if copied > max_bytes:
                        raise IngestionError(
                            "ARTIFACT_TOO_LARGE", "delivery artifact exceeds its size limit"
                        )
                    digest.update(chunk)
                    target.write(chunk)
                target.flush()
                os.fsync(target.fileno())
            if (copied, digest.hexdigest()) != (descriptor.byte_size, descriptor.sha256):
                raise IngestionError(
                    "ARTIFACT_HASH_MISMATCH", "stored object identity does not match metadata"
                )
        except FileExistsError:
            raise IngestionError(
                "ARTIFACT_CONFLICT", "materialization destination already exists"
            ) from None
        except Exception:
            if destination_created:
                destination.unlink(missing_ok=True)
            raise
        finally:
            if output_descriptor >= 0:
                os.close(output_descriptor)

    def _descriptor_from_head(
        self,
        address: ObjectAddress,
        original_filename: str,
        head: BackendObject,
    ) -> StoredObject:
        if head.object_key != self._policy.build(address):
            raise IngestionError("ARTIFACT_REFERENCE_INVALID", "stored object key is invalid")
        metadata = dict(head.metadata)
        if not _REQUIRED_METADATA.issubset(metadata):
            raise IngestionError(
                "ARTIFACT_METADATA_INVALID", "stored object metadata is incomplete"
            )
        _file_kind, _extension, expected_mime = ROLE_CONTRACTS[address.role]
        expected_values: Mapping[str, str] = {
            "arr-object-schema": OBJECT_METADATA_SCHEMA,
            "arr-state": address.state,
            "arr-role": address.role,
            "arr-job-id": address.job_id,
            "arr-attempt-no": str(address.attempt_no),
            "arr-mime-type": expected_mime,
        }
        if any(metadata.get(key) != value for key, value in expected_values.items()):
            raise IngestionError(
                "ARTIFACT_METADATA_INVALID", "stored object metadata conflicts with its key"
            )
        sha256 = metadata.get("arr-sha256", "")
        size_text = metadata.get("arr-byte-size", "")
        if not SHA256_RE.fullmatch(sha256) or not size_text.isdigit():
            raise IngestionError(
                "ARTIFACT_METADATA_INVALID", "stored object identity metadata is invalid"
            )
        byte_size = int(size_text)
        if byte_size != head.byte_size or byte_size > role_limit(address.role):
            raise IngestionError(
                "ARTIFACT_METADATA_INVALID", "stored object size metadata is invalid"
            )
        backend_content_type = metadata.get("arr-backend-content-type")
        if backend_content_type is not None and backend_content_type != expected_mime:
            raise IngestionError(
                "ARTIFACT_METADATA_INVALID", "stored object content type is invalid"
            )
        return StoredObject(
            object_key=head.object_key,
            job_id=address.job_id,
            attempt_no=address.attempt_no,
            state=address.state,
            role=address.role,
            original_filename=original_filename,
            sha256=sha256,
            byte_size=byte_size,
            mime_type=expected_mime,
            etag=head.etag,
            version_id=head.version_id,
        )

    @staticmethod
    def _metadata(
        address: ObjectAddress,
        sha256: str,
        byte_size: int,
        mime_type: str,
    ) -> Dict[str, str]:
        return {
            "arr-object-schema": OBJECT_METADATA_SCHEMA,
            "arr-state": address.state,
            "arr-role": address.role,
            "arr-job-id": address.job_id,
            "arr-attempt-no": str(address.attempt_no),
            "arr-sha256": sha256,
            "arr-byte-size": str(byte_size),
            "arr-mime-type": mime_type,
        }

    @staticmethod
    def _snapshot(source: Path, destination: Path, max_bytes: int) -> Tuple[str, int]:
        source_descriptor = -1
        destination_descriptor = -1
        copied = 0
        digest = hashlib.sha256()
        try:
            if source.is_symlink():
                raise IngestionError(
                    "ARTIFACT_REFERENCE_INVALID", "upload source must be a regular file"
                )
            source_descriptor = os.open(
                source,
                os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0),
            )
            if not stat.S_ISREG(os.fstat(source_descriptor).st_mode):
                raise IngestionError(
                    "ARTIFACT_REFERENCE_INVALID", "upload source must be a regular file"
                )
            destination_descriptor = os.open(
                destination,
                os.O_WRONLY | os.O_CREAT | os.O_EXCL,
                0o600,
            )
            with os.fdopen(source_descriptor, "rb") as reader, os.fdopen(
                destination_descriptor, "wb"
            ) as target:
                source_descriptor = -1
                destination_descriptor = -1
                while True:
                    chunk = reader.read(1024 * 1024)
                    if not chunk:
                        break
                    copied += len(chunk)
                    if copied > max_bytes:
                        raise IngestionError(
                            "ARTIFACT_TOO_LARGE", "delivery artifact exceeds its size limit"
                        )
                    digest.update(chunk)
                    target.write(chunk)
                target.flush()
                os.fsync(target.fileno())
        except FileNotFoundError:
            destination.unlink(missing_ok=True)
            raise IngestionError(
                "ARTIFACT_NOT_FOUND", "upload source is unavailable"
            ) from None
        except OSError:
            destination.unlink(missing_ok=True)
            raise IngestionError(
                "OBJECT_STORE_UNAVAILABLE", "object storage upload did not complete"
            ) from None
        except Exception:
            destination.unlink(missing_ok=True)
            raise
        finally:
            if source_descriptor >= 0:
                os.close(source_descriptor)
            if destination_descriptor >= 0:
                os.close(destination_descriptor)
        return digest.hexdigest(), copied

    def _hash_object(self, object_key: str, max_bytes: int) -> Tuple[str, int]:
        copied = 0
        digest = hashlib.sha256()
        with self._open_reader(object_key) as reader:
            while True:
                chunk = self._read_chunk(reader)
                if not chunk:
                    break
                if not isinstance(chunk, bytes):
                    raise IngestionError(
                        "OBJECT_STORE_UNAVAILABLE", "object storage read did not complete"
                    )
                copied += len(chunk)
                if copied > max_bytes:
                    raise IngestionError(
                        "ARTIFACT_TOO_LARGE", "delivery artifact exceeds its size limit"
                    )
                digest.update(chunk)
        return digest.hexdigest(), copied

    def _open_reader(self, object_key: str):
        try:
            return self._backend.open_reader(object_key)
        except BackendError as error:
            self._raise_backend(error)

    @staticmethod
    def _read_chunk(reader) -> bytes:
        try:
            return reader.read(1024 * 1024)
        except BackendError as error:
            ManagedObjectStore._raise_backend(error)
        except OSError:
            raise IngestionError(
                "OBJECT_STORE_UNAVAILABLE", "object storage read did not complete"
            ) from None

    def _validate_descriptor(self, descriptor: StoredObject, required_state: str) -> None:
        address = self._policy.parse(descriptor.object_key)
        if (
            address.job_id != descriptor.job_id
            or address.attempt_no != descriptor.attempt_no
            or address.state != descriptor.state
            or address.role != descriptor.role
            or descriptor.state != required_state
            or not valid_delivery_filename(descriptor.role, descriptor.original_filename)
            or not SHA256_RE.fullmatch(descriptor.sha256)
            or descriptor.byte_size < 0
            or descriptor.mime_type != ROLE_CONTRACTS[descriptor.role][2]
        ):
            raise IngestionError(
                "ARTIFACT_REFERENCE_INVALID", "stored object descriptor is invalid"
            )

    @staticmethod
    def _require_same_identity(left: StoredObject, right: StoredObject) -> None:
        if (
            left.job_id,
            left.attempt_no,
            left.role,
            left.sha256,
            left.byte_size,
            left.mime_type,
        ) != (
            right.job_id,
            right.attempt_no,
            right.role,
            right.sha256,
            right.byte_size,
            right.mime_type,
        ):
            raise IngestionError(
                "ARTIFACT_CONFLICT", "stored object identity conflicts"
            )

    @staticmethod
    def _raise_backend(error: BackendError) -> None:
        if error.kind == "not_found":
            raise IngestionError("ARTIFACT_NOT_FOUND", "stored object is unavailable") from None
        if error.kind == "conflict":
            raise IngestionError("ARTIFACT_CONFLICT", "object key already exists") from None
        if error.kind == "invalid_request":
            raise IngestionError("ARTIFACT_REFERENCE_INVALID", "object request is invalid") from None
        raise IngestionError(
            "OBJECT_STORE_UNAVAILABLE", "object storage operation did not complete"
        ) from None