"""
Types and classes for dialog tapes and annotators.
"""

from __future__ import annotations

from typing import Any, Literal, TypeAlias

from pydantic import BaseModel

from tapeagents.agent import Annotator
from tapeagents.core import (
    Action,
    AgentEvent,
    Call,
    FinalStep,
    Observation,
    Pass,
    Respond,
    SetNextNode,
    StopStep,
    Tape,
    Thought,
)
from tapeagents.tool_calling import ToolCalls, ToolResult, ToolSpec


class SystemStep(Observation):
    """
    Step rendered into system message of the prompt.

    Attributes:
        content (str): The content of the system step.
        kind (Literal["system"]): A literal indicating the type of step, which is always "system".
    """

    content: str
    kind: Literal["system"] = "system"


class UserStep(Observation):
    """
    Represents a step taken by a user in a dialog.

    Attributes:
        content (str): The content of the user's step.
        kind (Literal["user"]): The type of step, which is always "user".
    """

    content: str
    kind: Literal["user"] = "user"


class AssistantThought(Thought):
    """
    Represents a thought generated by an assistant.

    Attributes:
        content (Any): The content of the assistant's thought.
        kind (Literal["assistant_thought"]): A literal string indicating the type of thought.
    """

    content: Any
    kind: Literal["assistant_thought"] = "assistant_thought"


class AssistantStep(Action):
    """
    Represents a step taken by an assistant in a dialog.

    Attributes:
        content (str): The content of the assistant's response.
        kind (Literal["assistant"]): The type of step, which is always "assistant".
    """

    content: str
    kind: Literal["assistant"] = "assistant"


class AssistantAnswer(StopStep):
    """
    Represents an answer generated by an assistant.

    Attributes:
        kind (Literal["assistant_answer"]): The type of step, which is always "assistant_answer".
        success (bool): A boolean indicating whether the assistant's response was successful.
        answer (str): The assistant's answer.
    """

    kind: Literal["assistant_answer"] = "assistant_answer"
    success: bool
    answer: str


DialogStep: TypeAlias = (
    # observations
    UserStep
    | ToolResult
    | SystemStep
    # thoughts
    | AssistantThought
    | SetNextNode
    | Pass
    | Call
    | Respond
    # actions
    | FinalStep
    | AssistantStep
    | ToolCalls
)
"""Type alias for dialog steps."""


class DialogContext(BaseModel):
    """
    Context for dialog agents, containing tools and other information.
    """

    # TODO: define type signature for tools including JSONSchema and etc
    tools: list[ToolSpec]


# Tape for dialog agents, containing dialog steps and optional context
DialogTape = Tape[DialogContext | None, DialogStep]
"""Type alias for dialog tapes."""


DialogEvent: TypeAlias = AgentEvent[DialogTape]
"""Type alias for dialog events."""


class AnnotatorFreeFormThought(Thought):
    """
    AnnotatorFreeFormThought is a subclass of Thought that represents a free-form thought provided by an annotator.

    Attributes:
        kind (Literal["annotator_free_form_thought"]): A constant string that identifies the type of thought.
        content (str): The content of the free-form thought provided by the annotator.
    """

    kind: Literal["annotator_free_form_thought"] = "annotator_free_form_thought"
    content: str


class AnnotationAction(Action):
    """
    AnnotationAction is a subclass of Action that represents an action produced by an annotator.

    Attributes:
        kind (Literal["annotation_action"]): A string literal indicating the type of action.
        annotation (dict): A dictionary containing annotation data.
    """

    kind: Literal["annotation_action"] = "annotation_action"
    annotation: dict


DialogAnnotatorTape: TypeAlias = Tape[DialogTape, AnnotatorFreeFormThought | AnnotationAction]
"""Type alias for dialog annotator tapes."""


class DialogAnnotator(Annotator[DialogTape, DialogAnnotatorTape]):
    """
    DialogAnnotator is a class that extends the Annotator agent with specific types DialogTape and DialogAnnotatorTape.

    Methods:
        make_own_tape(tape: DialogTape) -> DialogAnnotatorTape:
            Creates and returns a DialogAnnotatorTape instance using the provided DialogTape instance.
    """

    def make_own_tape(self, tape: DialogTape) -> DialogAnnotatorTape:
        """
        Creates a DialogAnnotatorTape instance using given DialogTape as context.

        Args:
            tape (DialogTape): The DialogTape instance to be converted.

        Returns:
            DialogAnnotatorTape: A new instance of DialogAnnotatorTape with the provided context.
        """
        return DialogAnnotatorTape(context=tape)
