# SPDX-FileCopyrightText: 2026 Blender Authors
# SPDX-License-Identifier: GPL-3.0-or-later

"""BAT v2 packing to a Shaman server."""

from __future__ import annotations

__all__ = ("pack_start",)

import dataclasses
import email.header
import logging
import random
from functools import partial
from pathlib import Path, PurePosixPath
from typing import TYPE_CHECKING, Any, TypeAlias

import bpy

if TYPE_CHECKING:
    # _BATPacker: TypeAlias = pack.BATPacker
    from ..manager import ApiClient as _ApiClient
    from ..manager.apis import ShamanApi as _ShamanApi
    from ..manager.models import (
        ShamanCheckoutResult as _ShamanCheckoutResult,
    )
    from ..manager.models import (
        ShamanFileSpec as _ShamanFileSpec,
    )
    from ..manager.models import (
        ShamanRequirementsRequest as _ShamanRequirementsRequest,
    )
    from .submodules.file_usage import FileInfo as _FileInfo
    from .submodules.pack import BATPacker as _BATPacker
    from .submodules.pack import QueueingExecutor as _QueueingExecutor
else:
    _ApiClient = object
    _ShamanApi = object
    _ShamanCheckoutResult = object
    _ShamanRequirementsRequest = object
    _ShamanFileSpec = object
    _BATPacker = object
    _FileInfo = object
    _QueueingExecutor = object

log = logging.getLogger(__name__)

MAX_DEFERRED_PATHS = 8
MAX_FAILED_PATHS = 8
HASH_STORAGE_PATH = Path(bpy.app.cachedir) / "flamenco/shaman"
HASH_METHOD = "sha256"
SHAMAN_JOBS_VARIABLE = "{jobs}"

HashableShamanFileSpec = tuple[str, int, str]
"""Tuple of the 'sha', 'size', and 'path' fields of a ShamanFileSpec."""

# Alias some types from blender_asset_tracer so that we can use type annotations
# without having to import from BATv2.
BATPackReporter: TypeAlias = Any
BATPacker: TypeAlias = Any


def pack_start(
    project_root: Path,
    reporter: BATPackReporter,
    *,
    use_relative_only: bool,
    api_client: _ApiClient,
    checkout_path: PurePosixPath,
    ignore_globs: set[str] = set(),
) -> BATPacker:
    """Investigate what's needed to create a BAT pack."""
    from ..manager.apis import ShamanApi
    from .submodules import file_usage, pack

    shaman_api = ShamanApi(api_client)
    executor = pack.QueueingExecutor()
    shaman_transferer = ShamanPacker(shaman_api, checkout_path, executor, reporter)

    batpacker = pack.BATPacker(
        project_root,
        file_usage.Options(
            use_relative_only=use_relative_only,
            ignore_globs=ignore_globs,
        ),
        reporter,
        file_transfer=shaman_transferer,
    )

    return batpacker


@dataclasses.dataclass
class ShamanUploadProgress:
    # When another client is already uploading a file that we also want to
    # upload, we defer the file. That way, when we do get around to uploading
    # it, the other person may already have finished their upload, saving us
    # time.
    #
    # Mapping from 'path in pack' to the BAT FileInfo and Shaman FileSpec.
    deferred: dict[PurePosixPath, tuple[_FileInfo, _ShamanFileSpec]] = (
        dataclasses.field(default_factory=dict)
    )

    # When a file doesn't want to get uploaded, it's stored here to retry. If
    # too many files fail, or the retry counter reaches its max, it'll get
    # reported as an actual error.
    # The string value is the error message.
    failures: dict[PurePosixPath, tuple[_FileInfo, _ShamanFileSpec, str]] = (
        dataclasses.field(default_factory=dict)
    )

    retry_counter: int = 0
    max_retries: int = 50

    def is_deferred(self, relpath_in_pack: PurePosixPath) -> bool:
        return relpath_in_pack in self.deferred

    @property
    def num_deferred(self) -> int:
        return len(self.deferred)

    @property
    def num_failed(self) -> int:
        return len(self.failures)

    def defer(
        self,
        bat_file_info: _FileInfo,
        shaman_file_spec: _ShamanFileSpec,
    ) -> None:
        # A file should only be deferred once. Once an upload has been deferred,
        # the next attempt shouldn't be deferred again.
        assert bat_file_info.relpath_in_pack not in self.deferred

        self.deferred[bat_file_info.relpath_in_pack] = (bat_file_info, shaman_file_spec)

    def failed(
        self,
        bat_file_info: _FileInfo,
        shaman_file_spec: _ShamanFileSpec,
        errormsg: str,
    ) -> None:
        # A file should only be added to the 'failures' dict once. When its
        # upload is retried, it should be removed from the 'failures' dict first.
        assert bat_file_info.relpath_in_pack not in self.failures

        self.failures[bat_file_info.relpath_in_pack] = (
            bat_file_info,
            shaman_file_spec,
            errormsg,
        )


