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

import boto3
import jingrow
import jingrow.utils
import pytz
from botocore.exceptions import ClientError
from jingrow.model.page import Page
from oci.core import BlockstorageClient

from jcloud.utils import log_error


class VirtualDiskSnapshot(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

		cluster: DF.Link | None
		created_for_site_update: DF.Check
		duration: DF.Duration | None
		mariadb_root_password: DF.Password | None
		progress: DF.Data | None
		region: DF.Link
		size: DF.Int
		snapshot_id: DF.Data
		start_time: DF.Datetime | None
		status: DF.Literal["Pending", "Completed", "Error", "Recovering", "Recoverable", "Unavailable"]
		virtual_machine: DF.Link
		volume_id: DF.Data | None
	# end: auto-generated types

	def before_insert(self):
		self.set_credentials()

	def after_insert(self):
		self.sync()

	def set_credentials(self):
		series = jingrow.db.get_value("Virtual Machine", self.virtual_machine, "series")
		if 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"
			)

	def on_update(self):
		if self.has_value_changed("status") and self.status == "Unavailable":
			site_backup_name = jingrow.db.exists("Site Backup", {"database_snapshot": self.name})
			if site_backup_name:
				jingrow.db.set_value("Site Backup", site_backup_name, "files_availability", "Unavailable")

		if self.has_value_changed("status") and self.status == "Completed":
			old_pg = self.get_pg_before_save()
			if old_pg is None or old_pg.status != "Pending":
				return

			self.duration = jingrow.utils.cint(
				jingrow.utils.time_diff_in_seconds(jingrow.utils.now_datetime(), self.creation)
			)
			self.save()

	@jingrow.whitelist()
	def sync(self):
		cluster = jingrow.get_pg("Cluster", self.cluster)
		if cluster.cloud_provider == "AWS EC2":
			try:
				snapshots = self.client.describe_snapshots(SnapshotIds=[self.snapshot_id])["Snapshots"]
				if snapshots:
					snapshot = snapshots[0]
					self.volume_id = snapshot["VolumeId"]
					self.snapshot_id = snapshot["SnapshotId"]

					self.status = self.get_aws_status_map(snapshot["State"])
					self.description = snapshot["Description"]
					self.size = snapshot["VolumeSize"]
					self.start_time = jingrow.utils.format_datetime(
						snapshot["StartTime"], "yyyy-MM-dd HH:mm:ss"
					)
					self.progress = snapshot["Progress"]
			except Exception:
				self.status = "Unavailable"
		elif cluster.cloud_provider == "OCI":
			if ".bootvolumebackup." in self.snapshot_id:
				snapshot = self.client.get_boot_volume_backup(self.snapshot_id).data
				self.volume_id = snapshot.boot_volume_id
			else:
				snapshot = self.client.get_volume_backup(self.snapshot_id).data
				self.volume_id = snapshot.volume_id
			self.status = self.get_oci_status_map(snapshot.lifecycle_state)
			self.description = snapshot.display_name
			self.size = snapshot.size_in_gbs

			self.start_time = jingrow.utils.format_datetime(
				snapshot.time_created.astimezone(pytz.timezone(jingrow.utils.get_system_timezone())),
				"yyyy-MM-dd HH:mm:ss",
			)
		self.save()

	@jingrow.whitelist()
	def delete_snapshot(self):
		self.sync()
		if self.status == "Unavailable":
			return
		cluster = jingrow.get_pg("Cluster", self.cluster)
		if cluster.cloud_provider == "AWS EC2":
			try:
				self.client.delete_snapshot(SnapshotId=self.snapshot_id)
			except ClientError as e:
				if e.response["Error"]["Code"] == "InvalidSnapshot.InUse":
					jingrow.msgprint("Snapshot is in use", alert=True)
				else:
					raise e
		elif cluster.cloud_provider == "OCI":
			if ".bootvolumebackup." in self.snapshot_id:
				self.client.delete_boot_volume_backup(self.snapshot_id)
			else:
				self.client.delete_volume_backup(self.snapshot_id)
		self.sync()

	def get_aws_status_map(self, status):
		return {
			"pending": "Pending",
			"completed": "Completed",
			"error": "Error",
			"recovering": "Recovering",
			"recoverable": "Recoverable",
		}.get(status, "Unavailable")

	def get_oci_status_map(self, status):
		return {
			"CREATING": "Pending",
			"AVAILABLE": "Completed",
			"TERMINATING": "Pending",
			"TERMINATED": "Unavailable",
			"FAULTY": "Error",
			"REQUEST_RECEIVED": "Pending",
		}.get(status, "Unavailable")

	def create_volume(self, availability_zone: str, iops: int = 3000, throughput: int | None = None) -> str:
		self.sync()
		if self.status != "Completed":
			raise Exception("Snapshot is unavailable")
		if throughput is None:
			throughput = 125
		response = self.client.create_volume(
			SnapshotId=self.snapshot_id,
			AvailabilityZone=availability_zone,
			VolumeType="gp3",
			TagSpecifications=[
				{
					"ResourceType": "volume",
					"Tags": [{"Key": "Name", "Value": f"Jingrow Snapshot - {self.name}"}],
				},
			],
			Iops=iops,
			Throughput=throughput,
		)
		return response["VolumeId"]

	@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 BlockstorageClient(cluster.get_oci_config())
		return None


