from pathlib import Path
from typing import Union

import pydantic
from aind_behavior_services.schema import BonsaiSgenSerializers, convert_pydantic_to_bonsai
from aind_behavior_services.session import Session

from aind_behavior_just_frames.message_protocol import MessageProtocol
from aind_behavior_just_frames.rig import AindJustFramesRig

SCHEMA_ROOT = Path("./src/DataSchemas/")
EXTENSIONS_ROOT = Path("./src/Extensions/")
NAMESPACE_PREFIX = "AindJustFramesSchemas"


def main():
    models = [
        AindJustFramesRig,
        Session,
    ]
    model = pydantic.RootModel[Union[tuple(models)]]

    convert_pydantic_to_bonsai(
        model,
        model_name="aind_behavior_just_frames",
        root_element="Root",
        cs_namespace=NAMESPACE_PREFIX,
        json_schema_output_dir=SCHEMA_ROOT,
        cs_output_dir=EXTENSIONS_ROOT,
        cs_serializer=[BonsaiSgenSerializers.JSON],
    )

    convert_pydantic_to_bonsai(
        MessageProtocol,
        model_name="aind_physiology_fip_message_protocol",
        root_element="Root",
        cs_namespace=NAMESPACE_PREFIX + ".MessageProtocol",
        json_schema_output_dir=SCHEMA_ROOT,
        cs_output_dir=EXTENSIONS_ROOT,
        cs_serializer=[BonsaiSgenSerializers.JSON],
    )


if __name__ == "__main__":
    main()
