"""Checkpoint decorator for durable work boundaries.

A checkpoint is a unit of work inside a flow whose outcome is persisted.
Successful outputs become artifacts; failures are recorded for retry.
"""

from __future__ import annotations

import keyword
import sys
from collections.abc import Callable, Mapping, Sequence
from contextlib import ExitStack
from functools import update_wrapper, wraps
from itertools import islice
from typing import Any, Literal, cast, overload

from zenml.config.retry_config import StepRetryConfig
from zenml.enums import StepRuntime, StepType
from zenml.execution.pipeline.dynamic.outputs import OutputArtifact
from zenml.execution.pipeline.dynamic.run_context import DynamicPipelineRunContext
from zenml.pipelines.compilation_context import PipelineCompilationContext
from zenml.steps.step_context import StepContext
from zenml.steps.step_decorator import step

from kitaru import events
from kitaru._source_aliases import (
    build_checkpoint_registration_name,
    build_checkpoint_source_alias,
    callable_name,
)
from kitaru.analytics import AnalyticsEvent, track
from kitaru.errors import KitaruContextError, KitaruUsageError
from kitaru.runtime import (
    _checkpoint_scope,
    _flow_scope,
    _get_current_checkpoint_id,
    _get_current_execution_id,
    _get_current_flow,
    _get_zenml_checkpoint_id,
    _get_zenml_execution_id,
    _get_zenml_flow_name,
    _is_inside_flow,
)

CheckpointSurface = Literal["call", "submit", "map", "product"]

_CHECKPOINT_OUTSIDE_FLOW_ERROR = "Checkpoints can only run inside a @flow."
_CHECKPOINT_NESTED_ERROR = (
    "Nested checkpoint calls are not supported in the Kitaru MVP."
)
_CHECKPOINT_CONCURRENT_OUTSIDE_FLOW_ERROR = (
    "Concurrent checkpoint execution is only available inside a running @flow."
)
_CHECKPOINT_OUTPUT_HANDLE_MESSAGE = (
    "This is a Kitaru checkpoint output handle for `{checkpoint}.{output}`, "
    "not the materialized value. Call `.load()` to materialize the value."
)
_KITARU_EXTRA_NAMESPACE = "kitaru"
_CHECKPOINT_BOUNDARY_KEY = "boundary"
_CHECKPOINT_BOUNDARY_VALUE = "checkpoint"
_FLOW_RESULT_CANDIDATE_KEY = "flow_result_candidate"


def _is_checkpoint_output_handle(value: Any) -> bool:
    return isinstance(value, OutputArtifact)


def _checkpoint_output_handle_validation_message(
    value: Any,
    *,
    field_name: str,
    expected: str,
) -> str:
    checkpoint_name = getattr(value, "step_name", None)
    output_name = getattr(value, "output_name", None)
    handle_context = "a Kitaru checkpoint output handle"
    if checkpoint_name and output_name:
        handle_context = (
            f"a Kitaru checkpoint output handle for `{checkpoint_name}.{output_name}`"
        )

    return (
        f"{field_name} received {handle_context}, not {expected}. "
        "In a flow body, call `.load()` before constructing this request when "
        "you need the concrete value. To preserve the durable data edge, pass "
        "the original handle into a downstream `@checkpoint` instead."
    )


def _raise_if_checkpoint_output_handle(
    value: Any,
    *,
    field_name: str,
    expected: str,
) -> None:
    if _is_checkpoint_output_handle(value):
        raise ValueError(
            _checkpoint_output_handle_validation_message(
                value,
                field_name=field_name,
                expected=expected,
            )
        )


_CHECKPOINT_HANDLE_SCAN_MAX_VALUES = 1_000


def _format_checkpoint_value_path(parent: str, key: Any) -> str:
    if isinstance(key, str) and key.isidentifier() and not keyword.iskeyword(key):
        return f"{parent}.{key}"
    return f"{parent}[{key!r}]"


def _raise_checkpoint_handle_scan_limit_exceeded(field_name: str) -> None:
    raise ValueError(
        "Kitaru could not validate "
        f"{field_name} for checkpoint output handles because the value contains "
        f"more than {_CHECKPOINT_HANDLE_SCAN_MAX_VALUES} nested values. "
        "Pass a smaller value, or load any checkpoint output handles with `.load()` "
        "before passing the value here."
    )


