from __future__ import annotations

import datetime
import dataclasses
import enum
import time
import logging
from typing import Literal, Self, cast
import ntnx_vmm_py_client as vmm
import ntnx_prism_py_client as prism

from nutanix_shim_server import server

logger = logging.getLogger(__name__)

try:
    from IPython.terminal.embed import embed
except ImportError:

    def embed():
        pass


class VirtualMachineMgmt:
    config: vmm.Configuration

    def __init__(self, ctx: server.Context):
        self.config = vmm.Configuration()
        self.config.host = ctx.nutanix_host
        self.config.scheme = ctx.nutanix_host_scheme
        self.config.set_api_key(ctx.nutanix_api_key)
        self.config.max_retry_attempts = 3
        self.config.backoff_factor = 3
        self.config.verify_ssl = ctx.nutanix_host_verify_ssl
        self.config.port = ctx.nutanix_host_port
        self.config.client_certificate_file = ctx.nutanix_client_certificate_file
        self.config.root_ca_certificate_file = ctx.nutanix_root_ca_certificate_file

        # Prism config for task polling
        self.prism_config = prism.Configuration()
        self.prism_config.host = ctx.nutanix_host
        self.prism_config.scheme = ctx.nutanix_host_scheme
        self.prism_config.set_api_key(ctx.nutanix_api_key)
        self.prism_config.max_retry_attempts = 3
        self.prism_config.backoff_factor = 3
        self.prism_config.verify_ssl = ctx.nutanix_host_verify_ssl
        self.prism_config.port = ctx.nutanix_host_port
        self.config.client_certificate_file = ctx.nutanix_client_certificate_file
        self.config.root_ca_certificate_file = ctx.nutanix_root_ca_certificate_file

    @property
    def client(self) -> vmm.ApiClient:
        if not hasattr(self, "_client"):
            self._client = vmm.ApiClient(self.config)
            self._client.add_default_header(
                header_name="Accept-Encoding", header_value="gzip, deflate, br"
            )
        return self._client

    @property
    def prism_client(self) -> prism.ApiClient:
        if not hasattr(self, "_prism_client"):
            self._prism_client = prism.ApiClient(self.prism_config)
            self._prism_client.add_default_header(
                header_name="Accept-Encoding", header_value="gzip, deflate, br"
            )
        return self._prism_client

    @property
    def tasks_api(self) -> prism.TasksApi:
        if not hasattr(self, "_tasks_api"):
            self._tasks_api = prism.TasksApi(self.prism_client)
        return self._tasks_api

    @property
    def images_api(self) -> vmm.ImagesApi:
        if not hasattr(self, "_images_api"):
            self._images_api = vmm.ImagesApi(self.client)
        return self._images_api

    @property
    def vms_api(self) -> vmm.VmApi:
        if not hasattr(self, "_vms_api"):
            self._vms_api = vmm.VmApi(self.client)
        return self._vms_api

    def list_images(self) -> list[ImageMetadata]:
        # TODO: paginate
        resp: vmm.ListImagesApiResponse = self.images_api.list_images(_limit=100)  # type: ignore
        images: None | list[vmm.Image] = resp.data  # type: ignore
        if images:
            return [ImageMetadata.from_nutanix_image(img) for img in images]
        return []

    def list_vms(self) -> list["VmListMetadata"]:
        """List all VMs in the Nutanix environment"""
        # TODO: paginate
        resp: vmm.ListVmsApiResponse = self.vms_api.list_vms(_limit=100)  # type: ignore
        vms: None | list[vmm.AhvConfigVm] = resp.data  # type: ignore
        if vms:
            return [VmListMetadata.from_nutanix_vm(vm) for vm in vms]
        return []

    def get_vm_details(self, vm_ext_id: str) -> "VmDetailsMetadata":
        """
        Get detailed information about a VM including MAC address and IP addresses.

        Args:
            vm_ext_id: The external ID of the VM

        Returns:
            VmDetailsMetadata with full VM details
        """
        resp: vmm.AhvConfigGetVmApiResponse = self.vms_api.get_vm_by_id(extId=vm_ext_id)  # type: ignore
        vm: vmm.AhvConfigVm = resp.data  # type: ignore
        return VmDetailsMetadata.from_nutanix_vm(vm)

    def get_vm_power_state(self, vm_ext_id: str) -> "VmPowerStateResponse":
        """
        Get the current power state of a VM.

        Args:
            vm_ext_id: The external ID of the VM

        Returns:
            VmPowerStateResponse with the current power state
        """
        resp: vmm.AhvConfigGetVmApiResponse = self.vms_api.get_vm_by_id(extId=vm_ext_id)  # type: ignore
        vm: vmm.AhvConfigVm = resp.data  # type: ignore
        power_state = str(vm.power_state) if vm.power_state else "UNDETERMINED"
        return VmPowerStateResponse(
            ext_id=vm_ext_id,
            name=cast(str, vm.name),
            power_state=power_state,
        )

    def set_vm_power_state(
        self, vm_ext_id: str, action: "PowerAction"
    ) -> "VmPowerStateResponse":
        """
        Change the power state of a VM.

        Args:
            vm_ext_id: The external ID of the VM
            action: The power action to perform

        Returns:
            VmPowerStateResponse with the new power state
        """
        # Fetch VM to get ETag (required for all power state operations)
        get_resp = self.vms_api.get_vm_by_id(extId=vm_ext_id)
        etag = self.client.get_etag(get_resp)

        # Perform the requested power action with ETag
        if action == PowerAction.POWER_ON:
            self.vms_api.power_on_vm(extId=vm_ext_id, if_match=etag)
        elif action == PowerAction.POWER_OFF:
            self.vms_api.power_off_vm(extId=vm_ext_id, if_match=etag)
        elif action == PowerAction.SHUTDOWN:
            self.vms_api.shutdown_vm(extId=vm_ext_id, if_match=etag)
        elif action == PowerAction.REBOOT:
            self.vms_api.reboot_vm(extId=vm_ext_id, if_match=etag)
        elif action == PowerAction.RESET:
            self.vms_api.reset_vm(extId=vm_ext_id, if_match=etag)
        else:
            raise ValueError(f"Unknown power action: {action}")

        # Get updated power state
        return self.get_vm_power_state(vm_ext_id)

    def delete_vm(self, vm_ext_id: str) -> None:
        """
        Delete a virtual machine.

        Args:
            vm_ext_id: The external ID of the VM to delete

        Raises:
            ApiException: If the VM cannot be deleted
        """
        # First fetch the VM to get its ETag (required for deletion)
        get_resp = self.vms_api.get_vm_by_id(extId=vm_ext_id)
        etag = self.client.get_etag(get_resp)

        # Delete with the ETag header
        self.vms_api.delete_vm_by_id(extId=vm_ext_id, if_match=etag)

    def provision_vm(self, request: "VmProvisionRequest") -> "VmMetadata":
        """
        Provision a new VM with network (not image-based), CPU, memory, and disk configuration.

        Args:
            request: VM provisioning request with all required specifications

        Returns:
            VmMetadata with information about the created VM
        """
        # Create cluster reference
        cluster_ref = vmm.AhvConfigClusterReference(ext_id=request.cluster_ext_id)

        # Create network configuration with subnet reference and DHCP
        subnet_ref = vmm.SubnetReference(ext_id=request.subnet_ext_id)
        ipv4_config = vmm.Ipv4Config(should_assign_ip=True)  # Enable DHCP
        network_info = vmm.AhvConfigNicNetworkInfo(
            subnet=subnet_ref,
            ipv4_config=ipv4_config,
        )

        # Create NIC with VIRTIO emulation
        emulated_nic = vmm.EmulatedNic(
            model=vmm.EmulatedNicModel.VIRTIO,
            is_connected=True,
        )
        nic = vmm.AhvConfigNic(
            backing_info=emulated_nic,
            network_info=network_info,
        )

        # Create disk configuration (empty disk on storage container)
        storage_container = vmm.AhvConfigVmDiskContainerReference(
            ext_id=request.storage_container_ext_id
        )
        vm_disk = vmm.AhvConfigVmDisk(
            disk_size_bytes=request.disk_size_bytes,
            storage_container=storage_container,
        )
        disk_address = vmm.AhvConfigDiskAddress(
            bus_type=vmm.AhvConfigDiskBusType.SCSI,
            index=0,
        )
        disk = vmm.AhvConfigDisk(
            backing_info=vm_disk,
            disk_address=disk_address,
        )

        if request.boot_method == "bios":
            boot_config = vmm.LegacyBoot()
        else:
            boot_config = vmm.UefiBoot(is_secure_boot_enabled=request.secure_boot)

        # Create VM specification
        vm_spec = vmm.AhvConfigVm(
            name=request.name,
            description=request.description,
            cluster=cluster_ref,
            num_sockets=request.num_sockets,
            num_cores_per_socket=request.num_cores_per_socket,
            memory_size_bytes=request.memory_size_bytes,
            nics=[nic],
            disks=[disk],
            boot_config=boot_config,
        )

        # Create the VM - this returns a task reference, not the VM directly
        resp: vmm.CreateVmApiResponse = self.vms_api.create_vm(body=vm_spec)  # type: ignore

        # Get task ID - keep the full format including prefix for the tasks API
        task_ext_id = cast(str, resp.data.ext_id)  # type: ignore

        logger.info(f"Waiting for VM creation task {task_ext_id} to complete...")

        # Poll for task completion
        max_wait_seconds = 120
        poll_interval = 2
        elapsed = 0

        while elapsed < max_wait_seconds:
            task_resp = self.tasks_api.get_task_by_id(extId=task_ext_id)
            task = task_resp.data  # type: ignore
            status = str(task.status) if task.status else "UNKNOWN"

            logger.info(f"Task status: {status}")

            if status == "SUCCEEDED":
                # Extract VM ext_id from entities_affected
                if task.entities_affected:
                    for entity in task.entities_affected:
                        # The VM entity will have the ext_id we need
                        vm_ext_id = entity.ext_id
                        logger.info(f"VM created successfully with ext_id: {vm_ext_id}")

                        # Handle power-on if requested
                        if request.power_on:
                            try:
                                logger.info(
                                    f"Power-on requested, powering on VM {vm_ext_id}..."
                                )
                                # Fetch VM to get ETag (required for power-on operation)
                                get_resp = self.vms_api.get_vm_by_id(extId=vm_ext_id)
                                etag = self.client.get_etag(get_resp)
                                self.vms_api.power_on_vm(extId=vm_ext_id, if_match=etag)

                                # Verify VM actually powered on
                                logger.info(f"Verifying VM power state...")
                                power_on_verified = False
                                power_check_timeout = 60  # 60 seconds to power on
                                power_check_interval = 3
                                power_elapsed = 0

                                while power_elapsed < power_check_timeout:
                                    try:
                                        power_resp = self.vms_api.get_vm_by_id(
                                            extId=vm_ext_id
                                        )
                                        vm_data = power_resp.data  # type: ignore
                                        current_power_state = (
                                            str(vm_data.power_state)
                                            if vm_data.power_state
                                            else "UNKNOWN"
                                        )
                                        logger.info(
                                            f"Current power state: {current_power_state}"
                                        )

                                        if current_power_state == "ON":
                                            logger.info(
                                                f"VM {vm_ext_id} successfully powered on"
                                            )
                                            power_on_verified = True
                                            break

                                        time.sleep(power_check_interval)
                                        power_elapsed += power_check_interval
                                    except Exception as check_error:
                                        logger.warning(
                                            f"Error checking power state: {check_error}"
                                        )
                                        time.sleep(power_check_interval)
                                        power_elapsed += power_check_interval

                                if not power_on_verified:
                                    # Power-on failed, rollback by deleting the VM
                                    error_msg = f"VM created but failed to power on within {power_check_timeout} seconds"
                                    logger.error(
                                        f"{error_msg}, rolling back by deleting VM {vm_ext_id}"
                                    )

                                    try:
                                        # Delete the VM
                                        get_resp = self.vms_api.get_vm_by_id(
                                            extId=vm_ext_id
                                        )
                                        etag = self.client.get_etag(get_resp)
                                        self.vms_api.delete_vm_by_id(
                                            extId=vm_ext_id, if_match=etag
                                        )
                                        logger.info(
                                            f"VM {vm_ext_id} deleted successfully during rollback"
                                        )
                                    except Exception as delete_error:
                                        logger.error(
                                            f"Failed to delete VM during rollback: {delete_error}"
                                        )
                                        error_msg += f". Additionally, failed to delete VM: {delete_error}"

                                    raise RuntimeError(error_msg)

                            except RuntimeError:
                                # Re-raise RuntimeError from power-on verification failure
                                raise
                            except Exception as power_error:
                                # Power-on command itself failed, rollback
                                error_msg = f"Failed to power on VM: {power_error}"
                                logger.error(
                                    f"{error_msg}, rolling back by deleting VM {vm_ext_id}"
                                )

                                try:
                                    # Delete the VM
                                    get_resp = self.vms_api.get_vm_by_id(
                                        extId=vm_ext_id
                                    )
                                    etag = self.client.get_etag(get_resp)
                                    self.vms_api.delete_vm_by_id(
                                        extId=vm_ext_id, if_match=etag
                                    )
                                    logger.info(
                                        f"VM {vm_ext_id} deleted successfully during rollback"
                                    )
                                except Exception as delete_error:
                                    logger.error(
                                        f"Failed to delete VM during rollback: {delete_error}"
                                    )
                                    error_msg += f". Additionally, failed to delete VM: {delete_error}"

                                raise RuntimeError(error_msg)
                        else:
                            logger.info(f"Power-on not requested, VM will remain off")

                        return VmMetadata(
                            ext_id=vm_ext_id,
                            name=request.name,
                            description=request.description,
                            num_sockets=request.num_sockets,
                            num_cores_per_socket=request.num_cores_per_socket,
                            memory_size_bytes=request.memory_size_bytes,
                            disk_size_bytes=request.disk_size_bytes,
                        )

                raise ValueError("Task succeeded but no VM entity found in response")

            elif status == "FAILED":
                error_msg = (
                    task.error_messages
                    if hasattr(task, "error_messages")
                    else "Unknown error"
                )
                raise ValueError(f"VM creation task failed: {error_msg}")

            time.sleep(poll_interval)
            elapsed += poll_interval

        raise TimeoutError(
            f"VM creation task did not complete within {max_wait_seconds} seconds"
        )


