# License: Apache-2.0 License
# Created by: Patrick Lumbantobing, VertoX-AI
# Copyright (c) 2026 VertoX-AI. All rights reserved.
#
# This work is licensed under the Apache-2.0 License.
# To view a copy of this license, visit
# https://www.apache.org/licenses/LICENSE-2.0

"""Manage safe per-run temporary workspaces and optional diagnostic retention."""

from __future__ import annotations

import shutil
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path

from speech_text_data_aligners.core.errors import ConfigurationError


class ArtifactPolicy(StrEnum):
    """Select when a completed temporary workspace should be retained."""

    NEVER = "never"
    ON_FAILURE = "on_failure"
    ALWAYS = "always"


@dataclass(slots=True)
class Workspace:
    """Expose a scoped temporary path and its eventual retained copy.

    Parameters
    ----------
        path: Unique temporary work directory.
        retained_path: Final copied artifact directory, populated during cleanup.

    Invariants:
        ``path`` exists only within the surrounding :func:`workspace` context.

    """

    path: Path
    retained_path: Path | None = None
    failed: bool = False

    def mark_failed(self) -> None:
        """Mark a handled failure so ``ON_FAILURE`` policy retains artifacts.

        Side Effects:
            Changes this workspace's cleanup disposition.
        """
        self.failed = True


@contextmanager
def workspace(
    *,
    run_id: str,
    policy: ArtifactPolicy = ArtifactPolicy.NEVER,
    temporary_root: Path | None = None,
    retention_root: Path | None = None,
) -> Iterator[Workspace]:
    """Create a unique temporary workspace and safely apply retention policy.

    Parameters
    ----------
        run_id: Path-component-free identifier used for a retained directory name.
        policy: Retain never, only after failure, or always.
        temporary_root: Optional parent for temporary workspace creation.
        retention_root: Required destination parent when retention can occur.

    Yields
    ------
        Mutable workspace status object valid for the context lifetime.

    Raises
    ------
        ConfigurationError: If identifiers/roots are unsafe or retention would
            overwrite an existing artifact directory.

    Side Effects:
        Creates a temporary directory, optionally copies it to ``retention_root``, and
        always removes the temporary directory at context exit.

    """
    if not run_id or Path(run_id).name != run_id or run_id in {".", ".."}:
        raise ConfigurationError("workspace run_id must be one safe path component")
    if policy is not ArtifactPolicy.NEVER and retention_root is None:
        raise ConfigurationError("retention_root is required by artifact policy")
    if temporary_root is not None:
        temporary_root.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(
        prefix=f"speech-aligners-{run_id}-",
        dir=temporary_root,
    ) as temporary_name:
        state = Workspace(Path(temporary_name))
        raised = False
        try:
            yield state
        except BaseException:
            raised = True
            state.failed = True
            raise
        finally:
            should_retain = policy is ArtifactPolicy.ALWAYS or (
                policy is ArtifactPolicy.ON_FAILURE and (raised or state.failed)
            )
            if should_retain:
                assert retention_root is not None
                retention_root.mkdir(parents=True, exist_ok=True)
                destination = retention_root / run_id
                if destination.exists():
                    raise ConfigurationError(
                        f"retained artifact destination already exists: {destination}"
                    )
                shutil.copytree(state.path, destination)
                state.retained_path = destination