from typing import TYPE_CHECKING, Any, Sequence

from moxn.models import message as msg
from moxn.models import schema as sch
from moxn_types.base import BasePrompt, RenderableModel
from moxn_types.content import MessageRole

if TYPE_CHECKING:
    PromptSession = Any


class PromptTemplate(BasePrompt[msg.Message, sch.Schema]):
    """Immutable representation of a stored prompt configuration."""

    messages: Sequence[msg.Message]
    input_schema: sch.Schema

    def get_messages(self):
        return [message.model_copy(deep=True) for message in self.messages]

    def get_message_by_role(self, role: str | MessageRole) -> msg.Message | None:
        """Get the first message with the specified role."""
        _role = MessageRole(role) if isinstance(role, str) else role

        messages = [p for p in self.messages if p.role == role]
        if len(messages) == 1:
            return messages[0]
        elif len(messages) == 0:
            return None
        else:
            raise ValueError(
                f"get message is not deterministic, there are {len(messages)} {_role.value} messages in the prompt"
            )

    def get_messages_by_role(self, role: str | MessageRole) -> list[msg.Message]:
        """Get all messages with the specified role."""
        role = MessageRole(role) if isinstance(role, str) else role
        return [p for p in self.messages if p.role == role]

    def to_prompt_session(
        self,
        session_data: RenderableModel | None = None,
        render_kwargs: dict[str, Any] | None = None,
    ) -> "PromptSession":  # type: ignore
        from moxn.models.prompt.session import PromptSession

        return PromptSession.from_prompt_template(
            prompt=self,
            session_data=session_data,
            render_kwargs=render_kwargs,
        )