@dataclasses.dataclass(frozen=True)
class ImageMetadata:
    name: str
    description: None | str
    create_time: datetime.datetime
    last_update_time: datetime.datetime
    ext_id: str
    cluster_location_ext_ids: list[str]
    source: None | str
    placement_policy_status: None | str
    owner_ext_id: str
    tenant_id: None | str
    size_bytes: int
    type: str

    keys = __annotations__

    @classmethod
    def from_nutanix_image(cls, image: vmm.Image) -> Self:
        kwargs = {k: v for k, v in image.to_dict().items() if k in cls.keys}
        return cls(**kwargs)


@dataclasses.dataclass(frozen=True)
class VmListMetadata:
    """
    Metadata for listing VMs.

    Contains essential information about VMs for display in Foreman.
    """

    ext_id: str
    name: str
    cluster_ext_id: None | str
    power_state: None | str
    num_sockets: None | int
    num_cores_per_socket: None | int
    memory_size_bytes: None | int
    mac_address: None | str
    ip_addresses: list[str]
    create_time: None | datetime.datetime

    @classmethod
    def from_nutanix_vm(cls, vm: vmm.AhvConfigVm) -> Self:
        """Convert Nutanix SDK VM to our response model"""
        # Extract cluster ext_id from cluster reference
        cluster_ext_id = None
        if hasattr(vm, "cluster") and vm.cluster:
            cluster_ext_id = (
                vm.cluster.ext_id if hasattr(vm.cluster, "ext_id") else None
            )

        # Convert power state enum to string
        power_state = str(vm.power_state) if vm.power_state else None

        # Extract MAC address from first NIC
        mac_address = None
        ip_addresses = []
        if hasattr(vm, "nics") and vm.nics:
            first_nic = vm.nics[0]
            if hasattr(first_nic, "backing_info") and first_nic.backing_info:
                mac_address = getattr(first_nic.backing_info, "mac_address", None)
            # Extract IP addresses from NIC network info
            if hasattr(first_nic, "network_info") and first_nic.network_info:
                ipv4_config = getattr(first_nic.network_info, "ipv4_config", None)
                if ipv4_config and hasattr(ipv4_config, "ip_address"):
                    ip_addr = ipv4_config.ip_address
                    if ip_addr and hasattr(ip_addr, "value"):
                        ip_addresses.append(ip_addr.value)

        return cls(
            ext_id=cast(str, vm.ext_id),
            name=cast(str, vm.name),
            cluster_ext_id=cluster_ext_id,
            power_state=power_state,
            num_sockets=vm.num_sockets,
            num_cores_per_socket=vm.num_cores_per_socket,
            memory_size_bytes=vm.memory_size_bytes,
            mac_address=mac_address,
            ip_addresses=ip_addresses,
            create_time=vm.create_time if hasattr(vm, "create_time") else None,
        )


