from __future__ import annotations

import os
import typing

import click
import jingrow

if typing.TYPE_CHECKING:
	from jcloude.infrastructure.pagetype.arm_build_record.arm_build_record import ARMBuildRecord
	from jcloude.infrastructure.pagetype.arm_docker_image.arm_docker_image import ARMDockerImage
	from jcloude.infrastructure.pagetype.virtual_machine_migration.virtual_machine_migration import (
		VirtualMachineMigration,
	)
	from jcloude.jcloude.pagetype.server.server import Server


arm_machine_mappings = {
	"t2": "t4g",
	"c6i": "c8g",
	"m6i": "m8g",
	"m7i": "m8g",
	"r6i": "r8g",
	# Following are for Zurich due to lack of newer processors in that region
	"r5": "r7g",
	"m5": "m7g",
	"c5": "c7g",
}

amd_machine_mappings = {"r6i": "m6a", "m6i": "m6a", "c6i": "m6a", "m5": "m6a", "r7i": "m6a", "m7i": "m6a"}


def has_arm_build_record(server: str) -> bool:
	return bool(jingrow.get_value("ARM Build Record", {"server": server}))


def check_image_build_failure(arm_build_record: ARMBuildRecord) -> bool:
	return any(arm_image.status != "Success" for arm_image in arm_build_record.arm_images)


def create_vmm(server: str, virtual_machine_image: str, target_machine_type: str) -> VirtualMachineMigration:
	virtual_machine_migration: VirtualMachineMigration = jingrow.get_pg(
		{
			"pagetype": "Virtual Machine Migration",
			"virtual_machine_image": virtual_machine_image,
			"machine_type": target_machine_type,
			"virtual_machine": server,
		}
	)
	return virtual_machine_migration.insert()


def vmm(server, vmi, amd_conversion: bool = False) -> VirtualMachineMigration:
	machine_type = jingrow.db.get_value("Virtual Machine", {"name": server}, "machine_type")
	machine_series, machine_size = machine_type.split(".")

	machine_mappings = arm_machine_mappings if not amd_conversion else amd_machine_mappings
	if amd_conversion and ("r6i" in machine_series or "r7i" in machine_series):
		if machine_size == "xlarge":
			machine_size = "2xlarge"
		else:
			machine_size = machine_size.replace("2", "4")

	virtual_machine_migration: VirtualMachineMigration = create_vmm(
		server=server,
		virtual_machine_image=vmi,
		target_machine_type=f"{machine_mappings[machine_series]}.{machine_size}",
	)
	return virtual_machine_migration


def connect(bench_dir, site_dir):
	sites_dir = os.path.join(bench_dir, "sites")
	jingrow.init(site=site_dir, sites_path=sites_dir)
	jingrow.connect()


def load_servers_from_file(file_path: str) -> list[str]:
	with open(file_path) as server_file:
		return server_file.read().strip().split("\n")


@click.group()
@click.option("--site", "site_name", required=True, help="Jingrow site name")
def cli(site_name):
	"""CLI entry point."""
	bench_dir = os.path.dirname(__file__).split("apps")[0]
	site_dir = os.path.join(bench_dir, "sites", site_name)
	connect(bench_dir, site_dir)


@cli.command()
@click.option(
	"--server-file", type=click.Path(exists=True), help="Path to a file containing a list of servers."
)
@click.argument("servers", nargs=-1, type=str)
def trigger_arm_build(servers: list[str], server_file: str):
	"""Trigger ARM build for one or more servers."""
	if server_file:
		servers = load_servers_from_file(server_file)

	for server in servers:
		if has_arm_build_record(server):
			continue

		server: Server = jingrow.get_pg("Server", server)
		server.collect_arm_images()
		jingrow.db.commit()


@cli.command()
@click.option(
	"--server-file", type=click.Path(exists=True), help="Path to a file containing a list of servers."
)
@click.argument("servers", nargs=-1, type=str)
def pull_images_on_servers(servers: list[str], server_file: str):
	"""Trigger image pulls on Intel server to be converted"""
	if server_file:
		servers = load_servers_from_file(server_file)

	for server in servers:
		arm_build_record: ARMBuildRecord = jingrow.get_pg("ARM Build Record", {"server": server})

		try:
			arm_build_record.pull_images()
			print(f"Pulled image on {server}")
		except jingrow.ValidationError:
			print(f"Skipping server {server} due to failed builds")

		jingrow.db.commit()


