# Copyright (c) 2022, JINGROW
# For license information, please see license.txt
from __future__ import annotations

import boto3
import jingrow
from jingrow.core.utils import find
from jingrow.model.page import Page
from oci.core import ComputeClient
from oci.core.models import CreateImageDetails
from tenacity import retry, stop_after_attempt, wait_fixed
from tenacity.retry import retry_if_result


class VirtualMachineImage(Page):
	# begin: auto-generated types
	# This code is auto-generated. Do not modify anything in this block.

	from typing import TYPE_CHECKING

	if TYPE_CHECKING:
		from jingrow.types import DF

		from jcloud.jcloud.pagetype.virtual_machine_image_volume.virtual_machine_image_volume import (
			VirtualMachineImageVolume,
		)

		cluster: DF.Link
		copied_from: DF.Link | None
		has_data_volume: DF.Check
		image_id: DF.Data | None
		instance_id: DF.Data
		mariadb_root_password: DF.Password | None
		platform: DF.Data | None
		public: DF.Check
		region: DF.Link
		root_size: DF.Int
		series: DF.Literal["n", "f", "m", "c", "p", "e", "r"]
		size: DF.Int
		snapshot_id: DF.Data | None
		status: DF.Literal["Pending", "Available", "Unavailable"]
		virtual_machine: DF.Link
		volumes: DF.Table[VirtualMachineImageVolume]
	# end: auto-generated types

	PAGETYPE = "Virtual Machine Image"

	def after_insert(self):
		self.set_credentials()
		if self.copied_from:
			self.create_image_from_copy()
		else:
			self.create_image()

	def create_image(self):
		cluster = jingrow.get_pg("Cluster", self.cluster)
		if cluster.cloud_provider == "AWS EC2":
			volumes = self.get_volumes_from_virtual_machine()
			response = self.client.create_image(
				InstanceId=self.instance_id,
				Name=f"Jingrow {self.name} - {self.virtual_machine}",
				BlockDeviceMappings=volumes,
			)
			self.image_id = response["ImageId"]
		elif cluster.cloud_provider == "OCI":
			image = self.client.create_image(
				CreateImageDetails(
					compartment_id=cluster.oci_tenancy,
					display_name=f"Jingrow {self.name} - {self.virtual_machine}",
					instance_id=self.instance_id,
				)
			).data
			self.image_id = image.id
		self.sync()

	def create_image_from_copy(self):
		source = jingrow.get_pg("Virtual Machine Image", self.copied_from)
		response = self.client.copy_image(
			Name=f"Jingrow {self.name} - {self.virtual_machine}",
			SourceImageId=source.image_id,
			SourceRegion=source.region,
		)
		self.image_id = response["ImageId"]
		self.sync()

	def set_credentials(self):
		if self.series == "m" and jingrow.db.exists("Database Server", self.virtual_machine):
			self.mariadb_root_password = jingrow.get_pg("Database Server", self.virtual_machine).get_password(
				"mariadb_root_password"
			)

	@jingrow.whitelist()
	def sync(self):  # noqa: C901
		cluster = jingrow.get_pg("Cluster", self.cluster)
		if cluster.cloud_provider == "AWS EC2":
			images = self.client.describe_images(ImageIds=[self.image_id])["Images"]
			if images:
				image = images[0]
				self.status = self.get_aws_status_map(image["State"])
				self.platform = image["Architecture"]
				volume = find(image["BlockDeviceMappings"], lambda x: "Ebs" in x)
				# This information is not accurate for images created from multiple volumes
				attached_snapshots = []
				if volume and "SnapshotId" in volume["Ebs"]:
					self.snapshot_id = volume["Ebs"]["SnapshotId"]
				for volume in image["BlockDeviceMappings"]:
					if "Ebs" not in volume:
						# We don't care about non-EBS (instance store) volumes
						continue
					snapshot_id = volume["Ebs"].get("SnapshotId")
					if not snapshot_id:
						# We don't care about volumes without snapshots
						continue
					attached_snapshots.append(snapshot_id)
					existing = find(self.volumes, lambda x: x.snapshot_id == snapshot_id)
					device = volume["DeviceName"]
					volume_type = volume["Ebs"]["VolumeType"]
					size = volume["Ebs"]["VolumeSize"]
					if existing:
						existing.device = device
						existing.volume_type = volume_type
						existing.size = size
					else:
						self.append(
							"volumes",
							{
								"snapshot_id": snapshot_id,
								"device": device,
								"volume_type": volume_type,
								"size": size,
							},
						)
				for volume in list(self.volumes):
					if volume.snapshot_id not in attached_snapshots:
						self.remove(volume)

				self.size = self.get_data_volume().size
				self.root_size = self.get_data_volume().size
			else:
				self.status = "Unavailable"
		elif cluster.cloud_provider == "OCI":
			image = self.client.get_image(self.image_id).data
			self.status = self.get_oci_status_map(image.lifecycle_state)
			if image.size_in_mbs:
				self.size = image.size_in_mbs // 1024

		self.save()
		return self.status

	@retry(
		retry=retry_if_result(lambda result: result != "Available"),
		wait=wait_fixed(60),
		stop=stop_after_attempt(10),
	)
	def wait_for_availability(self):
		"""Retries sync until the image is available"""
		return self.sync()

	@jingrow.whitelist()
	def copy_image(self, cluster: str):
		image = jingrow.copy_pg(self)
		image.copied_from = self.name
		image.cluster = cluster
		return image.insert()

	@jingrow.whitelist()
	def delete_image(self):
		cluster = jingrow.get_pg("Cluster", self.cluster)
		if cluster.cloud_provider == "AWS EC2":
			self.client.deregister_image(ImageId=self.image_id)
			if self.snapshot_id:
				self.client.delete_snapshot(SnapshotId=self.snapshot_id)
		elif cluster.cloud_provider == "OCI":
			self.client.delete_image(self.image_id)
		self.sync()

	def get_aws_status_map(self, status):
		return {
			"pending": "Pending",
			"available": "Available",
		}.get(status, "Unavailable")

	def get_oci_status_map(self, status):
		return {
			"PROVISIONING": "Pending",
			"IMPORTING": "Pending",
			"AVAILABLE": "Available",
			"EXPORTING": "Pending",
			"DISABLED": "Unavailable",
			"DELETED": "Unavailable",
		}.get(status, "Unavailable")

	def get_volumes_from_virtual_machine(self):
		machine = jingrow.get_pg("Virtual Machine", self.virtual_machine)
		volumes = []
		for volume in machine.volumes:
			volumes.append(
				{
					"DeviceName": volume.device,
					"Ebs": {
						"DeleteOnTermination": True,
						"VolumeSize": volume.size,
						"VolumeType": volume.volume_type,
					},
				}
			)
		return volumes

	@property
	def client(self):
		cluster = jingrow.get_pg("Cluster", self.cluster)
		if cluster.cloud_provider == "AWS EC2":
			return boto3.client(
				"ec2",
				region_name=self.region,
				aws_access_key_id=cluster.aws_access_key_id,
				aws_secret_access_key=cluster.get_password("aws_secret_access_key"),
			)
		if cluster.cloud_provider == "OCI":
			return ComputeClient(cluster.get_oci_config())
		return None

	@classmethod
	def get_available_for_series(
		cls, series: str, region: str | None = None, platform: str | None = None
	) -> str | None:
		images = jingrow.qb.PageType(cls.PAGETYPE)
		get_available_images = (
			jingrow.qb.from_(images)
			.select("name")
			.where(images.status == "Available")
			.where(images.public == 1)
			.where(
				images.series == series,
			)
			.orderby(images.creation, order=jingrow.qb.desc)
		)
		if region:
			get_available_images = get_available_images.where(images.region == region)
		if platform:
			get_available_images = get_available_images.where(images.platform == platform)
		available_images = get_available_images.run(as_dict=True)
		if not available_images:
			return None
		return available_images[0].name

	def get_root_volume(self):
		# This only works for AWS
		if len(self.volumes) == 1:
			return self.volumes[0]

		volume = find(self.volumes, lambda v: v.device == "/dev/sda1")
		if volume:
			return volume
		return jingrow._dict({"size": 0})

	def get_data_volume(self):
		if not self.has_data_volume:
			return self.get_root_volume()

		# This only works for AWS
		if len(self.volumes) == 1:
			return self.volumes[0]

		volume = find(self.volumes, lambda v: v.device != "/dev/sda1")
		if volume:
			return volume
		return jingrow._dict({"size": 0})