@dataclasses.dataclass(frozen=True)
class VmDetailsMetadata:
    """
    Detailed metadata for a single VM.

    Contains all information needed for Foreman host details view.
    """

    ext_id: str
    name: str
    cluster_ext_id: None | str
    power_state: None | str
    num_sockets: None | int
    num_cores_per_socket: None | int
    memory_size_bytes: None | int
    mac_address: None | str
    ip_addresses: list[str]
    create_time: None | datetime.datetime

    @classmethod
    def from_nutanix_vm(cls, vm: vmm.AhvConfigVm) -> Self:
        """Convert Nutanix SDK VM to our detailed response model"""
        # Extract cluster ext_id from cluster reference
        cluster_ext_id = None
        if hasattr(vm, "cluster") and vm.cluster:
            cluster_ext_id = (
                vm.cluster.ext_id if hasattr(vm.cluster, "ext_id") else None
            )

        # Convert power state enum to string
        power_state = str(vm.power_state) if vm.power_state else None

        # Extract MAC address from first NIC
        mac_address = None
        ip_addresses = []
        if hasattr(vm, "nics") and vm.nics:
            first_nic = vm.nics[0]
            if hasattr(first_nic, "backing_info") and first_nic.backing_info:
                mac_address = getattr(first_nic.backing_info, "mac_address", None)
            # Extract IP addresses from NIC network info
            if hasattr(first_nic, "network_info") and first_nic.network_info:
                ipv4_config = getattr(first_nic.network_info, "ipv4_config", None)
                if ipv4_config and hasattr(ipv4_config, "ip_address"):
                    ip_addr = ipv4_config.ip_address
                    if ip_addr and hasattr(ip_addr, "value"):
                        ip_addresses.append(ip_addr.value)

        return cls(
            ext_id=cast(str, vm.ext_id),
            name=cast(str, vm.name),
            cluster_ext_id=cluster_ext_id,
            power_state=power_state,
            num_sockets=vm.num_sockets,
            num_cores_per_socket=vm.num_cores_per_socket,
            memory_size_bytes=vm.memory_size_bytes,
            mac_address=mac_address,
            ip_addresses=ip_addresses,
            create_time=vm.create_time if hasattr(vm, "create_time") else None,
        )