def _raise_if_checkpoint_output_handle_in_value(
    value: Any,
    *,
    field_name: str,
    expected: str,
) -> None:
    stack: list[tuple[Any, str]] = [(value, field_name)]
    seen_containers: set[int] = set()
    visited_values = 0
    scan_limit_exceeded = False

    while stack:
        item, item_field_name = stack.pop()
        visited_values += 1
        if visited_values > _CHECKPOINT_HANDLE_SCAN_MAX_VALUES:
            _raise_checkpoint_handle_scan_limit_exceeded(field_name)

        _raise_if_checkpoint_output_handle(
            item,
            field_name=item_field_name,
            expected=expected,
        )

        if isinstance(item, Mapping):
            container_id = id(item)
            if container_id in seen_containers:
                continue
            seen_containers.add(container_id)

            remaining_values = _CHECKPOINT_HANDLE_SCAN_MAX_VALUES - (
                visited_values + len(stack)
            )
            if remaining_values <= 0:
                scan_limit_exceeded = scan_limit_exceeded or bool(item)
                continue
            scan_limit_exceeded = scan_limit_exceeded or len(item) > remaining_values
            children = list(islice(item.items(), remaining_values))
            stack.extend(
                (child, _format_checkpoint_value_path(item_field_name, key))
                for key, child in reversed(children)
            )
            continue

        if isinstance(item, (list, tuple)):
            container_id = id(item)
            if container_id in seen_containers:
                continue
            seen_containers.add(container_id)

            remaining_values = _CHECKPOINT_HANDLE_SCAN_MAX_VALUES - (
                visited_values + len(stack)
            )
            if remaining_values <= 0:
                scan_limit_exceeded = scan_limit_exceeded or bool(item)
                continue
            scan_limit_exceeded = scan_limit_exceeded or len(item) > remaining_values
            stack.extend(
                (child, f"{item_field_name}[{index}]")
                for index, child in reversed(list(enumerate(item[:remaining_values])))
            )

    if scan_limit_exceeded:
        _raise_checkpoint_handle_scan_limit_exceeded(field_name)


class _KitaruOutputArtifact(OutputArtifact):
    """ZenML output artifact with a Kitaru-facing display string."""

    def __str__(self) -> str:
        """Return a short, helpful message instead of raw artifact metadata."""
        return _CHECKPOINT_OUTPUT_HANDLE_MESSAGE.format(
            checkpoint=self.step_name,
            output=self.output_name,
        )

    def __repr__(self) -> str:
        """Mirror ``str()`` for accidental repr conversions in user code."""
        return str(self)


def _with_kitaru_output_display(value: Any) -> Any:
    """Attach Kitaru-friendly display behavior to exposed output handles."""
    if isinstance(value, _KitaruOutputArtifact):
        return value
    if isinstance(value, OutputArtifact):
        return _KitaruOutputArtifact.model_construct(
            _fields_set=value.model_fields_set,
            **value.__dict__,
        )
    if isinstance(value, tuple):
        return tuple(_with_kitaru_output_display(item) for item in value)
    return value


def _register_checkpoint_source_alias(
    *,
    func: Callable[..., Any],
    alias: str,
    step_obj: Any,
) -> None:
    """Register the ZenML step object under a module-level alias."""
    module = sys.modules.get(func.__module__)
    if module is None:
        return
    setattr(module, alias, step_obj)


def _normalize_retries(retries: int) -> int:
    """Validate and normalize checkpoint retries."""
    if retries < 0:
        raise KitaruUsageError("Checkpoint retries must be >= 0.")
    return retries


def _to_retry_config(retries: int) -> StepRetryConfig | None:
    """Convert retry count to ZenML retry config."""
    if retries == 0:
        return None
    return StepRetryConfig(max_retries=retries)


def _build_checkpoint_extra(
    checkpoint_type: str | None,
    *,
    flow_result_candidate: bool | None = None,
) -> dict[str, Any]:
    """Build namespaced step metadata for dashboard rendering."""
    payload: dict[str, Any] = {_CHECKPOINT_BOUNDARY_KEY: _CHECKPOINT_BOUNDARY_VALUE}
    if checkpoint_type is not None:
        payload["type"] = checkpoint_type
    if flow_result_candidate is not None:
        payload[_FLOW_RESULT_CANDIDATE_KEY] = flow_result_candidate
    return {_KITARU_EXTRA_NAMESPACE: payload}


_KNOWN_STEP_TYPES: dict[str, StepType] = {
    "llm_call": StepType.LLM_CALL,
    "tool_call": StepType.TOOL_CALL,
}


def _to_step_type(checkpoint_type: str | None) -> StepType | None:
    """Map well-known checkpoint types to ZenML's StepType enum."""
    if checkpoint_type is None:
        return None
    return _KNOWN_STEP_TYPES.get(checkpoint_type)


def _normalize_flow_result_candidate(value: bool | None) -> bool | None:
    """Validate the private terminal-result eligibility flag."""
    if value is None:
        return None
    if not isinstance(value, bool):
        raise KitaruUsageError(
            "Synthetic checkpoint flow_result_candidate must be a bool when provided."
        )
    return value


