from typing import TypeVar
from redis_om import JsonModel
from .annotators import Annotator
from .env_agent_combo_storage import EnvAgentComboStorage
from .logs import AnnotationForEpisode, EpisodeLog
from .persistent_profile import (
    AgentProfile,
    EnvironmentProfile,
    RelationshipProfile,
    RelationshipType,
)
from .serialization import (
    agentprofiles_to_csv,
    agentprofiles_to_jsonl,
    envagnetcombostorage_to_csv,
    envagnetcombostorage_to_jsonl,
    environmentprofiles_to_csv,
    environmentprofiles_to_jsonl,
    episodes_to_csv,
    episodes_to_jsonl,
    get_rewards_from_episode,
    jsonl_to_agentprofiles,
    jsonl_to_envagnetcombostorage,
    jsonl_to_environmentprofiles,
    jsonl_to_episodes,
    jsonl_to_relationshipprofiles,
    relationshipprofiles_to_csv,
    relationshipprofiles_to_jsonl,
)
from .session_transaction import MessageTransaction, SessionTransaction
from .waiting_room import MatchingInWaitingRoom
from .aggregate_annotations import map_human_annotations_to_episode_logs

__all__ = [
    "AgentProfile",
    "EnvironmentProfile",
    "EpisodeLog",
    "EnvAgentComboStorage",
    "AnnotationForEpisode",
    "Annotator",
    "RelationshipProfile",
    "RelationshipType",
    "RedisCommunicationMixin",
    "SessionTransaction",
    "MessageTransaction",
    "MatchingInWaitingRoom",
    "agentprofiles_to_csv",
    "agentprofiles_to_jsonl",
    "environmentprofiles_to_csv",
    "environmentprofiles_to_jsonl",
    "relationshipprofiles_to_csv",
    "relationshipprofiles_to_jsonl",
    "envagnetcombostorage_to_csv",
    "envagnetcombostorage_to_jsonl",
    "episodes_to_csv",
    "episodes_to_jsonl",
    "map_human_annotations_to_episode_logs",
    "jsonl_to_agentprofiles",
    "jsonl_to_environmentprofiles",
    "jsonl_to_episodes",
    "jsonl_to_relationshipprofiles",
    "jsonl_to_envagnetcombostorage",
    "get_rewards_from_episode",
]

InheritedJsonModel = TypeVar("InheritedJsonModel", bound="JsonModel")


def _json_model_all(cls: type[InheritedJsonModel]) -> list[InheritedJsonModel]:
    return cls.find().all()  # type: ignore[return-value]


JsonModel.all = classmethod(_json_model_all)  # type: ignore[assignment,method-assign]
