from __future__ import annotations

import asyncio
import copy
import logging
from abc import ABCMeta, abstractmethod
from collections.abc import AsyncIterator, Mapping, Sequence
from contextlib import asynccontextmanager as actxmgr
from contextvars import ContextVar
from typing import (
    TYPE_CHECKING,
    Any,
    Final,
    cast,
)

import aiohttp
import aiotools
import sqlalchemy as sa
import trafaret as t
import yarl

from ai.backend.common.bgtask.reporter import ProgressReporter
from ai.backend.common.docker import (
    ImageRef,
    arch_name_aliases,
    validate_image_labels,
)
from ai.backend.common.docker import login as registry_login
from ai.backend.common.exception import (
    InvalidImageName,
    InvalidImageTag,
    ProjectMismatchWithCanonical,
)
from ai.backend.common.json import read_json
from ai.backend.common.types import SlotName, SSLContextType
from ai.backend.common.utils import join_non_empty
from ai.backend.logging import BraceStyleAdapter
from ai.backend.manager.data.image.types import (
    ImageData,
    ImageStatus,
    ImageType,
    RescanImagesResult,
)
from ai.backend.manager.defs import INTRINSIC_SLOTS_MIN
from ai.backend.manager.exceptions import ScanImageError, ScanTagError
from ai.backend.manager.models.image import ImageIdentifier, ImageRow
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine

log = BraceStyleAdapter(logging.getLogger(__spec__.name))
concurrency_sema: ContextVar[asyncio.Semaphore] = ContextVar("concurrency_sema")
progress_reporter: ContextVar[ProgressReporter | None] = ContextVar(
    "progress_reporter", default=None
)
all_updates: ContextVar[dict[ImageIdentifier, dict[str, Any]]] = ContextVar("all_updates")

if TYPE_CHECKING:
    from ai.backend.manager.models.container_registry import ContainerRegistryRow


