from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

from ai.backend.agent.kernel_registry.loader.container import ContainerBasedKernelRegistryLoader
from ai.backend.agent.kernel_registry.writer.container import ContainerBasedKernelRegistryWriter

if TYPE_CHECKING:
    from ai.backend.agent.agent import AbstractAgent


@dataclass
class ContainerBasedKernelRegistryCreatorArgs:
    scratch_root: Path
    agent: AbstractAgent[Any, Any]


class ContainerBasedLoaderWriterCreator:
    """
    Creates a loader and writer for container-based kernel registry.
    """

    def __init__(
        self,
        args: ContainerBasedKernelRegistryCreatorArgs,
    ) -> None:
        self._scratch_root = args.scratch_root
        self._agent = args.agent

    def create_loader(self) -> ContainerBasedKernelRegistryLoader:
        return ContainerBasedKernelRegistryLoader(
            scratch_root=self._scratch_root,
            agent=self._agent,
        )

    def create_writer(self) -> ContainerBasedKernelRegistryWriter:
        return ContainerBasedKernelRegistryWriter(
            scratch_root=self._scratch_root,
        )