def sync_snapshots():
	snapshots = jingrow.get_all(
		"Virtual Disk Snapshot", {"status": "Pending", "created_for_site_update": ["!=", 1]}
	)
	for snapshot in snapshots:
		try:
			jingrow.get_pg("Virtual Disk Snapshot", snapshot.name).sync()
			jingrow.db.commit()
		except Exception:
			jingrow.db.rollback()
			log_error(title="Virtual Disk Snapshot Sync Error", virtual_snapshot=snapshot.name)


def delete_old_snapshots():
	snapshots = jingrow.get_all(
		"Virtual Disk Snapshot",
		{
			"status": "Completed",
			"creation": ("<=", jingrow.utils.add_days(None, -2)),
			"created_for_site_update": 0,
		},
		pluck="name",
		order_by="creation asc",
		limit=500,
	)
	for snapshot in snapshots:
		try:
			jingrow.get_pg("Virtual Disk Snapshot", snapshot).delete_snapshot()
			jingrow.db.commit()
		except Exception:
			log_error("Virtual Disk Snapshot Delete Error", snapshot=snapshot)
			jingrow.db.rollback()


def sync_all_snapshots_from_aws():
	regions = jingrow.get_all("Cloud Region", {"provider": "AWS EC2"}, pluck="name")
	for region in regions:
		if not jingrow.db.exists("Virtual Disk Snapshot", {"region": region}):
			continue
		random_snapshot = jingrow.get_pg(
			"Virtual Disk Snapshot",
			{
				"region": region,
			},
		)
		client = random_snapshot.client
		paginator = client.get_paginator("describe_snapshots")
		for page in paginator.paginate(OwnerIds=["self"], Filters=[{"Name": "tag-key", "Values": ["Name"]}]):
			for snapshot in page["Snapshots"]:
				if _should_skip_snapshot(snapshot):
					continue
				try:
					delete_duplicate_snapshot_docs(snapshot)
					if _update_snapshot_if_exists(snapshot, random_snapshot):
						continue
					tag_name = next(tag["Value"] for tag in snapshot["Tags"] if tag["Key"] == "Name")
					virtual_machine = tag_name.split(" - ")[1]
					_insert_snapshot(snapshot, virtual_machine, random_snapshot)
					jingrow.db.commit()
				except Exception:
					log_error(
						title="Virtual Disk Snapshot Sync Error",
						snapshot=snapshot,
					)
					jingrow.db.rollback()


def _insert_snapshot(snapshot, virtual_machine, random_snapshot):
	start_time = jingrow.utils.format_datetime(snapshot["StartTime"], "yyyy-MM-dd HH:mm:ss")
	new_snapshot = jingrow.get_pg(
		{
			"pagetype": "Virtual Disk Snapshot",
			"snapshot_id": snapshot["SnapshotId"],
			"virtual_machine": virtual_machine,
			"volume_id": snapshot["VolumeId"],
			"status": random_snapshot.get_aws_status_map(snapshot["State"]),
			"description": snapshot["Description"],
			"size": snapshot["VolumeSize"],
			"start_time": start_time,
			"progress": snapshot["Progress"],
		}
	).insert()
	jingrow.db.set_value(
		"Virtual Disk Snapshot",
		new_snapshot.name,
		{"creation": start_time, "modified": start_time},
		update_modified=False,
	)
	return new_snapshot


def _should_skip_snapshot(snapshot):
	tag_names = [tag["Value"] for tag in snapshot["Tags"] if tag["Key"] == "Name"]
	if not tag_names:
		return True
	tag_name_parts = tag_names[0].split(" - ")
	if len(tag_name_parts) != 3:
		return True
	identifier, virtual_machine, _ = tag_name_parts
	if identifier != "Jingrow":
		return True
	if not jingrow.db.exists("Virtual Machine", virtual_machine):
		return True

	return False


def delete_duplicate_snapshot_docs(snapshot):
	# Delete all except one snapshot document
	# It doesn't matter which one we keep
	snapshot_id = snapshot["SnapshotId"]
	snapshot_count = jingrow.db.count("Virtual Disk Snapshot", {"snapshot_id": snapshot_id})
	if snapshot_count > 1:
		jingrow.db.sql(
			"""
				DELETE
				FROM `tabVirtual Disk Snapshot`
				WHERE snapshot_id=%s
				LIMIT %s
			""",
			(snapshot_id, snapshot_count - 1),
		)


def _update_snapshot_if_exists(snapshot, random_snapshot):
	snapshot_id = snapshot["SnapshotId"]
	if jingrow.db.exists("Virtual Disk Snapshot", {"snapshot_id": snapshot_id}):
		jingrow.db.set_value(
			"Virtual Disk Snapshot",
			{"snapshot_id": snapshot_id},
			"status",
			random_snapshot.get_aws_status_map(snapshot["State"]),
		)
		return True
	return False
