import typing as t

from sarus_data_spec.manager.ops.base import ScalarImplementation
from sarus_data_spec.scalar import SampleType


class PretrainedModel(ScalarImplementation):
    async def value(self) -> t.Any:
        assert self.scalar.is_pretrained_model()
        foundation_model_name = (
            self.scalar.protobuf().spec.model.pretrained_model.foundation_model_name  # noqa: E501
        )
        checkpoint_path = (
            self.scalar.protobuf().spec.model.pretrained_model.checkpoint_path  # noqa: E501
        )
        sample_type = SampleType(
            self.scalar.protobuf().spec.model.pretrained_model.sample_type
        )
        return foundation_model_name, checkpoint_path, sample_type