# 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 base64
import time
from datetime import datetime, timezone
import logging
import os
from dataclasses import dataclass, field

from typing import Dict, Optional, Any, List

import binascii
from azure.core.exceptions import ResourceNotFoundError, AzureError, ODataV4Error
from azure.mgmt.compute.models import VirtualMachine, RunCommandInput
from invoke import Result

from sdcm.provision.azure.kms_provider import AzureKmsProvider
from sdcm.provision.provisioner import (
    InstanceDefinition,
    PricingModel,
    ProvisionError,
    OperationPreemptedError,
    StuckVMProvisioningError,
)
from sdcm.provision.user_data import UserDataBuilder
from sdcm.sct_events.system import SpotTerminationEvent
from sdcm.utils.azure_utils import AzureService

LOGGER = logging.getLogger(__name__)

SCT_RESOURCE_GROUP_PREFIX = "SCT-"

DEFAULT_STUCK_VM_TIMEOUT = 900
DEFAULT_STUCK_VM_POLL_INTERVAL = 30


@dataclass
class VirtualMachineProvider:
    _resource_group_name: str
    _region: str
    _az: str
    _enable_azure_kms: bool = False
    _azure_service: AzureService = AzureService()
    _cache: Dict[str, VirtualMachine] = field(default_factory=dict)
    _stuck_vm_timeout: int = DEFAULT_STUCK_VM_TIMEOUT
    _stuck_vm_poll_interval: int = DEFAULT_STUCK_VM_POLL_INTERVAL

    def __post_init__(self):
        """Discover existing virtual machines for resource group."""
        try:
            v_ms = self._azure_service.compute.virtual_machines.list(self._resource_group_name)
            for _v_m in v_ms:
                v_m = self._azure_service.compute.virtual_machines.get(self._resource_group_name, _v_m.name)
                if v_m.provisioning_state == "Succeeded":
                    self._cache[v_m.name] = v_m
                elif v_m.provisioning_state != "Deleting":
                    LOGGER.warning(
                        "Discovered VM %s in non-final provisioning state '%s'; not caching it as ready",
                        v_m.name,
                        v_m.provisioning_state,
                    )
        except ResourceNotFoundError:
            pass

    def get_or_create(
        self,
        definitions: List[InstanceDefinition],
        nics_ids: List[str],
        pricing_model: PricingModel,
        deadline: Optional[float] = None,
    ) -> List[VirtualMachine]:
        v_ms = []
        pollers = []
        error_to_raise = None
        for definition, nic_id in zip(definitions, nics_ids):
            if definition.name in self._cache:
                v_ms.append(self._cache[definition.name])
                continue
            LOGGER.info("Creating '%s' VM in resource group %s...", definition.name, self._resource_group_name)
            LOGGER.info("Instance params: %s", definition)
            tags = definition.tags | {
                "ssh_user": definition.user_name,
                "ssh_key": definition.ssh_key.name,
                "creation_time": datetime.now(timezone.utc).isoformat(sep=" ", timespec="seconds"),
            }
            tags = self._replace_null_value_from_tags_with_empty_string(tags)
            params = {
                "location": self._region,
                "zones": [self._az] if self._az else [],
                "tags": tags,
                "properties": {
                    "hardwareProfile": {
                        "vmSize": definition.type,
                    },
                    "networkProfile": {
                        "networkInterfaces": [{"id": nic_id, "properties": {"deleteOption": "Detach"}}],
                    },
                    "diagnosticsProfile": {
                        "bootDiagnostics": {
                            "enabled": True,
                        }
                    },
                },
            }

            if self._enable_azure_kms:
                self._kms_provider = AzureKmsProvider(
                    self._resource_group_name, self._region, self._az, self._azure_service
                )
                # Extract test_id from resource group name
                test_id = self._resource_group_name.split(SCT_RESOURCE_GROUP_PREFIX)[-1][:36]
                vault_info = self._kms_provider.get_or_create_keyvault_and_identity(test_id)
                if vault_info is None:
                    error_msg = (
                        f"Failed to setup Azure KMS for VM {definition.name}. "
                        "Key Vault creation failed. "
                        "Check logs for details."
                    )
                    LOGGER.error(error_msg)
                    raise ProvisionError(error_msg)
                params["identity"] = {
                    "type": "UserAssigned",
                    "userAssignedIdentities": {vault_info["identity_id"]: {}},
                }
                LOGGER.info(f"Azure Key Vault enabled for {definition.name}")
            else:
                LOGGER.info(f"Azure Key Vault disabled for {definition.name}")

            if definition.user_data is None:
                # in case we use specialized image, we don't change things like computer_name, usernames, ssh_keys
                os_profile = {}
            else:
                builder = UserDataBuilder(user_data_objects=definition.user_data)
                custom_data = builder.build_user_data_yaml()
                os_profile = self._get_os_profile(
                    computer_name=definition.name,
                    admin_username=definition.user_name,
                    admin_password=binascii.hexlify(os.urandom(20)).decode(),
                    ssh_public_key=definition.ssh_key.public_key.decode(),
                    custom_data=custom_data,
                )
                params["properties"].update(
                    {
                        "userData": base64.b64encode(builder.get_scylla_machine_image_json().encode("utf-8")).decode(
                            "latin-1"
                        )
                    }
                )
            storage_profile = self._get_scylla_storage_profile(
                image_id=definition.image_id, name=definition.name, disk_size=definition.root_disk_size
            )
            params["properties"].update(os_profile)
            params["properties"].update(storage_profile)
            params["properties"].update(self._get_pricing_params(pricing_model))
            try:
                poller = self._azure_service.compute.virtual_machines.begin_create_or_update(
                    resource_group_name=self._resource_group_name, vm_name=definition.name, parameters=params
                )
                pollers.append((definition, poller))
            except AzureError as err:
                LOGGER.error("Error when sending create vm request for VM %s: %s", definition.name, str(err))
                error_to_raise = err
        pending_names = []
        for definition, poller in pollers:
            try:
                # check for create-time errors quickly
                poller.wait(timeout=self._stuck_vm_poll_interval)
            except ODataV4Error as err:
                LOGGER.error("Error when waiting for VM %s: %s", definition.name, str(err))
                if err.code == "OperationPreempted":
                    raise OperationPreemptedError(err)  # spot instance preemption, abort provision immediately
                error_to_raise = err
                continue
            except AzureError as err:
                LOGGER.error("Error when waiting for creation of VM %s: %s", definition.name, str(err))
                error_to_raise = err
                continue
            pending_names.append(definition.name)
        provisioned, stuck_vm_names = self._wait_all_provisioned(pending_names, pricing_model, deadline)
        for v_m in provisioned:
            self._cache[v_m.name] = v_m
            v_ms.append(v_m)
        if error_to_raise:
            raise ProvisionError(error_to_raise)
        if stuck_vm_names:
            raise StuckVMProvisioningError(stuck_vm_names)
        return v_ms

    def _wait_all_provisioned(
        self, names: List[str], pricing_model: PricingModel, deadline: Optional[float] = None
    ) -> tuple[List[VirtualMachine], List[str]]:
        """Wait for VMs to reach ``ProvisioningState=Succeeded`` within one shared timeout.

        Returns two lists: provisioned VMs and still-pending VM names. All VMs share the same wait
        window based on ``self._stuck_vm_timeout``, optionally capped by ``deadline``.

        If a spot VM disappears during polling, ``OperationPreemptedError`` is raised by
        ``_poll_provisioned_vm``. If a non-spot VM disappears, ``ResourceNotFoundError`` is re-raised.
        Other Azure polling errors are logged and treated as transient until the timeout expires.
        """
        wait_deadline = time.monotonic() + self._stuck_vm_timeout
        if deadline is not None:
            wait_deadline = min(wait_deadline, deadline)

        pending = list(names)
        provisioned = []
        while pending and time.monotonic() < wait_deadline:
            for name in list(pending):
                try:
                    v_m = self._poll_provisioned_vm(name, pricing_model)
                except ResourceNotFoundError:
                    raise
                except AzureError as err:
                    LOGGER.warning("Transient error polling VM %s: %s; keeping it pending", name, err)
                    continue

                if v_m is not None:
                    provisioned.append(v_m)
                    pending.remove(name)
            if pending and time.monotonic() < wait_deadline:
                time.sleep(self._stuck_vm_poll_interval)

        return provisioned, pending

    def _poll_provisioned_vm(self, name: str, pricing_model: PricingModel) -> Optional[VirtualMachine]:
        """Poll a VM once: return it if 'Succeeded', None otherwise.

        Raises OperationPreemptedError (with a SpotTerminationEvent) if a spot VM vanished (404);
        re-raises ResourceNotFoundError for a non-spot VM.
        """
        try:
            vm = self._azure_service.compute.virtual_machines.get(
                self._resource_group_name,
                name,
                expand="instanceView",
            )
        except ResourceNotFoundError as exc:
            if not pricing_model.is_spot():
                raise
            message = (
                "Azure removed the spot VM during provisioning (ResourceNotFound 404) - "
                "spot eviction / capacity-allocation failure"
            )
            SpotTerminationEvent(node=name, message=message).publish_or_dump(default_logger=LOGGER)
            raise OperationPreemptedError(f"spot VM {name} evicted during provisioning: {message}") from exc
        if vm.instance_view and vm.instance_view.statuses:
            status_summary = ", ".join(
                f"{status.code}: {status.display_status}" for status in vm.instance_view.statuses
            )
            LOGGER.debug("VM %s instance state: %s", name, status_summary)
        if vm.provisioning_state == "Succeeded":
            return vm
        return None

    def list(self):
        return list(self._cache.values())

    def delete(self, name: str, wait: bool = True):
        LOGGER.info("Triggering termination of instance: %s", name)
        self._azure_service.compute.virtual_machines.begin_update(
            self._resource_group_name,
            vm_name=name,
            parameters={
                "properties": {"storageProfile": {"osDisk": {"createOption": "FromImage", "deleteOption": "Delete"}}}
            },
        )

        task = self._azure_service.compute.virtual_machines.begin_delete(self._resource_group_name, vm_name=name)
        if wait is True:
            LOGGER.info("Waiting for termination of instance: %s...", name)
            task.wait()
            LOGGER.info("Instance %s has been terminated.", name)
        self._cache.pop(name, None)

    def reboot(self, name: str, wait: bool = True, hard: bool = False) -> None:
        LOGGER.info("Triggering reboot of instance: %s", name)

        # Try walinuxagent approach first
        try:
            flags = "-ff" if hard else "-f"
            self.run_command(name, f"reboot {flags}")
            LOGGER.debug("Successfully triggered reboot via walinuxagent for instance: %s", name)
        except Exception as ex:  # noqa: BLE001
            # Fallback to Azure SDK when walinuxagent is not available or fails
            LOGGER.warning("Failed to reboot instance %s via walinuxagent (%s), falling back to Azure SDK", name, ex)
            if hard:
                LOGGER.warning("Hard reboot requested but Azure SDK does not distinguish between hard and soft reboot")
            task = self._azure_service.compute.virtual_machines.begin_restart(self._resource_group_name, vm_name=name)
            if wait:
                LOGGER.info("Waiting for SDK-triggered reboot of instance: %s...", name)
                task.wait()
                LOGGER.info("Instance %s has been rebooted via Azure SDK.", name)
                return

        # Wait for walinuxagent-triggered reboot to complete
        if wait:
            start_time = time.time()
            while time.time() - start_time < 600:  # 10 minutes
                time.sleep(10)
                instance_view = self._azure_service.compute.virtual_machines.instance_view(
                    self._resource_group_name, vm_name=name
                )
                if instance_view and instance_view.statuses[-1].display_status == "VM running":
                    break

    def add_tags(self, name: str, tags: Dict[str, str]) -> VirtualMachine:
        """Adds tags to instance (with waiting for completion)"""
        if name not in self._cache:
            raise AttributeError(f"Instance '{name}' does not exist in resource group '{self._resource_group_name}'")
        current_tags = self._cache[name].tags
        tags = self._replace_null_value_from_tags_with_empty_string(tags)
        current_tags.update(tags)
        self._azure_service.compute.virtual_machines.begin_update(
            self._resource_group_name, name, parameters={"tags": current_tags}
        ).wait()
        v_m = self._azure_service.compute.virtual_machines.get(self._resource_group_name, name)
        self._cache[v_m.name] = v_m
        return v_m

    @staticmethod
    def _get_os_profile(
        computer_name: str, admin_username: str, admin_password: str, ssh_public_key: str, custom_data: str
    ) -> Dict[str, Any]:
        os_profile = {
            "osProfile": {
                "computerName": computer_name,
                "adminUsername": admin_username,
                "adminPassword": admin_password,
                "customData": base64.b64encode(custom_data.encode("utf-8")).decode("latin-1"),
                "linuxConfiguration": {
                    "disablePasswordAuthentication": True,
                    "ssh": {
                        "publicKeys": [
                            {
                                "path": f"/home/{admin_username}/.ssh/authorized_keys",
                                "keyData": ssh_public_key,
                            }
                        ],
                    },
                },
            }
        }
        return os_profile

    @staticmethod
    def _get_scylla_storage_profile(image_id: str, name: str, disk_size: Optional[int] = None) -> Dict[str, Any]:
        """Creates storage profile based on image_id. image_id may refer to scylla-crafted images
        (starting with '/subscription') or to 'Urn' of image (see output of e.g. `az vm image list --output table`)"""
        storage_profile = {
            "storageProfile": {
                "osDisk": {
                    "name": f"{name}-os-disk",
                    "osType": "linux",
                    "deleteOption": "Delete",
                    "caching": "ReadWrite",
                    "createOption": "FromImage",
                    "managedDisk": {
                        "storageAccountType": "StandardSSD_LRS",  # SSD
                    },
                }
                | ({} if disk_size is None else {"diskSizeGb": disk_size}),
            }
        }
        if image_id.startswith("/subscriptions/"):
            storage_profile["storageProfile"].update({"imageReference": {"id": image_id}})
        elif image_id.startswith("/CommunityGalleries/"):
            storage_profile["storageProfile"].update({"imageReference": {"communityGalleryImageId": image_id}})
        else:
            image_reference_values = image_id.split(":")
            storage_profile["storageProfile"].update(
                {
                    "imageReference": {
                        "publisher": image_reference_values[0],
                        "offer": image_reference_values[1],
                        "sku": image_reference_values[2],
                        "version": image_reference_values[3],
                    },
                }
            )
        return storage_profile

    def run_command(self, name: str, command: str) -> Result:
        if name not in self._cache:
            raise AttributeError(f"Instance '{name}' does not exist in resource group '{self._resource_group_name}'")
        LOGGER.debug("Running command '%s' on instance: %s", command, name)
        command_object = RunCommandInput(command_id="RunShellScript", script=[command])
        result = self._azure_service.compute.virtual_machines.begin_run_command(
            self._resource_group_name, name, command_object
        ).result()
        LOGGER.debug("Finished running command '%s' on instance: %s: %s", command, name, result.value[0])
        try:
            stdout, stderr = result.value[0].message.split("[stdout]\n")[1].split("\n[stderr]\n")[0:2]
            exited = 0  # doesn't mean it passed, Azure does not provide rc
        except IndexError:
            stdout, stderr = result.value[0].message, ""
            exited = 1
        return Result(stdout=stdout, stderr=stderr, exited=exited, encoding="utf-8")

    @staticmethod
    def _get_pricing_params(pricing_model: PricingModel):
        if pricing_model != PricingModel.ON_DEMAND:
            return {
                "priority": "Spot",  # possible values are "Regular", "Low", or "Spot"
                "evictionPolicy": "Delete",  # can be "Deallocate" or "Delete", Deallocate leaves disks intact
                "billingProfile": {
                    "maxPrice": -1,  # -1 indicates the VM shouldn't be evicted for price reasons
                },
            }
        else:
            return {"priority": "Regular"}

    def clear_cache(self):
        self._cache = {}

    @staticmethod
    def _replace_null_value_from_tags_with_empty_string(tags: Dict[str, str]) -> Dict[str, str]:
        """Azure API does not accept 'null' as value for tags, so we replace it with empty string."""
        for key, value in tags.items():
            if value == "null":
                LOGGER.warning(
                    "Tag '%s' has value 'null' which is not allowed by Azure API. Replacing with empty string.", key
                )
                tags[key] = ""
        return tags