_KNOWN_STEP_RUNTIMES: dict[str, StepRuntime] = {
    "inline": StepRuntime.INLINE,
    "isolated": StepRuntime.ISOLATED,
}
_RUNTIME_OPTIONS = ", ".join(_KNOWN_STEP_RUNTIMES)


def _runtime_error(value: object) -> KitaruUsageError:
    return KitaruUsageError(
        f"Unsupported checkpoint runtime {value!r}. "
        f"Expected one of: {_RUNTIME_OPTIONS}."
    )


def _normalize_runtime(runtime: StepRuntime | str | None) -> StepRuntime | None:
    """Validate and normalize checkpoint runtime input.

    Accepts ``None``, a :class:`StepRuntime` enum member, or a case-insensitive
    string (``"inline"`` / ``"isolated"``).
    """
    if runtime is None:
        return None
    if isinstance(runtime, StepRuntime):
        return runtime
    if not isinstance(runtime, str):
        raise _runtime_error(runtime)
    normalized = runtime.strip().lower()
    resolved = _KNOWN_STEP_RUNTIMES.get(normalized)
    if resolved is None:
        raise _runtime_error(runtime)
    return resolved


def _wrap_entrypoint(
    func: Callable[..., Any],
    *,
    checkpoint_type: str | None,
) -> Callable[..., Any]:
    """Wrap the user function with Kitaru checkpoint runtime scope."""

    checkpoint_name = callable_name(func)

    @wraps(func)
    def _wrapped(*args: Any, **kwargs: Any) -> Any:
        current_flow = _get_current_flow()
        execution_id = _get_current_execution_id() or _get_zenml_execution_id()
        checkpoint_id = _get_current_checkpoint_id() or _get_zenml_checkpoint_id()

        with ExitStack() as scope_stack:
            if current_flow is None:
                scope_stack.enter_context(
                    _flow_scope(
                        name=_get_zenml_flow_name(),
                        execution_id=execution_id,
                    )
                )

            scope_stack.enter_context(
                _checkpoint_scope(
                    name=checkpoint_name,
                    checkpoint_type=checkpoint_type,
                    execution_id=execution_id,
                    checkpoint_id=checkpoint_id,
                )
            )
            events._publish_checkpoint_started()
            try:
                result = func(*args, **kwargs)
            except BaseException as exc:
                events._publish_checkpoint_failed(exc)
                raise
            events._publish_checkpoint_returned()
            return result

    return _wrapped


