# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright (c) 2022 ScyllaDB
import logging
import pprint
import re
from contextlib import suppress

from azure.core.exceptions import ResourceNotFoundError as AzureResourceNotFoundError
from azure.mgmt.compute.models import GalleryImageVersion, CommunityGalleryImageVersion

from sdcm.provision.provisioner import VmArch
from sdcm.utils.azure_utils import AzureService
from sdcm.utils.version_utils import (
    SCYLLA_VERSION_GROUPED_RE,
    is_enterprise,
    ComparableScyllaVersion,
    parse_scylla_version_tag,
)

LOGGER = logging.getLogger(__name__)

# Regex to detect ARM architecture from Azure instance type name.
# Azure ARM instances have a 'p' immediately after the vCPU digit(s),
# e.g. Standard_D8ps_v5, Standard_E4ps_v5, Standard_D2pds_v5
_AZURE_ARM_RE = re.compile(r"^Standard_[A-Z]+\d+p")


def get_arch_from_azure_instance_type(instance_type: str) -> VmArch:
    """Detect CPU architecture from Azure instance type name.

    Azure ARM instances use a 'p' suffix after the vCPU count digit(s)
    to denote Ampere Altra (ARM64) processors.

    Examples:
        Standard_D8ps_v5  -> ARM (has 'p' after '8')
        Standard_L8s_v3   -> X86 (no 'p' after '8')
        Standard_E4ps_v5  -> ARM
        Standard_D2pds_v5 -> ARM
    """
    if instance_type and _AZURE_ARM_RE.match(instance_type):
        return VmArch.ARM
    return VmArch.X86


def vmarch_to_azure(arch: VmArch) -> str:
    """Convert VmArch enum to Azure architecture format.

    Args:
        arch: VmArch enum value

    Returns:
        Azure architecture string (x64 or Arm64)

    Raises:
        ValueError: If architecture is not supported
    """
    if arch is VmArch.X86:
        return "x64"
    elif arch is VmArch.ARM:
        return "Arm64"
    else:
        raise ValueError(f"Unsupported architecture: {arch}")


def get_scylla_images_private_galleries(
    scylla_version: str, region_name: str, arch: VmArch = VmArch.X86, azure_service: AzureService = AzureService()
) -> list[GalleryImageVersion]:
    version_bucket = scylla_version.split(":", 1)
    only_latest = False
    tags_to_search = {"arch": arch.value}

    # Check if this is a full version tag (e.g., "2024.2.5-0.20250221.cb9e2a54ae6d-1")
    if parse_scylla_version_tag(scylla_version):
        # For full version tags, use exact matching with normalized version
        # Azure tags use ~ separator which gets normalized to - for matching
        tags_to_search["scylla_version"] = scylla_version.replace("~", "-")
        branch = "master"
    elif len(version_bucket) == 1:
        return []
    else:
        # Branched version, like master:latest
        branch, build_id = version_bucket
        tags_to_search["branch"] = branch
        if build_id == "latest":
            only_latest = True
        elif build_id == "all":
            pass
        else:
            tags_to_search["build-id"] = build_id

    output = []
    unparsable_scylla_versions = []

    with suppress(AzureResourceNotFoundError):
        try:
            gallery_image_versions = list(
                azure_service.compute.gallery_image_versions.list_by_gallery_image(
                    resource_group_name="SCYLLA-IMAGES",
                    gallery_name="scylladb_dev",
                    gallery_image_name=f"{branch}-{vmarch_to_azure(arch)}",
                )
            )
        except AzureResourceNotFoundError:
            gallery_image_versions = []

        if not gallery_image_versions:
            gallery_image_versions = list(
                azure_service.compute.gallery_image_versions.list_by_gallery_image(
                    resource_group_name="SCYLLA-IMAGES",
                    gallery_name="scylladb_dev",
                    gallery_image_name=branch,
                )
            )
        for image in gallery_image_versions:
            if image.location != region_name or image.tags.get("name", "").startswith("debug-"):
                continue
            image.tags["scylla_version"] = image.tags.get("scylla_version", image.tags.get("ScyllaVersion"))
            # Filter by tags
            for tag_name, expected_value in tags_to_search.items():
                actual_value = image.tags.get(tag_name)
                if callable(expected_value):
                    if not expected_value(actual_value):
                        break
                elif expected_value != actual_value:
                    break
            else:
                if SCYLLA_VERSION_GROUPED_RE.match(image.tags.get("scylla_version")):
                    output.append(image)
                else:
                    unparsable_scylla_versions.append(f"{image.name}: {image.tags.get('scylla_version')}")
    if unparsable_scylla_versions:
        LOGGER.warning("Couldn't parse scylla version from images: %s", str(unparsable_scylla_versions))
    output.sort(key=lambda img: int(SCYLLA_VERSION_GROUPED_RE.match(img.tags.get("scylla_version")).group("date")))

    if only_latest:
        return output[-1:]
    return output[::-1]


