from typing import TYPE_CHECKING, Any, BinaryIO, Dict, List

from .container import Container
from .image import Image, ImageReference
from .method import Method
from .service import Service


if TYPE_CHECKING:
    import tomli as tomllib
else:
    try:
        import tomllib
    except ImportError:
        import tomli as tomllib


class State:
    __slots__ = ('images', 'containers', 'services', )

    images: List[Image]
    containers: List[Container]
    services: List[Service]

    def __init__(
            self,
            images: List[Image],
            containers: List[Container],
            services: List[Service],
    ) -> None:
        self.images = images
        self.containers = containers
        self.services = services

    def __repr__(self) -> str:
        return f'State({repr(self.images)}, {repr(self.containers)}, {repr(self.services)})'


def image_from_toml(spec: Dict[str, Any]) -> Image:
    """
    Loads an Image spec from toml
    """
    return Image(str(spec['path']), str(spec['hash']))


def container_from_toml(spec: Dict[str, Any]) -> Container:
    """
    Loads a Container spec from toml
    """
    return Container(ImageReference(str(spec['image'])), str(spec['runtime']))


def service_from_toml(spec: Dict[str, Any]) -> Service:
    """
    Loads a Service spec from toml
    """
    methods: List[Method] = []

    return Service(str(spec['name']), methods)


def from_toml(toml: BinaryIO) -> State:
    """
    Loads the given toml and builds a state instance out of it
    """
    toml_dict = tomllib.load(toml)

    images: List[Image] = []
    images_by_name: Dict[str, Image] = {}
    containers: List[Container] = []
    services: List[Service] = []

    for name, spec in toml_dict.items():
        if spec['apiVersion'] != 'v0':
            raise NotImplementedError('apiVersion', spec['apiVersion'])

        if spec['kind'] == 'Image':
            image = image_from_toml(spec)
            images.append(image)

            assert name not in images_by_name, f'Duplicate image name: {name}'
            images_by_name[name] = image
            continue

        if spec['kind'] == 'Container':
            containers.append(container_from_toml(spec))
            continue

        if spec['kind'] == 'Service':
            services.append(service_from_toml(spec))
            continue

        raise NotImplementedError(spec)

    for container in containers:
        if not isinstance(container.image, ImageReference):
            continue

        image_opt = images_by_name.get(container.image.name, None)
        assert image_opt is not None, f'Image reference not resolved: {container.image.name}'
        container.image = image_opt

    return State(images, containers, services)