from typing import override

from ai.backend.common.config import ModelDefinition
from ai.backend.manager.data.deployment.types import (
    ModelRevisionSpec,
)
from ai.backend.manager.repositories.deployment import DeploymentRepository
from ai.backend.manager.sokovan.deployment.definition_generator.base import ModelDefinitionGenerator


class CustomModelDefinitionGenerator(ModelDefinitionGenerator):
    """Model definition generator implementation for custom runtime variant."""

    _deployment_repository: DeploymentRepository

    def __init__(self, deployment_repository: DeploymentRepository) -> None:
        self._deployment_repository = deployment_repository

    @override
    async def generate_model_definition(self, model_revision: ModelRevisionSpec) -> ModelDefinition:
        model_definition_content = await self._deployment_repository.fetch_model_definition(
            vfolder_id=model_revision.mounts.model_vfolder_id,
            model_definition_path=model_revision.mounts.model_definition_path,
        )
        return ModelDefinition.model_validate(model_definition_content)