def get_scylla_images(
    scylla_version: str, region_name: str, arch: VmArch = VmArch.X86, azure_service: AzureService = AzureService()
) -> list[GalleryImageVersion]:
    version_bucket = scylla_version.split(":", 1)
    only_latest = False
    tags_to_search = {"arch": arch.value}
    if output := get_scylla_images_private_galleries(
        scylla_version=scylla_version, region_name=region_name, arch=arch, azure_service=azure_service
    ):
        return output

    # Check if this is a full version tag (e.g., "2024.2.5-0.20250221.cb9e2a54ae6d-1")
    if parse_scylla_version_tag(scylla_version):
        # For full version tags, use exact matching with normalized version
        # Azure tags use ~ separator which gets normalized to - for matching
        tags_to_search["scylla_version"] = scylla_version.replace("~", "-")
    elif len(version_bucket) == 1:
        if "." in scylla_version:
            # Plain version, like 4.5.0
            tags_to_search["scylla_version"] = (
                lambda ver: ver and ver.replace("~", "-").startswith(scylla_version) and "-dev" not in ver
            )
        else:
            # commit id
            tags_to_search["scylla_version"] = lambda ver: ver and scylla_version[:9] in ver
    else:
        # Branched version, like master:latest
        branch, build_id = version_bucket
        tags_to_search["branch"] = branch
        if build_id == "latest":
            only_latest = True
        elif build_id == "all":
            pass
        else:
            tags_to_search["build-id"] = build_id
    output = []
    unparsable_scylla_versions = []
    with suppress(AzureResourceNotFoundError):
        gallery_image_versions = azure_service.compute.images.list_by_resource_group(
            resource_group_name="SCYLLA-IMAGES",
        )
        for image in gallery_image_versions:
            if image.location != region_name or image.name.startswith("debug-"):
                continue
            image.tags["scylla_version"] = image.tags.get("scylla_version", image.tags.get("ScyllaVersion"))
            # Filter by tags
            for tag_name, expected_value in tags_to_search.items():
                actual_value = image.tags.get(tag_name).replace("~", "-")
                if callable(expected_value):
                    if not expected_value(actual_value):
                        break
                elif expected_value != actual_value:
                    break
            else:
                if SCYLLA_VERSION_GROUPED_RE.match(image.tags.get("scylla_version")):
                    output.append(image)
                else:
                    unparsable_scylla_versions.append(f"{image.name}: {image.tags.get('scylla_version')}")
    if unparsable_scylla_versions:
        LOGGER.warning("Couldn't parse scylla version from images: %s", str(unparsable_scylla_versions))
    output.sort(key=lambda img: int(SCYLLA_VERSION_GROUPED_RE.match(img.tags.get("scylla_version")).group("date")))

    if only_latest:
        return output[-1:]
    return output[::-1]


def get_released_scylla_images(
    scylla_version: str, region_name: str, arch: VmArch = VmArch.X86, azure_service: AzureService = AzureService()
) -> list[CommunityGalleryImageVersion]:
    if arch != VmArch.X86:
        LOGGER.warning("--arch option not implemented currently for Azure machine images.")
    branch_version = ".".join(scylla_version.split(".")[:2])
    if is_enterprise(branch_version) and ComparableScyllaVersion(branch_version) < "2025.1.0":
        gallery_image_name = f"scylla-enterprise-{branch_version}"
    else:
        gallery_image_name = f"scylla-{branch_version}"
    community_gallery_images = azure_service.compute.community_gallery_image_versions.list(
        location=region_name,
        gallery_image_name=gallery_image_name,
        public_gallery_name="scylladb-7e8d8a04-23db-487d-87ec-0e175c0615bb",
    )
    community_gallery_images: list[CommunityGalleryImageVersion] = list(community_gallery_images)
    community_gallery_images.sort(key=lambda x: x.properties.published_date, reverse=True)

    # a specific version was asked, return only that version
    if branch_version != scylla_version:
        specific_version = [image for image in community_gallery_images if image.name == scylla_version]
        assert specific_version, f"`{scylla_version}` wasn't found in:\n{pprint.pformat(community_gallery_images)}"
        return specific_version

    return community_gallery_images


IMAGE_URL_REGEX = re.compile(
    r".*/resourceGroups/(?P<resource_group_name>.*)/providers/Microsoft.Compute/images/(?P<image_name>.*)"
)


COMMUNITY_IMAGE_URL_REGEX = re.compile(
    r".*/CommunityGalleries/(?P<public_gallery_name>.*)/Images/(?P<gallery_image_name>.*)/Versions/(?P<version>.*)"
)


PRIVATE_GALLERY_IMAGE_URL_REGEX = re.compile(
    r".*/subscriptions/([^/]+)/resourceGroups/(?P<resource_group_name>[^/]+)/providers/Microsoft.Compute/galleries/(?P<gallery_name>[^/]+)/images/(?P<gallery_image_name>[^/]+)/versions/(?P<gallery_image_version_name>[^/]+)"
)


def get_image_tags(link: str) -> dict:
    if match := IMAGE_URL_REGEX.search(link):
        params = match.groupdict()
        azure_image: GalleryImageVersion = AzureService().compute.images.get(**params)
        return azure_image.tags
    elif match := PRIVATE_GALLERY_IMAGE_URL_REGEX.search(link):
        params = match.groupdict()
        azure_image: GalleryImageVersion = AzureService().compute.gallery_image_versions.get(**params)
        return azure_image.tags
    elif match := COMMUNITY_IMAGE_URL_REGEX.search(link):
        params = match.groupdict()
        return dict(scylla_version=params.get("version"), user_data_format_version="3")
    return {}
