import re
import typing as t

import attrs
import requests
from packaging.version import Version

from tungstenkit._internal.utils.version import check_if_two_versions_compatible
from tungstenkit.exceptions import NoCompatibleCUDAImage

from .base_image import BaseImage, BaseImageCollection

RE_CUDA_IMAGE = (
    r"([0-9]+(?:[.][0-9]+){0,2})-cudnn([0-9]+(?:[.][0-9]+){0,2})-devel-ubuntu([0-9]+[.][0-9]{2})"
)


@attrs.frozen(order=True)
class CUDAImage(BaseImage):
    cuda_ver: Version
    cudnn_ver: Version
    ubuntu_ver: Version

    @staticmethod
    def get_repository():
        return "nvcr.io/nvidia/cuda"

    def get_tag(self):
        tag = (
            f"{self.cuda_ver}-cudnn{self.cudnn_ver}-devel-"
            f"ubuntu{self.ubuntu_ver.major}.{self.ubuntu_ver.minor:02}"
        )
        return tag


@attrs.define
class CUDAImageCollection(BaseImageCollection):
    images: t.List[CUDAImage]

    @classmethod
    def from_file(cls):
        # TODO implement this
        return CUDAImageCollection.from_remote()

    @classmethod
    def from_remote(cls):
        # Fetch all tags from docker hub
        tags = fetch_cuda_image_tags()

        # Parse
        images = []
        cuda_image_pattern = re.compile(RE_CUDA_IMAGE)
        for tag in tags:
            matches = re.findall(cuda_image_pattern, tag)
            if matches:
                cuda_ver_str, cudnn_ver_str, ubuntu_ver_str = matches[0]
                cuda_ver = Version(cuda_ver_str)
                cudnn_ver = Version(cudnn_ver_str)
                ubuntu_ver = Version(ubuntu_ver_str)
                images.append(
                    CUDAImage(
                        cuda_ver=cuda_ver,
                        cudnn_ver=cudnn_ver,
                        ubuntu_ver=ubuntu_ver,
                    )
                )
        return CUDAImageCollection(images=images)

    def get_cuda_image_by_cuda_cudnn_ver(
        self, cuda_ver: Version, cudnn_ver: t.Optional[Version]
    ) -> CUDAImage:
        candidates = [
            img
            for img in self.images
            if check_if_two_versions_compatible(img.cuda_ver, v2=cuda_ver)
            and check_if_two_versions_compatible(img.cudnn_ver, v2=cudnn_ver)
        ]
        if len(candidates) == 0:
            cudnn_ver_str = str(cudnn_ver) if cudnn_ver else ""
            raise NoCompatibleCUDAImage(f"CUDA version: {cuda_ver}" + cudnn_ver_str)
        return max(candidates)

    def get_latest_image(self) -> CUDAImage:
        return max(self.images)


def fetch_cuda_image_tags():
    url = "https://catalog.ngc.nvidia.com/_next/data/dj6JMKsmIGaYZMqoDIyx5/orgs/nvidia/containers/cuda/tags.json?orgName=nvidia&resourceType=containers&slug=cuda&slug=tags"
    r = requests.get(url)
    r.raise_for_status()
    body = r.json()
    return body["pageProps"]["data"]["tags"]