@dataclasses.dataclass
class VmProvisionRequest:
    """
    Request model for provisioning a new VM.

    Example:
        {
            "name": "my-vm-01",
            "description": "Development VM for testing",
            "cluster_ext_id": "00061663-9fa0-28ca-185b-ac1f6b6f97e2",
            "subnet_ext_id": "3d5d8e8b-f3e0-4f4e-8c5d-5b5c5d5e5f5a",
            "storage_container_ext_id": "1a2b3c4d-5e6f-7a8b-9c0d-1e2f3a4b5c6d",
            "num_sockets": 2,
            "num_cores_per_socket": 2,
            "memory_size_bytes": 8589934592,  # 8 GB
            "disk_size_bytes": 107374182400,   # 100 GB
            "power_on": true,                  # Auto power-on after creation
            "boot_method": "uefi",             # uefi | bios
            "secure_boot": true,               # secure boot conf - applicable only to UEFI
        }
    """

    name: str
    cluster_ext_id: str
    subnet_ext_id: str
    storage_container_ext_id: str
    num_sockets: int
    num_cores_per_socket: int
    memory_size_bytes: int
    disk_size_bytes: int
    description: str = ""
    power_on: bool = True
    boot_method: Literal["bios", "uefi"] = "uefi"
    secure_boot: bool = False


