"""
MIT License Copyright (c) 2025-present June

Permission is hereby granted, free of
charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without
restriction, including without limitation the rights to use, copy, modify, merge,
publish, distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to the
following conditions:

The above copyright notice and this permission notice
(including the next paragraph) shall be included in all copies or substantial
portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO
EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import logging
from dataclasses import dataclass, field
from enum import Enum
from functools import wraps
from typing import Any, Callable, Dict, List, Optional, Union

from whatsapp.bindings import ChatContext
from whatsapp.messages import InteractiveTextMessage, TextMessage
from whatsapp.middleware.storage import SessionStore

logger = logging.getLogger(__name__)


# --- Dialog Class ---
# Dialog class is used for creating dialogs and managing dialog flow.
class DialogTurnStatus(Enum):
    """Enumeration of possible dialog turn statuses.

    Attributes:
        Waiting: Dialog is waiting for user input.
        Complete: Dialog completed successfully.
        Empty: No active dialog (stack is empty).
        Cancelled: Dialog was cancelled by user or system.
        Invalid: User input was invalid.
    """

    Waiting = "Waiting"
    Complete = "Complete"
    Empty = "Empty"
    Cancelled = "Cancelled"
    Invalid = "Invalid"


class DialogTurnResult:
    """Result of executing a dialog turn/step.

    Attributes:
        status: The DialogTurnStatus indicating the outcome.
        result: Optional data returned from the dialog step.
    """

    def __init__(self, status: DialogTurnStatus, result: Any = None):
        """
        Initialize a dialog turn result.

        Args:
            status: The turn status.
            result: Optional result/data from the turn.
        """
        self.status = status
        self.result = result

    def __repr__(self):
        """String representation of the result."""
        return f"DialogTurnResult(status={self.status}, result={self.result})"


class Dialog:
    """Manages a conversational dialog with waterfall-style steps.

    A dialog is a sequence of steps that collect information from the user,
    validate responses, and manage conversation flow. Features include:
    - Waterfall pattern with sequential steps
    - Input validation with retry logic
    - Dialog transfers to other dialogs
    - Session state storage
    - Decorators for organizing steps

    Class-level attributes store step definitions and metadata for each dialog subclass.
    """

    # class-level registry for all dialogs
    _registry: Dict[str, type] = {}
    _main_dialog_id: Optional[str] = None
    _step_counter: int = 0  # counter to track step order

    # class-level attributes set by decorators (per-subclass)
    _steps: List[Callable] = []
    _store_mappings: Dict[str, str] = {}  # step_name -> store_key
    _validators: Dict[str, dict] = {}  # step_name -> {validator, prompt, retry}
    _transfers: Dict[str, str] = {}  # step_name -> target_dialog_id
    _media_steps: Dict[str, List[str]] = {}  # step_name -> accepted media types (empty = all)
    _is_main: bool = False

    def __init__(self, ctx: "DialogContext", options: dict = {}):
        """
        Initialize a dialog instance.

        Args:
            ctx: The dialog context managing this dialog.
            options: Optional initial options/parameters for the dialog.
        """
        self.id = self.__class__.__name__
        self.ctx = ctx
        self.options = options
        self.state = {}
        self.steps = self.__class__._steps.copy() if self.__class__._steps else []

    # --- Dialog Context Wrappers ---
    # - Orchestration Methods -

    async def begin_dialog(self, dc: "DialogContext", options: Any = None) -> DialogTurnResult:
        """
        Begin the dialog, initializing state and running the first step.

        Args:
            dc: The dialog context.
            options: Optional initial options for the dialog.

        Returns:
            DialogTurnResult from executing the first step.
        """
        # populate initial data in the active dialog instance (already on stack)
        state = {"step_index": 0, "values": {}, "options": options or {}}
        if dc.active_dialog:
            dc.active_dialog.state = state

        # initiate dialog
        return await self.on_begin_dialog(dc, options)

    async def on_begin_dialog(self, dc: "DialogContext", options: Any = None) -> DialogTurnResult:
        """Executes the initial step when dialog begins.

        Args:
            dc: The dialog context.
            options: Optional initial options.

        Returns:
            DialogTurnResult from the begin step.
        """
        return await self.run_step(dc, "begin", None)

    async def continue_dialog(self, dc: "DialogContext") -> DialogTurnResult:
        """Continue dialog execution with user's latest input.

        Args:
            dc: The dialog context.

        Returns:
            DialogTurnResult from the current step.
        """
        return await self.resume_dialog(dc, "continue", dc.ctx.last_incoming_text)

    async def resume_dialog(self, dc: "DialogContext", reason: Any, result: Any = None) -> DialogTurnResult:
        """Resume dialog after waiting for user input.

        Args:
            dc: The dialog context.
            reason: Reason for resuming (e.g., "continue", "end_dialog").
            result: The user's response/input.

        Returns:
            DialogTurnResult from the current step.
        """
        state = dc.active_dialog.state
        state["step_index"] += 1  # increments step after waiting for input
        return await self.run_step(dc, reason, result)

    async def run_step(self, dc: "DialogContext", reason: Any = None, result: Any = None) -> DialogTurnResult:
        """
        Executes the current waterfall step.

        Before proceeding, a validator is run if one is in place.
        Handles storing results, validation, transfers, and step execution.

        Args:
            dc: The dialog context.
            reason: Reason for running (e.g., "begin", "continue").
            result: User input/response to validate and process.

        Returns:
            DialogTurnResult indicating the outcome.
        """
        state = dc.active_dialog.state
        index = state["step_index"]

        if index >= len(self.steps):
            return await dc.end_dialog(result)  # end of waterfall

        step = self.steps[index]
        step_name = step.__name__

        # -- handle store mapping --
        # save previous result to dialog state values
        if index > 0 and result is not None:
            prev_step = self.steps[index - 1]
            prev_step_name = prev_step.__name__
            if prev_step_name in self._store_mappings:
                store_key = self._store_mappings[prev_step_name]
                state["values"][store_key] = result

        # handle validation for previous step
        if index > 0 and result is not None:
            # get previous step validator
            prev_step = self.steps[index - 1]
            prev_step_name = prev_step.__name__

            # check if validator exists and run
            if prev_step_name in self._validators:
                validator_info = self._validators[prev_step_name]
                validator_fn = validator_info.get("validator")
                prompt = validator_info.get("prompt", "Invalid input")  # set default prompt
                retry = validator_info.get("retry", 3)  # set default retry
                retry_count = state.get("_retry_count", 0)  # get retry count

                # validator check fails
                if validator_fn and not await self._run_validator(validator_fn, result):
                    if retry_count < retry:
                        state["_retry_count"] = retry_count + 1
                        state["step_index"] -= 1  # Go back to repeat
                        await dc.ctx.send_message(TextMessage(dc.ctx, prompt))
                        return DialogTurnResult(DialogTurnStatus.Waiting)
                    else:
                        state["_retry_count"] = 0
                        return DialogTurnResult(DialogTurnStatus.Invalid)
                state["_retry_count"] = 0

        # check for transfer after step execution
        if step_name in self._transfers:
            target_dialog_id = self._transfers[step_name]
            return await dc.replace_dialog(target_dialog_id)

        # execute the step
        step_result = await step(self, dc)

        # if step returned Complete, automatically end the dialog
        if step_result and step_result.status == DialogTurnStatus.Complete:
            return await dc.end_dialog(step_result.result)

        return step_result

    # --- Class Method Wrappers ---
    async def _run_validator(self, validator: Callable, value: Any) -> bool:
        """Run a validator function (sync or async).

        Args:
            validator: The validator function.
            value: The value to validate.

        Returns:
            True if validation passes, False otherwise.
        """
        import asyncio

        if asyncio.iscoroutinefunction(validator):
            return await validator(value)
        return validator(value)

    # --- Class Method Decorators ---
    @classmethod
    def set(cls, name: str = None, main: bool = False):
        """
        Class decorator to register a dialog.

        Initializes class-level collections for dialog metadata and registers
        the dialog in the global registry.

        Args:
            name: Optional custom dialog ID. Defaults to class name.
            main: If True, marks this as the main/entry dialog.

        Returns:
            The decorator function.

        Usage:
            @Dialog.set(name="MainMenu", main=True)
            class MainMenuDialog(Dialog):
                ...
        """

        def decorator(dialog_cls):
            dialog_id = name or dialog_cls.__name__
            dialog_cls.id = dialog_id
            dialog_cls._is_main = main

            # initializes class-level collections for other methods if not present

            # steps
            if not hasattr(dialog_cls, "_steps") or dialog_cls._steps is Dialog._steps:
                dialog_cls._steps = []

            # store mappings
            if not hasattr(dialog_cls, "_store_mappings") or dialog_cls._store_mappings is Dialog._store_mappings:
                dialog_cls._store_mappings = {}

            # validators
            if not hasattr(dialog_cls, "_validators") or dialog_cls._validators is Dialog._validators:
                dialog_cls._validators = {}

            # transfers
            if not hasattr(dialog_cls, "_transfers") or dialog_cls._transfers is Dialog._transfers:
                dialog_cls._transfers = {}

            # media steps
            if not hasattr(dialog_cls, "_media_steps") or dialog_cls._media_steps is Dialog._media_steps:
                dialog_cls._media_steps = {}

            # main registry
            Dialog._registry[dialog_id] = dialog_cls
            if main:
                Dialog._main_dialog_id = dialog_id

            # finalize the dialog to collect steps
            Dialog._finalize_dialog(dialog_cls)

            logger.info(f"Registered dialog: {dialog_id} (main={main})")
            return dialog_cls

        return decorator

    @classmethod
    def step(cls):
        """
        Mark a method as a waterfall step.

        Steps are executed in the order they are defined (based on decoration order).
        Each step can prompt for user input, validate responses, store values, etc.

        Usage:
            @Dialog.step()
            async def greeting_step(self, ctx):
                ...
        """
        # capture the current counter value at decoration time
        order = cls._step_counter
        cls._step_counter += 1

        def decorator(func):
            func._is_dialog_step = True
            func._step_order = order  # assign the order when decorator is applied

            @wraps(func)
            async def wrapper(self, dc: "DialogContext"):
                dc.get_value("")  # dummy call to warm up cache
                return await func(self, dc)

            wrapper._is_dialog_step = True
            wrapper._step_order = order  # Preserve order on wrapper

            # preserve metadata attributes from the original function
            # in order to keep step decorator while declaring dialog
            wrapper._store_key = getattr(func, "_store_key", None)
            wrapper._validator = getattr(func, "_validator", None)
            wrapper._transfer_to = getattr(func, "_transfer_to", None)

            return wrapper

        return decorator

    @classmethod
    def store(cls, key: str):
        """
        Store the user's response in session storage.

        The response will be saved to the dialog state under the given key.
        Can be retrieved in subsequent steps using dialog.get_value(key).

        Args:
            key: The storage key to save the value under.

        Usage:
            @Dialog.step()
            @Dialog.store("user_name")
            async def ask_name(self, ctx):
                await ctx.send_message("What is your name?")
        """

        def decorator(func):
            func._store_key = key

            @wraps(func)
            async def wrapper(self, dc: "DialogContext"):
                return await func(self, dc)

            wrapper._store_key = key
            wrapper._is_dialog_step = getattr(func, "_is_dialog_step", False)

            # preserve all metadata attributes from the original function
            wrapper._validator = getattr(func, "_validator", None)
            wrapper._transfer_to = getattr(func, "_transfer_to", None)
            return wrapper

        return decorator

    @classmethod
    def validate(cls, validator: Callable, prompt: str = "Invalid input", retry: int = 3):
        """
        Validate user input before proceeding to the next step.

        If validation fails, optionally prompts the user to retry up to N times.
        If all retries are exhausted, returns DialogTurnStatus.Invalid.

        Args:
            validator: Function to validate input. Can be sync or async. Should return bool.
            prompt: Message to send if validation fails. Defaults to "Invalid input".
            retry: Number of retry attempts allowed. Defaults to 3.

        Usage:
            def validate_phone(value):
                return len(value) >= 10 and value.isdigit()

            @Dialog.validate(validate_phone, prompt="Please enter a valid phone number", retry=3)
            @Dialog.step()
            async def ask_phone(self, ctx):
                await ctx.send_message("What is your phone number?")
        """

        def decorator(func):
            func._validator = {"validator": validator, "prompt": prompt, "retry": retry}

            @wraps(func)
            async def wrapper(self, dc: "DialogContext"):
                return await func(self, dc)

            wrapper._validator = func._validator
            wrapper._is_dialog_step = getattr(func, "_is_dialog_step", False)
            wrapper._store_key = getattr(func, "_store_key", None)
            # Preserve all metadata attributes from the original function
            wrapper._transfer_to = getattr(func, "_transfer_to", None)
            return wrapper

        return decorator

    @classmethod
    def transfer(cls, target_dialog_id: str):
        """
        Transfer to another dialog after this step executes.

        After the decorated step completes, the current dialog will be replaced
        with the target dialog. Useful for dialog flows and conditional branching.
        The target dialog must already be registered.

        Args:
            target_dialog_id: The ID of the dialog to transfer to.

        Usage:
            @Dialog.transfer("OnBoardingDialog")
            @Dialog.step()
            async def initiate(self, ctx):
                await ctx.send_message("Transferring...")
        """

        def decorator(func):
            func._transfer_to = target_dialog_id

            @wraps(func)
            async def wrapper(self, dc: "DialogContext"):
                return await func(self, dc)

            wrapper._transfer_to = target_dialog_id
            wrapper._is_dialog_step = getattr(func, "_is_dialog_step", False)
            wrapper._store_key = getattr(func, "_store_key", None)
            # Preserve all metadata attributes from the original function
            wrapper._validator = getattr(func, "_validator", None)
            return wrapper
            wrapper._validator = getattr(func, "_validator", None)
            return wrapper

        return decorator

    @classmethod
    def accept_media(cls, *types: str, prompt: str = "Please send a supported media file."):
        """
        Mark a step as accepting media messages.

        When applied, the step only executes its body when dc.ctx.incoming_media
        is populated. If a non-media message reaches this step (e.g. text while
        the dialog is waiting for a file), the framework automatically calls
        dc.await_step(prompt) to re-prompt the user — no boilerplate needed.

        Args:
            *types: Media type strings to accept (e.g. "image", "document").
                    Pass no arguments to accept all media types.
            prompt: Re-prompt message sent when no media is received.
                    Defaults to "Please send a supported media file."

        Usage:
            @Dialog.step()
            @Dialog.accept_media("image", "document", prompt="Send an image or PDF.")
            async def upload_doc(self, dc: DialogContext):
                media = dc.ctx.incoming_media  # always populated here
                ...
        """
        accepted = list(types)

        def decorator(func):
            func._accept_media_types = accepted

            @wraps(func)
            async def wrapper(self, dc: "DialogContext"):
                if dc.ctx.incoming_media is None:
                    return await dc.await_step(prompt)
                return await func(self, dc)

            wrapper._accept_media_types = accepted
            wrapper._is_dialog_step = getattr(func, "_is_dialog_step", False)
            wrapper._step_order = getattr(func, "_step_order", 0)
            wrapper._store_key = getattr(func, "_store_key", None)
            wrapper._validator = getattr(func, "_validator", None)
            wrapper._transfer_to = getattr(func, "_transfer_to", None)
            return wrapper

        return decorator

    @classmethod
    def _finalize_dialog(cls, dialog_cls: type):
        """
        Process a dialog class after all decorators have been applied.

        Collects and sorts steps, store mappings, validators, and transfers.
        This is called automatically by the @Dialog.set decorator.

        Args:
            dialog_cls: The dialog class to finalize.
        """
        steps = []
        store_mappings = {}
        validators = {}
        transfers = {}
        media_steps = {}

        # Collect decorated methods
        for attr_name in dir(dialog_cls):
            if attr_name.startswith("_"):
                continue
            attr = getattr(dialog_cls, attr_name, None)
            if callable(attr) and getattr(attr, "_is_dialog_step", False):
                steps.append(attr)

                if hasattr(attr, "_store_key") and attr._store_key:
                    store_mappings[attr.__name__] = attr._store_key

                if hasattr(attr, "_validator") and attr._validator:
                    validators[attr.__name__] = attr._validator

                if hasattr(attr, "_transfer_to") and attr._transfer_to:
                    transfers[attr.__name__] = attr._transfer_to

                if hasattr(attr, "_accept_media_types"):
                    media_steps[attr.__name__] = attr._accept_media_types

        # sort steps by their _step_order (the order decorators were applied)
        steps.sort(key=lambda s: getattr(s, "_step_order", 0))

        dialog_cls._steps = steps
        dialog_cls._store_mappings = store_mappings
        dialog_cls._validators = validators
        dialog_cls._transfers = transfers
        dialog_cls._media_steps = media_steps

    @classmethod
    def get_main_dialog_id(cls) -> Optional[str]:
        """Get the ID of the main dialog.

        Returns:
            The dialog ID marked as main, or None if no main dialog is registered.
        """
        return cls._main_dialog_id


# --- Dialog Context ---
# Dialog Context is used for current state of dialogs within a chat session.
@dataclass
class DialogInstance:
    """Represents an active dialog instance in a conversation stack.

    Attributes:
        id: The dialog class name/ID.
        state: Dictionary containing step index, values, and metadata for the dialog.
    """

    id: str
    state: Dict[str, Any] = field(default_factory=dict)


class DialogContext:
    """Manages the dialog execution context for a single conversation.

    Maintains the dialog stack, session storage, and provides methods for
    navigating the dialog flow (begin, end, transfer, skip steps, etc.).

    Attributes:
        ctx: The chat context for this conversation.
        dialogs: Registry of available dialog classes.
        stack: The dialog execution stack (root to current).
        session_store: Storage for persisting dialog state.
    """

    def __init__(
        self,
        ctx: "ChatContext",
        dialogs: Dict[str, Dialog],
        stack: List[DialogInstance] = None,
        session_store: "SessionStore" = None,
    ):
        """
        Initialize a dialog context.

        Args:
            ctx: The chat context.
            dialogs: Dictionary of available dialogs keyed by dialog ID.
            stack: Optional existing dialog stack to resume.
            session_store: Optional storage for persisting state.
        """
        self.ctx = ctx
        self.dialogs = dialogs
        self.stack = stack or []
        self.session_store = session_store
        self._finished = False
        self.message = self.MessageHandler(ctx)

    def get_value(self, key: str, active_dialog=False) -> Optional[Any]:
        """
        Get a stored value from dialog state.

        Args:
            key: The key to look up
            active_dialog: If True, only look in the active dialog's state.

        Returns:
            The stored value, or None if not found.
        """
        stack = self.stack

        # get from current dialog (O(1))
        if active_dialog:
            if stack and (values := stack[-1].state.get("values")):
                return values.get(key)
            return None

        # get from entire stack (O(n))
        for instance in reversed(stack):  # search from oldest to newest
            if (values := instance.state.get("values")) and (val := values.get(key)) is not None:
                return val
        return None

    def set_value(self, key: str, value: Any) -> None:
        """Store a value in the active dialog's state."""
        if self.active_dialog:
            if "values" not in self.active_dialog.state:
                self.active_dialog.state["values"] = {}
            self.active_dialog.state["values"][key] = value

    @property
    def active_dialog(self) -> Optional[DialogInstance]:
        if self.stack:
            return self.stack[-1]
        return None

    def find_dialog(self, dialog_id: str) -> Optional[Dialog]:
        """Find a dialog class by ID.

        Args:
            dialog_id: The dialog ID.

        Returns:
            The dialog class, or None if not found.
        """
        return self.dialogs.get(dialog_id)

    async def begin_dialog(self, dialog_id: str, options: Any = None) -> DialogTurnResult:
        """Begin a new dialog, pushing it onto the stack.

        Args:
            dialog_id: The ID of the dialog to begin.
            options: Optional initial options for the dialog.

        Returns:
            DialogTurnResult from the dialog's begin step.

        Raises:
            ValueError: If the dialog is not registered.
        """
        dialog_cls = self.find_dialog(dialog_id)
        if not dialog_cls:
            logger.error(f"Dialog with id '{dialog_id}' not found.")
            raise ValueError(f"Dialog with id '{dialog_id}' not found.")

        # Reset _finished since we're starting a new dialog
        self._finished = False

        # Push new instance
        instance = DialogInstance(id=dialog_id)
        self.stack.append(instance)

        # Instantiate the dialog class with this context
        dialog = dialog_cls(ctx=self, options=options or {})

        logger.info(f"Beginning dialog: {dialog_id}")
        return await dialog.begin_dialog(self, options)

    async def replace_dialog(self, dialog_id: str, options: Any = None) -> DialogTurnResult:
        """Replace the current dialog with a different one.

        Pops the current dialog off the stack and begins the new dialog.

        Args:
            dialog_id: The ID of the dialog to switch to.
            options: Optional initial options for the new dialog.

        Returns:
            DialogTurnResult from the new dialog's begin step.
        """
        # Pop current
        if self.stack:
            ended_dialog = self.stack.pop()
            logger.info(f"Replacing dialog: {ended_dialog.id} with {dialog_id}")

        return await self.begin_dialog(dialog_id, options)

    async def end_dialog(self, result: Any = None) -> DialogTurnResult:
        """End the current dialog and resume the parent or complete.

        Pops the current dialog off the stack. If there are parent dialogs,
        resumes the parent. Otherwise, marks the entire session as complete.

        Args:
            result: Optional result to return from the dialog.

        Returns:
            DialogTurnResult with Complete status if all dialogs are done,
            otherwise continues the parent dialog.
        """
        # Pop current
        if self.stack:
            ended_dialog = self.stack.pop()
            logger.info(f"Ending dialog: {ended_dialog.id}")

        # If stack is now empty, we are done
        if not self.stack:
            return DialogTurnResult(DialogTurnStatus.Complete, result)

        # Resume parent
        parent_instance = self.active_dialog
        parent = self.find_dialog(parent_instance.id)
        if parent:
            return await parent.resume_dialog(self, "end_dialog", result)

        return DialogTurnResult(DialogTurnStatus.Complete, result)

    async def continue_dialog(self) -> DialogTurnResult:
        """Continue the active dialog with user's latest input.

        Returns:
            DialogTurnResult from continuing the dialog.
        """
        if not self.active_dialog:
            return DialogTurnResult(DialogTurnStatus.Empty)

        dialog_cls = self.find_dialog(self.active_dialog.id)
        if dialog_cls:
            # Instantiate the dialog class
            dialog = dialog_cls(ctx=self, options={})
            return await dialog.continue_dialog(self)

        return DialogTurnResult(DialogTurnStatus.Empty)

    async def finish_chat(self, message: str = None) -> DialogTurnResult:
        """Completely end the chat session, clearing all dialog state.

        Args:
            message: Optional message to send before finishing.

        Returns:
            DialogTurnResult with Complete status.
        """
        self.stack.clear()
        self._finished = True
        logger.info(f"Chat finished for user: {self.context.from_id}")

        if message:
            content = TextMessage(self.bot_ctx, message)
            await self.context.send_message(content)

        return DialogTurnResult(DialogTurnStatus.Complete)

    # - Dialog Step Navigation -
    async def next_step(self, result: Any = None) -> DialogTurnResult:
        """Skip to the next step of the waterfall.

        Args:
            result: Optional result to pass to the next step.

        Returns:
            DialogTurnResult from the next step.
        """
        if self.active_dialog:
            self.active_dialog.state["step_index"] += 1
            dialog_cls = self.find_dialog(self.active_dialog.id)
            if dialog_cls:
                dialog = dialog_cls(ctx=self, options={})
                return await dialog.run_step(self, "next", result)
        return DialogTurnResult(DialogTurnStatus.Cancelled)

    async def previous_step(self) -> DialogTurnResult:
        """Go back to the previous step of the waterfall.

        Returns:
            DialogTurnResult from the previous step.
        """
        if self.active_dialog:
            self.active_dialog.state["step_index"] -= 1
            dialog_cls = self.find_dialog(self.active_dialog.id)
            if dialog_cls:
                dialog = dialog_cls(ctx=self, options={})
                return await dialog.run_step(self, "previous", None)
        return DialogTurnResult(DialogTurnStatus.Cancelled)

    async def repeat_step(self, prompt: str = None) -> DialogTurnResult:
        """Repeat the current step of the waterfall.

        Args:
            prompt: Optional prompt to display before repeating.

        Returns:
            DialogTurnResult from the repeated step.
        """
        if prompt:
            await self.ctx.client.message.send_message(TextMessage(self.ctx, prompt))

        if self.active_dialog:
            dialog_cls = self.find_dialog(self.active_dialog.id)
            if dialog_cls:
                dialog = dialog_cls(ctx=self, options={})
                return await dialog.run_step(self, "repeat", None)
        return DialogTurnResult(DialogTurnStatus.Cancelled)

    async def await_step(self, prompt: str = None) -> DialogTurnResult:
        """Wait for user input on the current step without advancing.

        Args:
            prompt: Optional prompt to display while waiting.
        """
        if prompt:
            await self.ctx.client.message.send_message(TextMessage(self.ctx, prompt))

        if self.active_dialog:
            self.active_dialog.state["step_index"] -= 1
        # TODO: work on non-dialog environments

        return DialogTurnResult(DialogTurnStatus.Waiting)

    async def skip_to(self, target_step, result: Any = None) -> DialogTurnResult:
        """Skip to a specific step of the waterfall.

        Args:
            target_step: Step index (int), function reference, or step name (str).
            result: Optional result to pass to the target step.

        Returns:
            DialogTurnResult from the target step.
        """
        if not self.active_dialog:
            return DialogTurnResult(DialogTurnStatus.Cancelled)

        dialog_cls = self.find_dialog(self.active_dialog.id)
        if not dialog_cls:
            return DialogTurnResult(DialogTurnStatus.Cancelled)

        # instantiate to access steps
        dialog = dialog_cls(ctx=self, options={})

        # find target step index
        if isinstance(target_step, int):
            target_index = target_step
        elif callable(target_step):
            target_name = target_step.__name__
            target_index = None
            for idx, step in enumerate(dialog.steps):
                if step.__name__ == target_name:
                    target_index = idx
                    break
            if target_index is None:
                logger.warning(f"skip_to: target step '{target_name}' not found in dialog steps")
                return DialogTurnResult(DialogTurnStatus.Invalid)
        elif isinstance(target_step, str):  # allow passing step name as string
            target_index = None
            for idx, step in enumerate(dialog.steps):
                if step.__name__ == target_step:
                    target_index = idx
                    break
            if target_index is None:
                logger.warning(f"skip_to: target step '{target_step}' not found in dialog steps")
                return DialogTurnResult(DialogTurnStatus.Invalid)
        else:
            return DialogTurnResult(DialogTurnStatus.Invalid)

        # validate index bounds
        if target_index < 0 or target_index >= len(dialog.steps):
            return DialogTurnResult(DialogTurnStatus.Invalid)

        # set step index and run step
        self.active_dialog.state["step_index"] = target_index
        return await dialog.run_step(self, "skip_to", result)

    async def create_active_dialog(self, dialog_id: str, state: Dict[str, Any]) -> DialogInstance:
        """
        Creates new active dialog instance and pushes it to the stack.

        Args:
            dialog_id: The dialog ID.
            state: Initial state for the dialog.

        Returns:
            The created DialogInstance.
        """
        instance = DialogInstance(id=dialog_id, state=state)
        self.stack.append(instance)
        return instance

    # --- Message Wrappers ---
    class MessageHandler:
        """
        Class to handle message parsing and sending within dialog context.
        Connects directly to the binding client.
        """

        def __init__(self, ctx: "ChatContext"):
            """
            Initialize the message handler.

            Args:
                ctx: The chat context.
            """
            self.client = ctx.client

        async def send(self, content: Union[TextMessage, InteractiveTextMessage]) -> Any:
            """
            Send a message or interactive content.

            Args:
                content: The message or interactive message to send.

            Returns:
                The result from the binding client's send_message method.
            """
            # send message directly - the binding's MessageHandler handles parsing
            return await self.client.message.send_message(content)
