from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, List, Optional

from rasa.dialogue_understanding.stack.frames.dialogue_stack_frame import (
    DialogueStackFrame,
)
from rasa.shared.core.flows import Flow, FlowsList, FlowStep
from rasa.shared.core.flows.steps.constants import START_STEP
from rasa.shared.exceptions import RasaException


class InvalidFlowStackFrameType(RasaException):
    """Raised if the stack frame type is invalid."""

    def __init__(self, frame_type: Optional[str]) -> None:
        """Creates a `InvalidFlowStackFrameType`.

        Args:
            frame_type: The invalid stack frame type.
        """
        super().__init__(f"Invalid stack frame type '{frame_type}'.")


class FlowStackFrameType(str, Enum):
    INTERRUPT = "interrupt"
    """The frame is an interrupt frame.

    This means that the previous flow was interrupted by this flow. An
    interrupt should be used for frames that span multiple turns and
    where we expect the user needing help to get back to the previous
    flow."""
    LINK = "link"
    """The frame is a link frame.

    This means that the previous flow linked to this flow."""

    CALL = "call"
    """The frame is a flow that is called from another flow."""

    REGULAR = "regular"
    """The frame is a regular frame.

    In all other cases, this is the case."""

    @staticmethod
    def from_str(typ: Optional[str]) -> "FlowStackFrameType":
        """Creates a `FlowStackFrameType` from a string.

        Args:
            typ: The string to create the `FlowStackFrameType` from.

        Returns:
        The created `FlowStackFrameType`.
        """
        if typ is None:
            return FlowStackFrameType.REGULAR
        elif typ == FlowStackFrameType.INTERRUPT.value:
            return FlowStackFrameType.INTERRUPT
        elif typ == FlowStackFrameType.CALL.value:
            return FlowStackFrameType.CALL
        elif typ == FlowStackFrameType.LINK.value:
            return FlowStackFrameType.LINK
        elif typ == FlowStackFrameType.REGULAR.value:
            return FlowStackFrameType.REGULAR
        else:
            raise InvalidFlowStackFrameType(typ)


class InvalidFlowIdException(Exception):
    """Raised if the flow ID is invalid."""

    def __init__(self, flow_id: str) -> None:
        """Creates a `InvalidFlowIdException`.

        Args:
            flow_id: The invalid flow ID.
        """
        super().__init__(f"Invalid flow ID '{flow_id}'.")


class InvalidFlowStepIdException(Exception):
    """Raised if the flow step ID is invalid."""

    def __init__(self, flow_id: str, step_id: str) -> None:
        """Creates a `InvalidFlowStepIdException`.

        Args:
            flow_id: The invalid flow ID.
            step_id: The invalid flow step ID.
        """
        super().__init__(f"Invalid flow step ID '{step_id}' for flow '{flow_id}'.")


class InvalidAgentState(RasaException):
    """Raised if the agent state is invalid."""

    def __init__(self, invalid_state: str) -> None:
        """Creates a `InvalidAgentState`.

        Args:
            invalid_state: The invalid agent state.
        """
        super().__init__(f"Invalid agent state '{invalid_state}'.")