class BaseContainerRegistry(metaclass=ABCMeta):
    db: ExtendedAsyncSAEngine
    registry_name: str
    registry_info: ContainerRegistryRow
    registry_url: yarl.URL
    max_concurrency_per_registry: int
    base_hdrs: dict[str, str]
    credentials: dict[str, str]
    ssl_verify: bool

    MEDIA_TYPE_OCI_INDEX: Final[str] = "application/vnd.oci.image.index.v1+json"
    MEDIA_TYPE_OCI_MANIFEST: Final[str] = "application/vnd.oci.image.manifest.v1+json"
    MEDIA_TYPE_DOCKER_MANIFEST_LIST: Final[str] = (
        "application/vnd.docker.distribution.manifest.list.v2+json"
    )
    MEDIA_TYPE_DOCKER_MANIFEST: Final[str] = "application/vnd.docker.distribution.manifest.v2+json"

    # Legacy manifest types (deprecated)
    MEDIA_TYPE_DOCKER_MANIFEST_V1_PRETTY_JWS: Final[str] = (
        "application/vnd.docker.distribution.manifest.v1+prettyjws"
    )
    MEDIA_TYPE_DOCKER_MANIFEST_V1_JSON: Final[str] = (
        "application/vnd.docker.distribution.manifest.v1+json"
    )

    def __init__(
        self,
        db: ExtendedAsyncSAEngine,
        registry_name: str,
        registry_info: ContainerRegistryRow,
        *,
        max_concurrency_per_registry: int = 4,
        ssl_verify: bool = True,
    ) -> None:
        self.db = db
        self.registry_name = registry_name
        self.registry_info = registry_info
        self.registry_url = yarl.URL(registry_info.url)
        self.max_concurrency_per_registry = max_concurrency_per_registry
        self.base_hdrs = {
            "Accept": self.MEDIA_TYPE_DOCKER_MANIFEST,
        }
        self.credentials = {}
        self.ssl_verify = ssl_verify

    @actxmgr
    async def prepare_client_session(self) -> AsyncIterator[tuple[yarl.URL, aiohttp.ClientSession]]:
        ssl_ctx: SSLContextType = True  # default
        if not self.registry_info.ssl_verify:
            ssl_ctx = False
        connector = aiohttp.TCPConnector(ssl=ssl_ctx)
        async with aiohttp.ClientSession(connector=connector) as sess:
            yield self.registry_url, sess

    async def rescan_single_registry(
        self,
        reporter: ProgressReporter | None = None,
    ) -> RescanImagesResult:
        log.info("rescan_single_registry()")
        errors: list[str] = []

        all_updates_token = all_updates.set({})
        concurrency_sema.set(asyncio.Semaphore(self.max_concurrency_per_registry))
        progress_reporter.set(reporter)
        try:
            username = self.registry_info.username
            if username is not None:
                self.credentials["username"] = username
            password = self.registry_info.password
            if password is not None:
                self.credentials["password"] = password
            async with self.prepare_client_session() as (url, client_session):
                self.registry_url = url

                tasks = []
                async for image in self.fetch_repositories(client_session):
                    task = asyncio.create_task(self._scan_image(client_session, image))
                    tasks.append(task)

                for fut in asyncio.as_completed(tasks):
                    try:
                        await fut
                    except Exception as e:
                        errors.append(f"Failed to scan image! Detail: {e!s}")

            scanned_images = await self.commit_rescan_result()
            return RescanImagesResult(images=scanned_images, errors=errors)
        finally:
            all_updates.reset(all_updates_token)

    async def commit_rescan_result(self) -> list[ImageData]:
        scanned_images: list[ImageData] = []
        _all_updates = all_updates.get()
        if not _all_updates:
            log.info("No images found in registry {0}", self.registry_url)
        else:
            image_identifiers = [(k.canonical, k.architecture) for k in _all_updates.keys()]
            async with self.db.begin_session() as session:
                existing_images = await session.scalars(
                    sa.select(ImageRow).where(
                        sa.func.ROW(ImageRow.name, ImageRow.architecture).in_(image_identifiers),
                    )
                )
                is_local = self.registry_name == "local"

                for image_row in existing_images:
                    image_ref = image_row.image_ref
                    update_key = ImageIdentifier(image_ref.canonical, image_ref.architecture)
                    if update := _all_updates.pop(update_key, None):
                        image_row.config_digest = update["config_digest"]
                        image_row.size_bytes = update["size_bytes"]
                        image_row.accelerators = update.get("accels")
                        image_row.labels = update["labels"]
                        image_row.is_local = is_local
                        scanned_images.append(image_row.to_dataclass())

                        if image_row.status == ImageStatus.DELETED:
                            image_row.status = ImageStatus.ALIVE

                            progress_msg = f"Restored deleted image - {image_ref.canonical}/{image_ref.architecture} ({update['config_digest']})"
                            log.info(progress_msg)

                            if (reporter := progress_reporter.get()) is not None:
                                await reporter.update(1, message=progress_msg)

                for image_identifier, update in _all_updates.items():
                    try:
                        parsed_img = ImageRef.from_image_str(
                            image_identifier.canonical,
                            self.registry_info.project,
                            self.registry_info.registry_name,
                            is_local=is_local,
                        )
                    except (ProjectMismatchWithCanonical, ValueError) as e:
                        skip_reason = str(e)
                        progress_msg = f"Skipped image - {image_identifier.canonical}/{image_identifier.architecture} ({skip_reason})"
                        log.warning(progress_msg)
                        if (reporter := progress_reporter.get()) is not None:
                            await reporter.update(1, message=progress_msg)
                        continue

                    image_row = ImageRow(
                        name=parsed_img.canonical,
                        project=self.registry_info.project,
                        registry=parsed_img.registry,
                        registry_id=self.registry_info.id,
                        image=join_non_empty(parsed_img.project, parsed_img.name, sep="/"),
                        tag=parsed_img.tag,
                        architecture=image_identifier.architecture,
                        is_local=is_local,
                        config_digest=update["config_digest"],
                        size_bytes=update["size_bytes"],
                        type=ImageType.COMPUTE,
                        accelerators=update.get("accels"),
                        labels=update["labels"],
                        status=ImageStatus.ALIVE,
                    )
                    session.add(image_row)
                    scanned_images.append(image_row.to_dataclass())
                    progress_msg = f"Updated image - {parsed_img.canonical}/{image_identifier.architecture} ({update['config_digest']})"
                    log.info(progress_msg)

                    if (reporter := progress_reporter.get()) is not None:
                        await reporter.update(1, message=progress_msg)

                await session.flush()
        return scanned_images

    async def scan_single_ref(self, image: str) -> RescanImagesResult:
        all_updates_token = all_updates.set({})
        sema_token = concurrency_sema.set(asyncio.Semaphore(1))
        try:
            username = self.registry_info.username
            if username is not None:
                self.credentials["username"] = username
            password = self.registry_info.password
            if password is not None:
                self.credentials["password"] = password
            async with self.prepare_client_session() as (url, sess):
                project_and_image_name, tag = ImageRef.parse_image_tag(image)
                rqst_args = await registry_login(
                    sess,
                    self.registry_url,
                    self.credentials,
                    f"repository:{project_and_image_name}:pull",
                )
                await self._scan_tag(sess, rqst_args, project_and_image_name, tag)
            scanned_images = await self.commit_rescan_result()
            return RescanImagesResult(images=scanned_images)
        finally:
            concurrency_sema.reset(sema_token)
            all_updates.reset(all_updates_token)

    async def _scan_image(
        self,
        sess: aiohttp.ClientSession,
        image: str,
    ) -> None:
        log.info("_scan_image()")
        rqst_args = await registry_login(
            sess,
            self.registry_url,
            self.credentials,
            f"repository:{image}:pull",
        )
        rqst_args["headers"].update(**self.base_hdrs)
        tags = []
        tag_list_url: yarl.URL | None
        tag_list_url = (self.registry_url / f"v2/{image}/tags/list").with_query(
            {"n": "10"},
        )
        while tag_list_url is not None:
            async with sess.get(tag_list_url, **rqst_args) as resp:
                data = await read_json(resp)
                tags_data = data.get("tags", [])
                # sometimes there are dangling image names in the hub.
                if not tags_data:
                    break

                tags.extend(tags_data)
                tag_list_url = None
                next_page_link = resp.links.get("next")
                if next_page_link:
                    next_page_url = cast(yarl.URL, next_page_link["url"])
                    tag_list_url = self.registry_url.with_path(next_page_url.path).with_query(
                        next_page_url.query
                    )
        if (reporter := progress_reporter.get()) is not None:
            reporter.total_progress += len(tags)

        try:
            async with aiotools.TaskGroup() as tg:
                for tag in tags:
                    tg.create_task(self._scan_tag(sess, rqst_args, image, tag))
        except aiotools.TaskGroupError as e:
            raise ScanImageError(
                f"Image scan failed, Details: {cast(ExceptionGroup, e).exceptions}"
            ) from e

    async def _scan_tag(
        self,
        sess: aiohttp.ClientSession,
        rqst_args: dict[str, Any],
        image: str,
        tag: str,
    ) -> None:
        async with concurrency_sema.get():
            rqst_args = copy.deepcopy(rqst_args)
            acceptables = [
                self.MEDIA_TYPE_DOCKER_MANIFEST_LIST,
                self.MEDIA_TYPE_DOCKER_MANIFEST,
                self.MEDIA_TYPE_OCI_INDEX,
                self.MEDIA_TYPE_OCI_MANIFEST,
                self.MEDIA_TYPE_DOCKER_MANIFEST_V1_JSON,
                self.MEDIA_TYPE_DOCKER_MANIFEST_V1_PRETTY_JWS,
            ]
            rqst_args["headers"]["Accept"] = ",".join(acceptables)
            async with sess.get(
                self.registry_url / f"v2/{image}/manifests/{tag}", **rqst_args
            ) as resp:
                if resp.status == 404:
                    # ignore missing tags
                    # (may occur after deleting an image from the docker hub)
                    return
                content_type = resp.headers["Content-Type"]
                resp.raise_for_status()
                resp_json = await read_json(resp)

                try:
                    async with aiotools.TaskGroup() as tg:
                        match content_type:
                            case self.MEDIA_TYPE_DOCKER_MANIFEST:
                                await self._process_docker_v2_image(
                                    tg, sess, rqst_args, image, tag, resp_json
                                )
                            case self.MEDIA_TYPE_DOCKER_MANIFEST_LIST:
                                await self._process_docker_v2_multiplatform_image(
                                    tg, sess, rqst_args, image, tag, resp_json
                                )
                            case self.MEDIA_TYPE_OCI_INDEX:
                                await self._process_oci_index(
                                    tg, sess, rqst_args, image, tag, resp_json
                                )
                            case self.MEDIA_TYPE_OCI_MANIFEST:
                                await self._process_oci_manifest(
                                    tg, sess, rqst_args, image, tag, resp_json
                                )
                            case (
                                self.MEDIA_TYPE_DOCKER_MANIFEST_V1_PRETTY_JWS
                                | self.MEDIA_TYPE_DOCKER_MANIFEST_V1_JSON
                            ):
                                await self._process_docker_v1_image(
                                    tg, sess, rqst_args, image, tag, resp_json
                                )

                            case _:
                                log.warning("Unknown content type: {}", content_type)
                                raise RuntimeError(
                                    "The registry does not support the standard way of "
                                    "listing multiarch images."
                                )
                except aiotools.TaskGroupError as e:
                    raise ScanTagError(
                        f"Tag scan failed, Details: {cast(ExceptionGroup, e).exceptions}"
                    ) from e

    async def _read_manifest_list(
        self,
        sess: aiohttp.ClientSession,
        manifest_list: Sequence[Any],
        rqst_args: dict[str, Any],
        image: str,
        tag: str,
    ) -> None:
        """
        Understands images defined under [OCI image manifest](https://github.com/opencontainers/image-spec/blob/main/manifest.md#example-image-manifest) or
        [Docker image manifest list](https://github.com/openshift/docker-distribution/blob/master/docs/spec/manifest-v2-2.md#example-manifest-list)
        and imports Backend.AI compatible images.
        """
        manifests = {}
        for manifest in manifest_list:
            platform_arg = f"{manifest['platform']['os']}/{manifest['platform']['architecture']}"
            if variant := manifest["platform"].get("variant", None):
                platform_arg += f"/{variant}"
            architecture = manifest["platform"]["architecture"]
            architecture = arch_name_aliases.get(architecture, architecture)

            async with sess.get(
                self.registry_url / f"v2/{image}/manifests/{manifest['digest']}",
                **rqst_args,
            ) as resp:
                manifest_info = await resp.json()

            manifests[architecture] = await self._preprocess_manifest(
                sess, manifest_info, rqst_args, image
            )

            if not manifests[architecture]["labels"]:
                log.warning(
                    "The image {}:{}/{} has no metadata labels -> treating as vanilla image",
                    image,
                    tag,
                    architecture,
                )
                manifests[architecture]["labels"] = {}

        await self._read_manifest(image, tag, manifests)

    async def _preprocess_manifest(
        self,
        sess: aiohttp.ClientSession,
        manifest: Mapping[str, Any],
        rqst_args: dict[str, Any],
        image: str,
    ) -> dict[str, Any]:
        """
        Extracts informations from
        [Docker iamge manifest](https://github.com/openshift/docker-distribution/blob/master/docs/spec/manifest-v2-2.md#example-image-manifest)
        required by Backend.AI.
        """
        config_digest = manifest["config"]["digest"]
        size_bytes = sum(layer["size"] for layer in manifest["layers"]) + manifest["config"]["size"]

        async with sess.get(
            self.registry_url / f"v2/{image}/blobs/{config_digest}", **rqst_args
        ) as resp:
            resp.raise_for_status()
            data = await read_json(resp)
        labels = {}

        # we should favor `config` instead of `container_config` since `config` can contain additional datas
        # set when commiting image via `--change` flag
        if _config_labels := (data.get("config") or {}).get("Labels"):
            labels = _config_labels
        elif _container_config_labels := (data.get("container_config") or {}).get("Labels"):
            labels = _container_config_labels

        return {
            "size": size_bytes,
            "labels": labels,
            "digest": config_digest,
        }

    async def _process_oci_index(
        self,
        _tg: aiotools.TaskGroup,
        sess: aiohttp.ClientSession,
        rqst_args: dict[str, Any],
        image: str,
        tag: str,
        image_info: Mapping[str, Any],
    ) -> None:
        manifest_list = [
            item
            for item in image_info["manifests"]
            if "annotations" not in item  # skip attestation manifests
        ]
        rqst_args = copy.deepcopy(rqst_args)
        rqst_args["headers"]["Accept"] = self.MEDIA_TYPE_OCI_MANIFEST

        await self._read_manifest_list(sess, manifest_list, rqst_args, image, tag)

    async def _process_oci_manifest(
        self,
        _tg: aiotools.TaskGroup,
        sess: aiohttp.ClientSession,
        rqst_args: dict[str, Any],
        image: str,
        tag: str,
        image_info: Mapping[str, Any],
    ) -> None:
        rqst_args = copy.deepcopy(rqst_args)
        rqst_args["headers"] = rqst_args.get("headers", {})
        rqst_args["headers"].update({
            "Accept": self.MEDIA_TYPE_OCI_MANIFEST,
        })

        if (reporter := progress_reporter.get()) is not None:
            reporter.total_progress += 1

        config_digest = image_info["config"]["digest"]
        size_bytes = (
            sum(layer["size"] for layer in image_info["layers"]) + image_info["config"]["size"]
        )

        async with sess.get(
            self.registry_url / f"v2/{image}/blobs/{config_digest}",
            **rqst_args,
        ) as resp:
            resp.raise_for_status()
            config_data = await read_json(resp)

        labels = {}
        if _config_labels := (config_data.get("config") or {}).get("Labels"):
            labels = _config_labels
        elif _container_config_labels := (config_data.get("container_config") or {}).get("Labels"):
            labels = _container_config_labels

        if not labels:
            log.warning(
                "The image {}:{} has no metadata labels -> treating as vanilla image",
                image,
                tag,
            )
            labels = {}

        architecture = config_data.get("architecture")
        if architecture:
            architecture = arch_name_aliases.get(architecture, architecture)
        else:
            if tag.endswith(("-arm64", "-aarch64")):
                architecture = "aarch64"
            else:
                architecture = "x86_64"

        manifests = {
            architecture: {
                "size": size_bytes,
                "labels": labels,
                "digest": config_digest,
            }
        }
        await self._read_manifest(image, tag, manifests)

    async def _process_docker_v2_multiplatform_image(
        self,
        _tg: aiotools.TaskGroup,
        sess: aiohttp.ClientSession,
        rqst_args: dict[str, Any],
        image: str,
        tag: str,
        image_info: Mapping[str, Any],
    ) -> None:
        manifest_list = image_info["manifests"]
        rqst_args = copy.deepcopy(rqst_args)
        rqst_args["headers"]["Accept"] = self.MEDIA_TYPE_DOCKER_MANIFEST

        await self._read_manifest_list(
            sess,
            manifest_list,
            rqst_args,
            image,
            tag,
        )

    async def _process_docker_v2_image(
        self,
        _tg: aiotools.TaskGroup,
        sess: aiohttp.ClientSession,
        rqst_args: dict[str, Any],
        image: str,
        tag: str,
        image_info: Mapping[str, Any],
    ) -> None:
        config_digest = image_info["config"]["digest"]
        rqst_args = copy.deepcopy(rqst_args)
        rqst_args["headers"]["Accept"] = self.MEDIA_TYPE_DOCKER_MANIFEST

        async with sess.get(
            self.registry_url / f"v2/{image}/blobs/{config_digest}",
            **rqst_args,
        ) as resp:
            resp.raise_for_status()
            blob_data = await read_json(resp)

        manifest_arch = blob_data["architecture"]
        architecture = arch_name_aliases.get(manifest_arch, manifest_arch)

        manifests = {
            architecture: await self._preprocess_manifest(sess, image_info, rqst_args, image),
        }
        await self._read_manifest(image, tag, manifests)

    async def _process_docker_v1_image(
        self,
        _tg: aiotools.TaskGroup,
        sess: aiohttp.ClientSession,
        rqst_args: dict[str, Any],
        image: str,
        tag: str,
        image_info: Mapping[str, Any],
    ) -> None:
        log.warning("Docker image manifest v1 is deprecated.")

        architecture = image_info["architecture"]

        manifest_list = [
            {
                "platform": {
                    "os": "linux",
                    "architecture": architecture,
                },
                "digest": tag,
            }
        ]

        rqst_args = copy.deepcopy(rqst_args)
        rqst_args["headers"]["Accept"] = self.MEDIA_TYPE_DOCKER_MANIFEST
        await self._read_manifest_list(sess, manifest_list, rqst_args, image, tag)

    def _read_supported_accelerators_from_labels(self, labels: Mapping[str, Any]) -> set[str]:
        """Returns set of accelerators, represented by the slot name, supported by the image (excluding cpu and mem)."""

        accels = set()
        if _accels_str := labels.get("ai.backend.accelerators"):
            accels |= set(_accels_str.split(","))
        for key in labels.keys():
            if key.startswith(("ai.backend.resource.min", "ai.backend.resource.max")):
                accel = key.split(
                    "."
                )[
                    4
                ]  # "ai" "backend" "resource" "min" "cuda/cpu/mem/..." "device/shares/...(if necessary)"
                if SlotName(accel) not in INTRINSIC_SLOTS_MIN:
                    accels.add(accel)

        return accels

    async def _read_manifest(
        self,
        image: str,
        tag: str,
        manifests: dict[str, dict[str, Any]],
        skip_reason: str | None = None,
    ) -> None:
        """
        Detects if image is compatible with Backend.AI and injects the matadata to database if it complies.
        """
        if not manifests:
            if not skip_reason:
                skip_reason = "missing/deleted"
            log.warning("Skipped image - {}:{} ({})", image, tag, skip_reason)
            progress_msg = f"Skipped {image}:{tag} ({skip_reason})"
            if (reporter := progress_reporter.get()) is not None:
                await reporter.update(1, message=progress_msg)
            return

        for architecture, manifest in manifests.items():
            try:
                try:
                    validate_image_labels(manifest["labels"])
                except t.DataError as e:
                    match e.as_dict():
                        case str() as error_msg:
                            skip_reason = error_msg
                        case dict() as error_data:
                            skip_reason = "; ".join(
                                f"{field} {reason}" for field, reason in error_data.items()
                            )
                    continue
                except ValueError as e:
                    skip_reason = str(e)
                    continue

                update_key = ImageIdentifier(
                    f"{self.registry_name}/{image}:{tag}",
                    architecture,
                )

                updates = {
                    "config_digest": manifest["digest"],
                    "size_bytes": manifest["size"],
                    "labels": manifest["labels"],  # keep the original form
                }
                if "ai.backend.kernelspec" in manifest["labels"]:
                    accels = self._read_supported_accelerators_from_labels(manifest["labels"])
                    if accels:
                        updates["accels"] = ",".join(accels)
                else:
                    # allow every accelerators for non-backend.ai image
                    updates["accels"] = "*"

                all_updates.get().update({
                    update_key: updates,
                })
            except (InvalidImageName, InvalidImageTag) as e:
                skip_reason = str(e)
            finally:
                if skip_reason:
                    log.warning(
                        "Skipped image (_read_manifest inner) - {}:{}/{} ({})",
                        image,
                        tag,
                        architecture,
                        skip_reason,
                    )
                    progress_msg = f"Skipped {image}:{tag}/{architecture} ({skip_reason})"
                else:
                    log.info(
                        "Scanned image - {0}:{1}/{2} ({3})",
                        image,
                        tag,
                        architecture,
                        manifest["digest"],
                    )
                    progress_msg = f"Updated {image}:{tag}/{architecture} ({manifest['digest']})"
                if (reporter := progress_reporter.get()) is not None:
                    await reporter.update(1, message=progress_msg)

    @abstractmethod
    async def fetch_repositories(
        self,
        sess: aiohttp.ClientSession,
    ) -> AsyncIterator[str]:
        yield ""
