from enum import StrEnum, auto
from typing import TYPE_CHECKING

from sqlalchemy.orm import Mapped, mapped_column, relationship

from .base import Base, DefaultFieldsMixin
from .types import EnumText, StringUUID

if TYPE_CHECKING:
    from human_input import HumanInputForm


class ExecutionContentType(StrEnum):
    HUMAN_INPUT_RESULT = auto()


class ExecutionExtraContent(DefaultFieldsMixin, Base):
    """ExecutionExtraContent stores extra contents produced during workflow / chatflow execution."""

    # The `ExecutionExtraContent` uses single table inheritance to model different
    # kinds of contents produced during message generation.
    #
    # See: https://docs.sqlalchemy.org/en/20/orm/inheritance.html#single-table-inheritance

    __tablename__ = "execution_extra_contents"
    __mapper_args__ = {
        "polymorphic_abstract": True,
        "polymorphic_on": "type",
        "with_polymorphic": "*",
    }
    # type records the type of the content. It serves as the `discriminator` for the
    # single table inheritance.
    type: Mapped[ExecutionContentType] = mapped_column(
        EnumText(ExecutionContentType, length=30),
        nullable=False,
    )

    # `workflow_run_id` records the workflow execution which generates this content, correspond to
    # `WorkflowRun.id`.
    workflow_run_id: Mapped[str] = mapped_column(StringUUID, nullable=False, index=True)

    # `message_id` records the messages generated by the execution associated with this `ExecutionExtraContent`.
    # It references to `Message.id`.
    #
    # For workflow execution, this field is `None`.
    #
    # For chatflow execution, `message_id`` is not None, and the following condition holds:
    #
    # The message referenced by `message_id` has `message.workflow_run_id == execution_extra_content.workflow_run_id`
    #
    message_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, index=True)


class HumanInputContent(ExecutionExtraContent):
    """HumanInputContent is a concrete class that represents human input content.
    It should only be initialized with the `new` class method."""

    __mapper_args__ = {
        "polymorphic_identity": ExecutionContentType.HUMAN_INPUT_RESULT,
    }

    # A relation to HumanInputForm table.
    #
    # While the form_id column is nullable in database (due to the nature of single table inheritance),
    # the form_id field should not be null for a given `HumanInputContent` instance.
    form_id: Mapped[str] = mapped_column(StringUUID, nullable=True)

    @classmethod
    def new(cls, form_id: str, message_id: str | None) -> "HumanInputContent":
        return cls(form_id=form_id, message_id=message_id)

    form: Mapped["HumanInputForm"] = relationship(
        "HumanInputForm",
        foreign_keys=[form_id],
        uselist=False,
        lazy="raise",
        primaryjoin="foreign(HumanInputContent.form_id) == HumanInputForm.id",
    )