class _CheckpointDefinition:
    """Callable wrapper returned by `@checkpoint`."""

    def __init__(
        self,
        func: Callable[..., Any],
        *,
        retries: int,
        checkpoint_type: str | None,
        runtime: StepRuntime | str | None,
        cache: bool | None,
        flow_result_candidate: bool | None,
    ) -> None:
        """Initialize a Kitaru checkpoint wrapper."""
        self._func = func
        self._checkpoint_type = checkpoint_type
        self._default_retries = _normalize_retries(retries)
        self._runtime = _normalize_runtime(runtime)
        self._cache = cache
        self._cache_explicitly_set = cache is not None
        flow_result_candidate = _normalize_flow_result_candidate(flow_result_candidate)

        wrapped_entrypoint = _wrap_entrypoint(
            func,
            checkpoint_type=checkpoint_type,
        )
        func_name = callable_name(func)
        registration_name = build_checkpoint_registration_name(func_name)
        source_alias = build_checkpoint_source_alias(func_name)
        aliasable_entrypoint = cast(Any, wrapped_entrypoint)
        aliasable_entrypoint.__name__ = source_alias
        aliasable_entrypoint.__qualname__ = source_alias

        self._step = step(
            name=registration_name,
            retry=_to_retry_config(self._default_retries),
            enable_cache=self._cache,
            extra=_build_checkpoint_extra(
                checkpoint_type,
                flow_result_candidate=flow_result_candidate,
            ),
            step_type=_to_step_type(checkpoint_type),
            runtime=self._runtime,
        )(wrapped_entrypoint)
        _register_checkpoint_source_alias(
            func=func,
            alias=source_alias,
            step_obj=self._step,
        )

        update_wrapper(self, func)

    def _assert_call_allowed(self) -> None:
        """Validate that checkpoint call semantics match Kitaru MVP rules."""
        if PipelineCompilationContext.is_active() and _is_inside_flow():
            return

        if StepContext.is_active():
            raise KitaruContextError(_CHECKPOINT_NESTED_ERROR)

        if DynamicPipelineRunContext.is_active() and _is_inside_flow():
            return

        raise KitaruContextError(_CHECKPOINT_OUTSIDE_FLOW_ERROR)

    def _assert_submit_allowed(self) -> None:
        """Validate that checkpoint submission is legal in the current context."""
        if StepContext.is_active():
            raise KitaruContextError(_CHECKPOINT_NESTED_ERROR)

        if not DynamicPipelineRunContext.is_active() or not _is_inside_flow():
            raise KitaruContextError(_CHECKPOINT_CONCURRENT_OUTSIDE_FLOW_ERROR)

    def _track_invocation(self, *, method: CheckpointSurface) -> None:
        """Emit checkpoint-level telemetry for invocation surfaces."""
        track(
            AnalyticsEvent.CHECKPOINT_INVOKED,
            {
                "method": method,
                "cache_explicitly_set": self._cache_explicitly_set,
            },
        )

    def __call__(
        self,
        *args: Any,
        id: str | None = None,
        after: Any | Sequence[Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        """Call the checkpoint with context guardrails."""
        self._assert_call_allowed()
        result = self._step(*args, id=id, after=after, **kwargs)
        self._track_invocation(method="call")
        return _with_kitaru_output_display(result)

    def submit(
        self,
        *args: Any,
        id: str | None = None,
        after: Any | Sequence[Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        """Submit the checkpoint concurrently inside a running flow."""
        self._assert_submit_allowed()
        result = self._step.submit(*args, id=id, after=after, **kwargs)
        self._track_invocation(method="submit")
        return result

    def map(
        self,
        *args: Any,
        after: Any | Sequence[Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        """Map checkpoint invocations inside a running flow."""
        self._assert_submit_allowed()
        result = self._step.map(*args, after=after, **kwargs)
        self._track_invocation(method="map")
        return result

    def product(
        self,
        *args: Any,
        after: Any | Sequence[Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        """Map checkpoint invocations as a cartesian product in a running flow."""
        self._assert_submit_allowed()
        result = self._step.product(*args, after=after, **kwargs)
        self._track_invocation(method="product")
        return result


@overload
def checkpoint(func: Callable[..., Any], /) -> _CheckpointDefinition: ...


@overload
def checkpoint(
    *,
    retries: int = 0,
    type: str | None = None,
    runtime: str | None = None,
    cache: bool | None = None,
) -> Callable[[Callable[..., Any]], _CheckpointDefinition]: ...


def _synthetic_checkpoint(
    func: Callable[..., Any] | None = None,
    *,
    retries: int = 0,
    type: str | None = None,
    runtime: str | None = None,
    cache: bool | None = None,
    flow_result_candidate: bool | None = None,
) -> _CheckpointDefinition | Callable[[Callable[..., Any]], _CheckpointDefinition]:
    """Create a checkpoint with optional internal metadata."""
    checkpoint_type = type

    def _decorate(target: Callable[..., Any]) -> _CheckpointDefinition:
        return _CheckpointDefinition(
            target,
            retries=retries,
            checkpoint_type=checkpoint_type,
            runtime=runtime,
            cache=cache,
            flow_result_candidate=flow_result_candidate,
        )

    if func is not None:
        return _decorate(func)

    return _decorate


def checkpoint(
    func: Callable[..., Any] | None = None,
    *,
    retries: int = 0,
    type: str | None = None,
    runtime: str | None = None,
    cache: bool | None = None,
) -> _CheckpointDefinition | Callable[[Callable[..., Any]], _CheckpointDefinition]:
    """Mark a function as a durable checkpoint.

    Can be used as a bare decorator or with arguments:

    ```python
    from kitaru import checkpoint

    @checkpoint
    def my_step(): ...

    @checkpoint(retries=3, type="llm_call")
    def my_step(): ...

    @checkpoint(runtime="isolated")
    def heavy_step(): ...
    ```

    Args:
        func: Optional function for bare decorator use.
        retries: Number of checkpoint-level retries on failure.
        type: Checkpoint type for dashboard visualization.
        runtime: Execution runtime for this checkpoint. Accepts ``"inline"``
            or ``"isolated"`` (case-insensitive). When set to ``"isolated"``,
            the checkpoint runs in its own container on remote orchestrators
            that support it. ``None`` (the default) lets the orchestrator
            decide.
        cache: Optional per-checkpoint cache override. ``True`` enables cache
            for this checkpoint, ``False`` disables it, and ``None`` defers to
            higher-level flow/default cache behavior. Useful for forcing one
            expensive step to re-run while the rest of the flow stays cached.
    Returns:
        The wrapped checkpoint object or a decorator that returns it.
    """
    return _synthetic_checkpoint(
        func,
        retries=retries,
        type=type,
        runtime=runtime,
        cache=cache,
    )