@dataclasses.dataclass
class ShamanPacker:
    shaman_api: _ShamanApi
    checkout_path: PurePosixPath
    executor: _QueueingExecutor
    reporter: BATPackReporter

    # Shaman may decide to create the checkout at another path than requested.
    # This will be set to the actually-used path on the farm, relative to the
    # Shaman's "jobs" directory.
    #
    # NOTE: It is the checkout path of the job, NOT the path to the blend file.
    _checkout_path_final: PurePosixPath | None = None

    _source_file_relpath_in_pack: PurePosixPath | None = None

    _num_files_to_transfer_total: int = -1
    _num_files_to_transfer_done: int = 0

    @property
    def is_succes(self) -> bool:
        """Return whether the Shaman operation was completed succesfully."""
        return bool(self._checkout_path_final)

    def start(self, batpacker: _BATPacker) -> None:
        files_to_copy = batpacker.all_files_to_copy()

        # Initial value is the total number of files to copy. Once the Shaman
        # server has told us how many files to submit, this will be adjusted.
        self._num_files_to_transfer_total = len(files_to_copy)
        self._num_files_to_transfer_done = 0

        # Remember where the main blend file sits in the BAT pack.
        source_file_info = batpacker.deps_repo.source_file_info()
        assert source_file_info.relpath_in_pack is not None
        self._source_file_relpath_in_pack = PurePosixPath(
            source_file_info.relpath_in_pack.as_posix()
        )

        self.executor.queue(partial(self._step_queue_hashing, files_to_copy))

    def step(self) -> bool:
        """Perform a single step in the Shaman file transfer.

        Returns whether there are more steps to do (True) or the process is done (False).
        """
        if self.executor.is_done:
            return False
        self.executor.run_step()
        return not self.executor.is_done

    def blendfile_location_in_pack(self) -> PurePosixPath:
        assert self._checkout_path_final is not None
        assert self._source_file_relpath_in_pack is not None
        return (
            PurePosixPath(SHAMAN_JOBS_VARIABLE)
            / self._checkout_path_final
            / self._source_file_relpath_in_pack
        )

    def num_files_to_transfer(self) -> tuple[int, int]:
        """Return the number of files that need to be transferred.

        This is a tuple [total, done] with the total number of files to
        transfer, and the number of transferred files so far.

        The number may change during the packing process, as it takes time
        for the Shaman protocol to get this information. Or some paths may
        turn out to be multiple paths (UDIMs for example).
        """
        return self._num_files_to_transfer_total, self._num_files_to_transfer_done

    def _step_queue_hashing(self, files_to_copy: dict[Path, _FileInfo]) -> None:
        from ..manager.models import ShamanRequirementsRequest

        # Shaman Spec that's shared between the queued function calls. They
        # can all just append to the same list.
        shaman_spec = ShamanRequirementsRequest(files=[])
        assert isinstance(shaman_spec, ShamanRequirementsRequest)

        # Tracks deferred files and failed uploads.
        upload_progress = ShamanUploadProgress()

        # Queue up all the hash computations.
        for file_info in files_to_copy.values():
            self.executor.queue(
                partial(
                    self._step_hash_file,
                    file_info,
                    shaman_spec,
                )
            )

        # After the hashes are gathered in 'filespecs', send the spec to Shaman.
        self.executor.queue(
            partial(
                self._step_queue_uploads_of_files,
                files_to_copy,
                shaman_spec,
                upload_progress,
            )
        )

    def _step_hash_file(
        self,
        file_info: _FileInfo,
        shaman_spec: _ShamanRequirementsRequest,
    ) -> None:
        from _bpy_internal import disk_file_hash_service

        from ..manager.models import ShamanFileSpec

        path_to_pack = file_info.path_to_pack

        if not path_to_pack.exists():
            # If the file is missing, there's little else to do than reporting
            # it as such and continue with the next file.
            if path_to_pack == file_info.source_path:
                log.info("File missing: %s", path_to_pack)
            else:
                log.info(
                    "File missing after rewriting %s to %s",
                    file_info.source_path,
                    path_to_pack,
                )
            self.reporter.on_missing_file(
                file_info.source_path, file_info.relpath_in_pack
            )
            return

        # It might be tempting to use the same Disk File Hash Service as BAT's
        # path rewriting system is using. However, that only hashes the files
        # that need rewriting, and the code below only deals with paths after
        # rewriting (or where rewriting was not necessary). That means that
        # there is no benefit in sharing the same database.
        dfhs = disk_file_hash_service.get_service(HASH_STORAGE_PATH)
        checksum = dfhs.get_hash(path_to_pack, HASH_METHOD)

        filesize = path_to_pack.stat().st_size

        filespec = ShamanFileSpec(
            sha=checksum,
            size=filesize,
            path=str(file_info.relpath_in_pack),
        )
        assert isinstance(filespec, ShamanFileSpec)
        shaman_spec.files.append(filespec)

    def _step_queue_uploads_of_files(
        self,
        files_to_copy: dict[Path, _FileInfo],
        shaman_spec: _ShamanRequirementsRequest,
        upload_progress: ShamanUploadProgress,
    ) -> None:
        """Send the spec file to Shaman, and queue file uploads."""

        # Query Shaman to figure out which files still need uploading.
        to_upload = self._send_spec_to_shaman(shaman_spec)
        if to_upload is None:
            # Errors have been reported already, so just stop.
            return

        log.info(
            "Feeding %d/%d files to the Shaman", len(to_upload), len(shaman_spec.files)
        )
        self._num_files_to_transfer_total = len(to_upload)

        # Create a mapping from the path in the pack (which is used in
        # `filespecs`) to the FileInfo.
        path_in_pack_to_abs: dict[str, _FileInfo] = {
            str(file_info.relpath_in_pack): file_info
            for file_info in files_to_copy.values()
        }

        # Queue the file uploads.
        for index, file_spec in enumerate(to_upload):
            file_info = path_in_pack_to_abs[file_spec.path]
            is_last_file = index == len(to_upload)
            self.executor.queue(
                partial(
                    self._step_queue_upload_file,
                    file_info,
                    file_spec,
                    is_last_file,
                    upload_progress,
                )
            )
        self.executor.queue(
            partial(
                self._step_check_upload_success,
                files_to_copy,
                shaman_spec,
                upload_progress,
            )
        )

    def _step_queue_upload_file(
        self,
        file_info: _FileInfo,
        file_spec: _ShamanFileSpec,
        is_last_file: bool,
        upload_progress: ShamanUploadProgress,
    ) -> None:
        # Pre-flight check. The generated API code will load the entire file
        # into memory before sending it to the Shaman. It's faster to do a check
        # at Shaman first, to see if we need uploading at all.
        check_resp = self.shaman_api.shaman_file_store_check(
            checksum=file_spec.sha,
            filesize=file_spec.size,
        )
        if check_resp.status.value == "stored":
            log.info("  %s: skipping, already on server", file_spec.path)
            return

        # Do the 'start' reporting in a separate step, so that the Blender UI
        # can be updated for it. The 'done'/'error' reports are done at the end
        # of the file upload step, and so these don't need a separate step.
        self.executor.queue(partial(self._step_report_upload_file_start, file_info))
        self.executor.queue(
            partial(
                self._step_upload_file,
                file_info,
                file_spec,
                is_last_file,
                upload_progress,
            )
        )

    def _step_report_upload_file_start(self, file_info: _FileInfo) -> None:
        self.reporter.on_copy_start(file_info.source_path, file_info.relpath_in_pack)

    def _step_upload_file(
        self,
        file_info: _FileInfo,
        file_spec: _ShamanFileSpec,
        is_last_file: bool,
        upload_progress: ShamanUploadProgress,
    ) -> None:
        from ..manager.exceptions import ApiException

        # See whether we may be able to defer uploading this file or not.
        can_defer = bool(
            not is_last_file
            and upload_progress.num_deferred < MAX_DEFERRED_PATHS
            and not upload_progress.is_deferred(file_info.relpath_in_pack)
        )

        filename_header = _encode_original_filename_header(file_spec.path)
        try:
            with file_info.path_to_pack.open("rb") as file_reader:
                self.shaman_api.shaman_file_store(
                    checksum=file_spec.sha,
                    filesize=file_spec.size,
                    body=file_reader,
                    x_shaman_can_defer_upload=can_defer,
                    x_shaman_original_filename=filename_header,
                )
        except ApiException as ex:
            if ex.status == 425:
                # Too Early, i.e. defer uploading this file.
                log.info(
                    "  %s: someone else is uploading this file, deferring",
                    file_spec.path,
                )
                upload_progress.defer(file_info, file_spec)
                return
            elif ex.status == 417:
                # Expectation Failed; mismatch of checksum or file size.
                msg = "Error from Shaman uploading %s, code %d: %s" % (
                    file_spec.path,
                    ex.status,
                    ex.body,
                )
            else:  # Unknown error
                msg = "API exception\nHeaders: %s\nBody: %s\n" % (
                    ex.headers,
                    ex.body,
                )

            log.error(msg)
            upload_progress.failed(file_info, file_spec, msg)
            return

        self._num_files_to_transfer_done += 1
        self.reporter.on_copy_done(file_info.source_path, file_info.relpath_in_pack)

    def _step_check_upload_success(
        self,
        files_to_copy: dict[Path, _FileInfo],
        shaman_spec: _ShamanRequirementsRequest,
        upload_progress: ShamanUploadProgress,
    ) -> None:
        """See if there were any deferred or failed files.

        If there were, re-queue the uploading of the remaining files.
        Unless the number of retries has been exceeded, in which case the
        failures are final.
        """

        if upload_progress.num_deferred == 0 and upload_progress.num_failed == 0:
            # Nothing left to do, so move on to the next stage.
            self.executor.queue(partial(self._step_request_checkout, shaman_spec))
            return

        upload_progress.retry_counter += 1
        if upload_progress.retry_counter >= upload_progress.max_retries:
            # Failed uploads have really failed now.
            #
            # Deferred uploads shouldn't be mentioned, because they only get
            # deferred on the first upload attempt. After that, if they fail,
            # they get into the failures.
            for fileinfo, _, errormsg in upload_progress.failures.values():
                self.reporter.on_copy_error(
                    fileinfo.source_path, fileinfo.relpath_in_pack, errormsg
                )
            return

        # Retry uploading.
        self.executor.queue(
            partial(
                self._step_queue_uploads_of_files,
                files_to_copy,
                shaman_spec,
                upload_progress,
            )
        )

    def _step_request_checkout(self, shaman_spec: _ShamanRequirementsRequest) -> None:
        """Ask the Shaman to create a checkout of this BAT pack."""
        assert self.checkout_path

        from ..manager.exceptions import ApiException
        from ..manager.models import ShamanCheckout, ShamanCheckoutResult

        log.info(
            "Requesting checkout at Shaman for checkout_path=%s", self.checkout_path
        )

        checkoutRequest = ShamanCheckout(
            files=shaman_spec.files,
            checkout_path=str(self.checkout_path),
        )

        try:
            result: ShamanCheckoutResult = self.shaman_api.shaman_checkout(
                checkoutRequest
            )
        except ApiException as ex:
            if ex.status == 424:  # Files were missing
                msg = "We did not upload some files, checkout aborted"
            elif ex.status == 409:  # Checkout already exists
                msg = "There is already an existing checkout at %s" % self.checkout_path
            else:  # Unknown error
                msg = "API exception\nHeaders: %s\nBody: %s\n" % (
                    ex.headers,
                    ex.body,
                )
            log.error(msg)
            self.reporter.on_error(msg)
            return

        log.info("Shaman created checkout at %s", result.checkout_path)
        self._checkout_path_final = result.checkout_path

    def _send_spec_to_shaman(
        self,
        requirements: _ShamanRequirementsRequest,
    ) -> list[_ShamanFileSpec] | None:
        """Send the checkout definition file to the Shaman.

        :return: A list of file specs that still need to be uploaded, or
            None if there was an error.
        """
        from ..manager.exceptions import ApiException
        from ..manager.models import ShamanRequirementsResponse

        requested_relpaths = {file.path for file in requirements.files}

        try:
            resp = self.shaman_api.shaman_checkout_requirements(requirements)
        except ApiException as ex:
            # TODO: the body should be JSON of a predefined type, parse it to get the actual message.
            msg = "Error from Shaman, code %d: %s" % (ex.status, ex.body)
            log.error(msg)
            self.reporter.on_error(msg)
            return None
        assert isinstance(resp, ShamanRequirementsResponse)

        # Go over the response, and create two queues for uploading. Any file
        # that's already being uploaded by somebody else will be put in the
        # low-priority queue.
        to_upload_normal_prio: list[_ShamanFileSpec] = []
        to_upload_low_prio: list[_ShamanFileSpec] = []
        for file_spec in resp.files:
            if file_spec.path not in requested_relpaths:
                msg = (
                    "Shaman requested path we did not intend to upload: %r" % file_spec
                )
                log.error(msg)
                self.reporter.on_error(msg)
                return None

            log.debug("   %s: %s", file_spec.status, file_spec.path)
            status = file_spec.status.value
            if status == "unknown":
                to_upload_normal_prio.append(file_spec)
            elif status == "uploading":
                to_upload_low_prio.append(file_spec)
            else:
                msg = "Unknown status in response from Shaman: %r" % file_spec
                log.error(msg)
                self.reporter.on_error(msg)
                return None

        # Randomize the two lists, so that when two clients upload similar sets
        # of files, collissions are minimized.
        random.shuffle(to_upload_normal_prio)
        random.shuffle(to_upload_low_prio)

        return to_upload_normal_prio + to_upload_low_prio


def _encode_original_filename_header(filename: str) -> str:
    """Encode the 'original filename' as valid HTTP Header.

    See the specs for the X-Shaman-Original-Filename header in the OpenAPI
    operation `shamanFileStore`, defined in flamenco-openapi.yaml.
    """

    # This is a no-op when the filename is already in ASCII.
    fake_header = email.header.Header(maxlinelen=0)
    fake_header.append(filename, charset="utf-8")
    encoded_header = fake_header.encode()

    # Make sure that there are no newlines in the returned value.
    # HTTP Header line folding is obsolete, see RFC9112 section 5.2 in
    # https://www.rfc-editor.org/rfc/rfc9112#name-obsolete-line-folding
    assert "\n" not in encoded_header

    return encoded_header