@cli.command()
@click.option("--vmi", default="f377-mumbai.jingrow.cloud")
@click.option("--vmi-cluster", required=True)
@click.option(
	"--server-file", type=click.Path(exists=True), help="Path to a file containing a list of servers."
)
@click.argument("servers", nargs=-1, type=str)
def update_image_and_create_migration(
	vmi: str,
	vmi_cluster: str,
	servers: list[str],
	server_file: str,
):
	"""Update docker image on bench config and create virtual machine migration"""
	vmi = jingrow.get_value("Virtual Machine Image", {"virtual_machine": vmi, "cluster": vmi_cluster}, "name")
	if not vmi:
		print(f"Aborting VMI not found {vmi}!")
		return

	if server_file:
		servers = load_servers_from_file(server_file)

	for server in servers:
		arm_build_record: ARMBuildRecord = jingrow.get_pg("ARM Build Record", {"server": server})
		try:
			arm_build_record.update_image_tags_on_benches()
			virtual_machine_migration = vmm(server, vmi)
			jingrow.db.commit()
			print(f"Created {virtual_machine_migration.name}")
		except jingrow.ValidationError as e:
			print(f"Aborting: {e}!")
			break


@cli.command()
@click.option("--vmi", default="m263-mumbai.jingrow.cloud")
@click.option("--vmi-cluster", required=True)
@click.option(
	"--server-file",
	type=click.Path(exists=True),
	help="Path to a file containing a list of servers.",
)
@click.option("--start", type=bool, default=False)
@click.argument("servers", nargs=-1, type=str)
def convert_database_servers(
	vmi: str, vmi_cluster: str, servers: list[str], server_file: str, start: bool = False
):
	vmi = jingrow.get_value("Virtual Machine Image", {"virtual_machine": vmi, "cluster": vmi_cluster}, "name")
	if not vmi:
		print(f"Aborting VMI not found {vmi}!")
		return

	if server_file:
		servers = load_servers_from_file(server_file)

	for server in servers:
		virtual_machine_migration = vmm(server, vmi, amd_conversion=True)
		jingrow.db.commit()
		print(f"Created {virtual_machine_migration.name}")

	if start:
		for server in servers:
			virtual_machine_migration: VirtualMachineMigration = jingrow.get_pg(
				"Virtual Machine Migration", {"virtual_machine": server}
			)
			virtual_machine_migration.execute()
			jingrow.db.commit()


@cli.command()
@click.option(
	"--server-file", type=click.Path(exists=True), help="Path to a file containing a list of servers."
)
@click.argument("servers", nargs=-1, type=str)
def arm_build_info(servers: list[str], server_file: str):
	total, successful, failed, running = 0, 0, 0, 0
	if server_file:
		servers = load_servers_from_file(server_file)

	def _status_info(images: list[ARMDockerImage], status: str):
		return len([image for image in images if image.status == status])

	for server in servers:
		arm_build_record: ARMBuildRecord = jingrow.get_pg("ARM Build Record", {"server": server})
		arm_build_record.sync_status()
		total += len(arm_build_record.arm_images)
		running += _status_info(arm_build_record.arm_images, "Running")
		successful += _status_info(arm_build_record.arm_images, "Success")
		failed += _status_info(arm_build_record.arm_images, "Failure")

	print(f"Total: {total}\nSuccessful: {successful}\nRunning: {running}\nFailed: {failed}")


@cli.command()
@click.option(
	"--server-file", type=click.Path(exists=True), help="Path to a file containing a list of servers."
)
@click.option("--vmi", default="f436-mumbai.jingrow.cloud")
@click.option("--vmi-cluster", required=True)
@click.argument("servers", nargs=-1, type=str)
def convert_to_amd(servers: list[str], vmi: str, server_file: str, vmi_cluster: str):
	"""Update docker image on bench config and create virtual machine migration"""
	vmi = jingrow.get_value("Virtual Machine Image", {"virtual_machine": vmi, "cluster": vmi_cluster}, "name")
	if not vmi:
		print(f"Aborting VMI not found {vmi}!")
		return

	if server_file:
		servers = load_servers_from_file(server_file)

	for server in servers:
		try:
			virtual_machine_migration = vmm(server, vmi, amd_conversion=True)
			jingrow.db.commit()
			print(f"Created {virtual_machine_migration.name}")
		except jingrow.ValidationError as e:
			print(f"Aborting: {e}!")
			break


@cli.command()
@click.option(
	"--server-file", type=click.Path(exists=True), help="Path to a file containing a list of servers."
)
@click.argument("servers", nargs=-1, type=str)
def database_post_migration_steps(servers: list[str], server_file: str):
	"""Not a part of the migration script since"""
	if server_file:
		servers = load_servers_from_file(server_file)

	for server in servers:
		server = jingrow.get_pg("Database Server", server)
		server.set_swappiness()
		server.add_glass_file()
		server.install_filebeat()
		server.adjust_memory_config()
		server.setup_logrotate()
		server.save()


@cli.result_callback()
def cleanup(*args, **kwargs):
	jingrow.destroy()


def main():
	cli(obj={})


if __name__ == "__main__":
	main()