@dataclasses.dataclass(frozen=True)
class VmMetadata:
    """
    Response model containing information about a provisioned VM.
    """

    ext_id: str
    name: str
    description: str
    num_sockets: int
    num_cores_per_socket: int
    memory_size_bytes: int
    disk_size_bytes: int


class PowerAction(str, enum.Enum):
    """
    Available power actions for VMs.

    - POWER_ON: Turn on the VM (hard power on)
    - POWER_OFF: Turn off the VM (hard power off)
    - SHUTDOWN: Gracefully shut down the VM
    - REBOOT: Reboot the VM (hard reboot)
    - RESET: Reset the VM (hard reset/power cycle)
    """

    POWER_ON = "POWER_ON"
    POWER_OFF = "POWER_OFF"
    SHUTDOWN = "SHUTDOWN"
    REBOOT = "REBOOT"
    RESET = "RESET"


@dataclasses.dataclass
class PowerStateChangeRequest:
    """
    Request model for changing VM power state.

    Example:
        {
            "action": "POWER_ON"
        }
    """

    action: PowerAction


@dataclasses.dataclass(frozen=True)
class VmPowerStateResponse:
    """
    Response model containing VM power state information.

    Power states:
    - ON: VM is powered on
    - OFF: VM is powered off
    - PAUSED: VM is paused
    - UNDETERMINED: Power state cannot be determined
    """

    ext_id: str
    name: str
    power_state: str