@dataclass
class BaseFlowStackFrame(DialogueStackFrame):
    flow_id: str = ""  # needed to avoid "default arg before non-default" error
    """The ID of the current flow."""
    step_id: str = START_STEP
    """The ID of the current step."""

    def flow(self, all_flows: FlowsList) -> Flow:
        """Returns the current flow.

        Args:
            all_flows: All flows in the assistant.

        Returns:
        The current flow.
        """
        flow = all_flows.flow_by_id(self.flow_id)
        if not flow:
            # we shouldn't ever end up with a frame that belongs to a non
            # existing flow, but if we do, we should raise an error
            raise InvalidFlowIdException(self.flow_id)
        return flow

    def step(self, all_flows: FlowsList) -> FlowStep:
        """Returns the current flow step.

        Args:
            all_flows: All flows in the assistant.

        Returns:
            The current flow step.
        """
        flow = self.flow(all_flows)

        step_id = self.step_id
        # in 3.11.4 we added the flow_id as a prefix to the step_id
        # this causes issues when loading old dialogues as the prefix is missing
        # (see https://rasahq.atlassian.net/jira/software/c/projects/ENG/boards/43?selectedIssue=ENG-1939)
        # so we try to find the step by adding the flow prefix to old step_ids as well
        # TODO: remove this in 4.0.0
        alternative_step_id = f"{self.flow_id}_{self.step_id}"

        step = flow.step_by_id(step_id) or flow.step_by_id(alternative_step_id)

        if not step:
            # we shouldn't ever end up with a frame that belongs to a non
            # existing step, but if we do, we should raise an error
            raise InvalidFlowStepIdException(self.flow_id, self.step_id)
        return step


@dataclass
class UserFlowStackFrame(BaseFlowStackFrame):
    frame_type: FlowStackFrameType = FlowStackFrameType.REGULAR
    """The type of the frame. Defaults to `StackFrameType.REGULAR`."""
    suspended_agent_frames: Optional[List[Dict[str, Any]]] = None
    """Agent stack frames (as dicts) to restore when reaching that agent's step.

    Used when rewinding the stack (e.g. slot correction or agent restart) so
    interrupted agents can be resumed. Old trackers may have no key; treat as None.
    """

    @classmethod
    def type(cls) -> str:
        """Returns the type of the frame."""
        return "flow"

    def as_dict(self) -> Dict[str, Any]:
        """Serialize to dict; omit suspended_agent_frames when None to avoid bloat."""
        data = super().as_dict()
        if data.get("suspended_agent_frames") is None:
            data.pop("suspended_agent_frames", None)
        return data

    @staticmethod
    def from_dict(data: Dict[str, Any]) -> UserFlowStackFrame:
        """Creates a `DialogueStackFrame` from a dictionary.

        Args:
            data: The dictionary to create the `DialogueStackFrame` from.

        Returns:
            The created `DialogueStackFrame`.
        """
        suspended = data.get("suspended_agent_frames")
        if suspended is not None and not isinstance(suspended, list):
            suspended = None
        return UserFlowStackFrame(
            frame_id=data["frame_id"],
            flow_id=data["flow_id"],
            step_id=data["step_id"],
            frame_type=FlowStackFrameType.from_str(data.get("frame_type")),
            suspended_agent_frames=suspended,
        )


class AgentState(str, Enum):
    INTERRUPTED = "interrupted"
    WAITING_FOR_INPUT = "waiting_for_input"

    @staticmethod
    def from_str(state: Optional[str]) -> AgentState:
        if state == AgentState.WAITING_FOR_INPUT.value:
            return AgentState.WAITING_FOR_INPUT
        elif state == AgentState.INTERRUPTED.value:
            return AgentState.INTERRUPTED
        else:
            raise InvalidAgentState(state)


@dataclass
class AgentStackFrame(BaseFlowStackFrame):
    agent_id: str = ""
    state: AgentState = AgentState.WAITING_FOR_INPUT
    metadata: Optional[Dict[str, Any]] = None

    @classmethod
    def type(cls) -> str:
        """Returns the type of the frame."""
        return "agent"

    @staticmethod
    def from_dict(data: Dict[str, Any]) -> AgentStackFrame:
        """Creates a `AgentStackFrame` from a dictionary.

        Args:
            data: The dictionary to create the `AgentStackFrame` from.

        Returns:
            The created `AgentStackFrame`.
        """
        return AgentStackFrame(
            frame_id=data["frame_id"],
            flow_id=data["flow_id"],
            step_id=data["step_id"],
            agent_id=data["agent_id"],
            state=AgentState.from_str(data["state"]),
            metadata=data.get("metadata"),
        )
