"""
Scaleway instances provider for Enoslib
"""

import time
import subprocess
from typing import Iterable, Optional, Tuple

import enoslib as en
from enoslib import sync_info
from enoslib.infra.provider import Provider
from enoslib.objects import Host, Networks, Roles
from enoslib.log import getLogger

from scaleway import Client, Profile, WaitForOptions, ScalewayException
from scaleway.instance.v1 import (
    Server,
    ServerAction,
    ServerState,
)
from scaleway.instance.v1.custom_api import InstanceUtilsV1API
from scaleway.instance.v1.types import (
    IpType,
    Ip,
    VolumeServerTemplate,
    VolumeVolumeType,
)
from scaleway.vpc.v2.api import PrivateNetwork, VpcV2API
from scaleway.vpcgw.v2.api import Gateway, GatewayNetwork, PatRuleProtocol, VpcgwV2API
from scaleway.block.v1.api import BlockV1API
from scaleway.block.v1.types import VolumeStatus

from .configuration import Configuration, Architecture
from .constants import (
    DEFAULT_GATEWAY_PORT,
    DEFAULT_GATEWAY_USER,
    DEFAULT_VOLUMES,
    DEFAULT_GPU_VOLUMES,
)
from .error import EnosScalewayError

logger = getLogger(__name__, ["ScwInstances"])

MAX_BLOCK_POLL = 20

STATE_POLLING_INTERVAL = 5  # seconds

SERVER_ID = "server_id"


def _get_client(
    file_path: Optional[str] = None, profile: Optional[str] = None
) -> Client:
    """
    Create and return a Scaleway client using the specified profile.
    If no profile is specified, the default profile is used.
    """
    profile = "default" if profile is None else profile
    try:
        p = Profile.from_config_file(filepath=file_path, profile_name=profile)
        logger.info(f'Loaded scaleway profile: "{profile}" from config file')
        return Client.from_profile(p)
    except ValueError as e:
        logger.debug(f"Error loading profile from config file: {e}")
        logger.debug("Error loading profile from config file")
    try:
        logger.debug("Loading profile from environment variables")
        p = Profile.from_env()
        logger.info("Loaded scaleway profile from env")
        client = Client.from_profile(p)
        return client
    except ValueError as e:
        logger.error(f"Error loading profile from env: {e}")
        raise e


def _get_instances_api(client: Client) -> InstanceUtilsV1API:
    """
    Create and return the InstanceUtilsV1API for the given Scaleway client.
    """
    try:
        api = InstanceUtilsV1API(client=client)
    except ValueError as e:
        logger.error(f"Error loading instances api: {e}")
        raise EnosScalewayError(e)
    return api


def _get_vpc_api(client: Client) -> VpcV2API:
    try:
        api = VpcV2API(client=client)
    except ValueError as e:
        logger.error(f"Error loading VPC api: {e}")
        raise EnosScalewayError(e)
    return api


def _get_vpc_gw_api(client: Client) -> VpcgwV2API:
    try:
        api = VpcgwV2API(client=client)
    except ValueError as e:
        logger.error(f"Error loading VPC api: {e}")
        raise EnosScalewayError(e)
    return api


def _get_block_api(client: Client) -> BlockV1API:
    try:
        api = BlockV1API(client=client)
    except ValueError as e:
        logger.error(f"Error loading Block api: {e}")
        raise EnosScalewayError(e)
    return api


def _create_private_network(
    client: Client, net_name: Optional[str] = None, region: Optional[str] = None
) -> PrivateNetwork:
    """Create a private network

    Args:
        client (Client): Scaleway client
        net_name (str, optional): private network name.
        region (str, optional): private network region.

    Returns:
        PrivateNetwork: private network object
    """
    vpc_api = _get_vpc_api(client)
    l_existing_net = vpc_api.list_private_networks(region=region)
    if net_name not in [net.name for net in l_existing_net.private_networks]:
        net = vpc_api.create_private_network(
            region=region, name=net_name, default_route_propagation_enabled=True
        )
        logger.info(f"Created private network {net.name} with ID {net.id}")
    else:
        net = [net for net in l_existing_net.private_networks if net.name == net_name][
            0
        ]
        logger.info(f'Reusing existing private network "{net.name}" with ID {net.id}')
    return net


def _create_public_gateway(
    client: Client,
    name: str,
    type: str,
    zone: Optional[str] = None,
    enable_bastion: bool = True,
    enable_smtp: bool = False,
) -> Gateway:
    vpcgw_api = _get_vpc_gw_api(client=client)
    l_existing_gw = vpcgw_api.list_gateways(zone=zone)
    if l_existing_gw.total_count == 0 or name not in [
        gw.name for gw in l_existing_gw.gateways
    ]:
        gw = vpcgw_api.create_gateway(
            type_=type,
            enable_bastion=enable_bastion,
            enable_smtp=enable_smtp,
            zone=zone,
            name=name,
        )
        logger.info(f'Created gateway "{gw.name}" with ID {gw.id}')
    else:
        gw = [gw for gw in l_existing_gw.gateways if gw.name == name][0]
        logger.info(f"Reusing existing gateway {gw.name} with ID {gw.id}")

    gw = vpcgw_api.wait_for_gateway(gateway_id=gw.id, options=WaitForOptions())

    return gw


def _create_gateway_network(
    client: Client,
    private_network: PrivateNetwork,
    gateway: Gateway,
    enable_masquerade: bool = True,
    push_default_route: bool = True,
) -> GatewayNetwork:
    vpcgw_api = _get_vpc_gw_api(client=client)
    l_gw_net = vpcgw_api.list_gateway_networks()

    if l_gw_net.total_count == 0 or (gateway.id, private_network.id) not in [
        (gw_net.gateway_id, gw_net.private_network_id)
        for gw_net in l_gw_net.gateway_networks
    ]:
        gw_net = vpcgw_api.create_gateway_network(
            gateway_id=gateway.id,
            private_network_id=private_network.id,
            enable_masquerade=enable_masquerade,
            push_default_route=push_default_route,
        )
        logger.info(f"Created gateway network with ID {gw_net.id}")
    else:
        gw_net = [
            gw_net
            for gw_net in l_gw_net.gateway_networks
            if (gw_net.gateway_id, gw_net.private_network_id)
            == (gateway.id, private_network.id)
        ][0]
        logger.info(f"Reusing existing gateway network with ID {gw_net.id}")

    gw_net = vpcgw_api.wait_for_gateway_network(gateway_network_id=gw_net.id)

    return gw_net


def _attach_servers_to_network(
    client: Client,
    private_network: PrivateNetwork,
    servers: Iterable[Server],
):
    instance_api = _get_instances_api(client=client)
    for server in servers:
        try:
            # res = instance_api.create_private_nic(
            instance_api.create_private_nic(
                server_id=server.id,
                private_network_id=private_network.id,
                zone=server.zone,
            )
            logger.debug(
                f'Attached inst "{server.name}" '
                f'to private network "{private_network.name}"'
            )
        # TODO: improve error handling
        except Exception as e:
            logger.warning(
                f'Failed attaching server "{server.name}" '
                f'to network "{private_network.name}": {e}'
            )


def _open_ports(
    provider_conf: Configuration,
    client: Client,
    roles: Roles,
    gateway: Gateway,
    gateway_network: GatewayNetwork,
):
    vpc_gw_api = _get_vpc_gw_api(client=client)
    instance_api = _get_instances_api(client=client)

    for pconf in provider_conf.get_ports_configurations():
        private_ip = None
        if len(roles[pconf.role]) > 1:
            logger.warning(
                f"Can't open port for role {pconf.role} with multiple hosts in the role"
            )
            continue
        elif len(roles[pconf.role]) == 0:
            logger.warning(f"No host found for role {pconf.role} to open port")
            continue
        host: Host = roles[pconf.role][0]
        server_id = host.extra[SERVER_ID]
        private_nic_resp = instance_api.list_private_ni_cs(server_id=server_id)
        if int(private_nic_resp.total_count) < 1:
            logger.error(
                f"Instance in role {pconf.role} is not connected to any private network"
            )
            continue
        mac_addr = private_nic_resp.private_nics[0].mac_address
        # Hacky way to fetch private address
        # API does not seem to pick it up
        cmd = (
            "ip addr show "
            f"| awk -v mac={mac_addr} '"
            "    tolower($0) ~ mac {found=1}"
            '    found && $1=="inet" {print $2; exit}\''
            "| cut -d/ -f1"
        )
        ret = en.run(cmd, roles=host, task_name="Fetching private address")
        if ret[0].status == "OK":
            private_ip = ret[0].stdout
        if private_ip is None:
            logger.error("Failed retrieving private IP for ports opening")
        l_pat_rules = vpc_gw_api.list_pat_rules(gateway_ids=[gateway.id])
        if (pconf.public_port, pconf.private_port, private_ip) in [
            (pr.public_port, pr.private_port, pr.private_ip)
            for pr in l_pat_rules.pat_rules
        ]:
            logger.info(
                f"PAT rule already exists for public port {pconf.public_port}, "
                f"private port {pconf.private_port}, role {pconf.role} with private IP {private_ip}"
            )
            continue
        protocol = pconf.protocol
        if protocol is not None:
            protocol = PatRuleProtocol(protocol)
        # TODO: add zone support
        try:
            vpc_gw_api.create_pat_rule(
                gateway_id=gateway.id,
                public_port=pconf.public_port,
                private_port=pconf.private_port,
                private_ip=private_ip,
                protocol=protocol,
            )
        except ScalewayException as e:
            logger.error(e)
        logger.info(
            f"Opened port {pconf.public_port} on gateway for role {pconf.role}, "
            f'access at "{gateway.ipv4.address}:{pconf.public_port}"'
        )


def _terminate_server(instances_api: InstanceUtilsV1API, server: Server):
    """Terminates server

    Args:
        instances_api (InstanceUtilsV1API): scaleway instances API
        server (Server): Scaleway server
    """
    try:
        # Send TERMINATE action prefered over "delete_server"
        #  as it also destroys storage
        instances_api.server_action(
            server_id=server.id,
            action=ServerAction.TERMINATE,
            zone=server.zone,
        )
    except Exception as e:
        logger.error(f"Error terminating server {server.name}: {e}")
        raise e


def _delete_server(instances_api: InstanceUtilsV1API, server: Server):
    """
    Deletes undeployed servers
    """
    try:
        instances_api.delete_server(server_id=server.id)
        for volume in server.volumes.values():
            instances_api.delete_volume(volume_id=volume.id)
    except Exception as e:
        logger.error(f"Error deleting server {server.name}: {e}")
        raise e


def _clean_server_block_storage(block_api: BlockV1API, server: Server):
    """Removing block storage, automatically used with GPU and not removed
    by terminating the instance

    Args:
        block_api (BlockV1API): block storage api
        server (Server): server object
    """
    for storage in server.volumes.values():
        try:
            ret = block_api.get_volume(volume_id=storage.id)
        except ScalewayException:
            # Not found
            continue
        counter = 0
        block_freed = True
        # fetcher = block_api.get_volume(volume_id=storage.id)
        # wait_for_resource()
        # Polling status until block storage is not in use
        # TODO: Check documentation to see if there is a better way to do it
        while ret.status != VolumeStatus.AVAILABLE:
            time.sleep(1)
            ret = block_api.get_volume(volume_id=storage.id)
            counter += 1
            if counter == MAX_BLOCK_POLL:
                block_freed = False
                break
        if block_freed:
            block_api.delete_volume(volume_id=storage.id)
            logger.info("Cleaned block storage")
        else:
            logger.error(
                f'Failed deleteing block storage "{storage.id}", timed out or storage is still in use'
            )


def _create_server(
    client: Client, provider_conf: Configuration, architecture=Architecture
) -> list[Server]:
    instances_api = _get_instances_api(client)

    servers = []
    total = 0

    for machine_conf in provider_conf.machines:
        zone = machine_conf.zone
        # Get current deployed instances:
        existing_servers = instances_api.list_servers(zone=zone).servers

        existing_servers_collection = {s.name: s for s in existing_servers}
        for _ in range(machine_conf.number):
            # Giving different names to servers
            server_name = f"{machine_conf.name}-{total}"

            # TODO: improve retrieval of existing servers
            if server_name in existing_servers_collection:
                logger.info(
                    f"Server {server_name} found in existing servers. Reusing it."
                )
                instance = existing_servers_collection[server_name]
            else:
                logger.debug(f"Server {server_name} not found. Creating a new one.")
                # Create server
                kwargs = machine_conf.to_instance_kwargs()
                kwargs["name"] = server_name
                kwargs["tags"] = machine_conf.roles
                kwargs["volumes"] = __configure_volumes(
                    machine_conf.volumes, "gpu" in kwargs["image"]
                )
                # Request public IP for instance
                kwargs["dynamic_ip_required"] = provider_conf.dynamic_ip_required
                if (
                    architecture == Architecture.Standalone
                    and not provider_conf.dynamic_ip_required
                ):
                    logger.info(
                        f"{server_name}: when deploying instances standalone (no private network), a dynamic IP is required"
                    )
                    kwargs["dynamic_ip_required"] = True

                    # This option does not seem to work
                    # kwargs["enable_ipv6"] = True
                try:
                    instance = instances_api.create_server(**kwargs).server
                except Exception as e:
                    logger.error(f"Error creating server {server_name}: {e}")
                    raise e
            total += 1
            servers.append(instance)
    return servers


def __configure_volumes(
    volumes_conf: Optional[list[dict]],
    is_gpu: bool = False,
):
    volumes = {}
    if is_gpu:
        default_volumes = DEFAULT_GPU_VOLUMES
    else:
        default_volumes = DEFAULT_VOLUMES
    volumes = default_volumes
    default_values = default_volumes["0"]
    if volumes_conf is None:
        return volumes
    for i, conf in enumerate(volumes_conf):
        key = str(i)
        volume_type = conf.get("volume_type", default_values.volume_type.value)
        volumes[key] = VolumeServerTemplate(
            volume_type=VolumeVolumeType(volume_type),
            size=conf.get("size", default_values.size),
            boot=conf.get("boot", default_values.boot),
        )
    return volumes


def start_servers(client: Client, servers: list[Server]):
    instances_api = _get_instances_api(client=client)
    total = len(servers)

    _poweron_server(servers, instances_api)

    # Waiting for runnig state
    # Could also use :
    # instances_api.wait_instance_server()
    while True:
        deployed = []
        undeployed = []
        for server in servers:
            ret = instances_api.get_server(server_id=server.id, zone=server.zone)
            if ret.server.state == ServerState.RUNNING:
                deployed.append(ret.server)
            if ret.server.state in (ServerState.LOCKED, ServerState.STOPPED):
                undeployed.append(ret.server)
        logger.info("Polling deployment status...")
        logger.info("%s/%s deployed servers", len(deployed), total)
        logger.info("%s/%s blocked servers", len(undeployed), total)
        if len(deployed) + len(undeployed) >= total:
            break
        time.sleep(STATE_POLLING_INTERVAL)
    return deployed, undeployed


def _poweron_server(
    servers: list[Server],
    instances_api: InstanceUtilsV1API,
):
    """Send POWERON action to all possible servers

    Args:
        servers (list[Server]): Scaleway server obj
        instances_api (InstanceUtilsV1API): Instances api
    """
    for server in servers:
        if ServerAction.POWERON in server.allowed_actions:
            try:
                instances_api.server_action(
                    server_id=server.id, action=ServerAction.POWERON, zone=server.zone
                )
            except ScalewayException as e:
                logger.error(f'Failed deploying server: "{server.name}": {e}')
        else:
            if server.state == ServerState.RUNNING:
                logger.debug(f"Server {server.name} is already running.")
            else:
                logger.warning(
                    f"Server {server.name} cannot be powered on. \
                    Current state: {server.state}"
                )


def _finalize(
    arch: Architecture,
    servers: Iterable[Server],
    gateway: Optional[Gateway],
    private_network: Optional[PrivateNetwork],
) -> Tuple[Roles, Networks]:
    if arch == Architecture.PrivateNetwork:
        extra = {
            "gateway": gateway.ipv4.address,
            "gateway_user": DEFAULT_GATEWAY_USER,
            "gateway_port": DEFAULT_GATEWAY_PORT,
            "forward_agent": True,
        }
    else:
        extra = {}
    roles = Roles()
    for server in servers:
        for tag in server.tags:
            if server.public_ip is not None:
                address = server.public_ip.address
            else:
                internal_address = ".".join(
                    [server.name, private_network.name, "internal"]
                )
                extra.update(
                    {"internal_address": internal_address, SERVER_ID: server.id}
                )
                address = internal_address
            roles[tag] += [
                Host(
                    address,
                    alias=str(server.name),
                    # TODO: support different users
                    user="root",
                    extra=extra,
                )
            ]
    # TODO: private networks handling
    networks = Networks()
    sync_info(roles, networks, inplace=True)
    return roles, networks


# TODO: Maybe only enable ipv6 for some roles only
def enable_ipv6(
    client: Client,
    servers: Iterable[Server],
) -> Iterable[Ip]:
    """THIS FUNCTION IS HACKY AS THE SDK DOES NOT SEEM TO SUPPORT
    SOME OF THE OPTION DISPLAYED IN THE DOCUMENTATION AND AVAILABLE IN
    THE CLI LIKE THE cloud-ini OPTION

    Args:
        client (Client): client
        servers (Iterable[Server]): list of servers to enable ipv6 on

    Returns:
        Iterable[Ip]: list of reserevrd ipv6 instances
    """
    instance_api = _get_instances_api(client=client)
    # https://www.scaleway.com/en/docs/instances/how-to/use-flexips/#flexible-ipv6
    # num_servers = len(servers)
    # if num_servers > 5:
    #     logger.error("Can't attach more than 5 IPv6 addresses at the same time, not enabling IPv6")
    #     return
    ips = []
    for server in servers:
        resp = instance_api.create_ip(zone=server.zone, type_=IpType.ROUTED_IPV6)
        if resp.ip is None:
            logger.error(f'Failed creating IPv6 IP for "{server.name}"')
            continue

        ipv6 = resp.ip
        # could not find sdk function
        subprocess.run(
            [
                "scw",
                "instance",
                "ip",
                "attach",
                ipv6.id,
                f"server-id={server.id}",
            ],
            stdout=subprocess.DEVNULL,
        )
        ips.append(ipv6)
    return ips


def _remove_ip_instance_ips(client=Client, ips=Iterable[Ip]):
    instance_api = _get_instances_api(client=client)
    all_ips = instance_api.list_ips_all()
    for ip in ips:
        try:
            instance_api.delete_ip(ip=ip.id)
        except Exception as e:
            logger.error(e)


class ScalewayInstancesProvider(Provider):
    # TODO Improve IP address handling to remove addresses once they are not needed
    ENOS_META = "Provider"

    deployed_instances = []
    undeployed_instances = []

    private_network = None
    gateway = None
    gateway_network = None

    def init(
        self, force_deploy: bool = False, start_time: Optional[int] = None, **kwargs
    ) -> Tuple[Roles, Networks]:
        # Implementation for initializing Scaleway instances
        self.provider_conf: Configuration

        self.arch = Architecture(self.provider_conf.architecture)

        client = _get_client(
            file_path=self.provider_conf.config_file,
            profile=self.provider_conf.profile,
        )

        if self.arch == Architecture.PrivateNetwork:
            logger.info("Deploying instances in a private network")
            private_network = _create_private_network(
                client=client,
                region=self.provider_conf.region,
                net_name=self.provider_conf.network_name,
            )
            gateway = _create_public_gateway(
                client=client,
                name=self.provider_conf.gateway_name,
                type=self.provider_conf.gateway_type,
                zone=self.provider_conf.gateway_zone,
            )
            gateway_network = _create_gateway_network(
                client=client, private_network=private_network, gateway=gateway
            )
        else:
            logger.info("Deploying instances standalone")
            private_network = None
            gateway = None
            gateway_network = None

        servers = _create_server(
            client=client,
            provider_conf=self.provider_conf,
            architecture=self.arch,
        )

        if self.arch == Architecture.PrivateNetwork:
            _attach_servers_to_network(
                client=client, private_network=private_network, servers=servers
            )

        logger.info("Starting Scaleway instances")
        deployed, undeployed = start_servers(client=client, servers=servers)

        ipv6s: Optional[Iterable[Ip]] = None
        if self.provider_conf.enable_ipv6:
            ipv6s = enable_ipv6(client=client, servers=deployed)
            logger.info("Enabled IPv6")

        # TODO: maybe not store those in attributes
        self.deployed_instances = deployed
        self.undeployed_instances = undeployed
        self.gateway = gateway
        self.private_network = private_network
        self.gateway_network = gateway_network
        self.ipv6s = ipv6s

        roles, networks = _finalize(self.arch, deployed, gateway, private_network)

        if self.arch == Architecture.PrivateNetwork:
            _open_ports(
                provider_conf=self.provider_conf,
                client=client,
                roles=roles,
                gateway=gateway,
                gateway_network=gateway_network,
            )
        return roles, networks

    def destroy(self, wait: bool = False, **kwargs):
        client = _get_client(
            file_path=self.provider_conf.config_file,
            profile=self.provider_conf.profile,
        )
        instances_api = _get_instances_api(client=client)
        block_api = _get_block_api(client=client)
        for server in self.deployed_instances:
            logger.info(f"Terminating server {server.name}")
            _terminate_server(instances_api=instances_api, server=server)
            _clean_server_block_storage(block_api=block_api, server=server)
        for server in self.undeployed_instances:
            logger.info(f"Deleting undeployed server {server.name}")
            _delete_server(instances_api=instances_api, server=server)
            _clean_server_block_storage(block_api=block_api, server=server)

        if self.ipv6s:
            _remove_ip_instance_ips(client=client, ips=self.ipv6s)
            logger.info("Done deleting IPs")

        if self.arch == Architecture.Standalone:
            return

        logger.info("Deleting gateway")
        vpc_gw_api = _get_vpc_gw_api(client=client)
        vpc_gw_api.delete_gateway_network(
            gateway_network_id=self.gateway_network.id, zone=self.gateway.zone
        )
        vpc_gw_api.delete_gateway(gateway_id=self.gateway.id, delete_ip=True)
        logger.info("Deleted gateway")

        logger.info("Deleting private network")
        vpc_api = _get_vpc_api(client=client)
        for _ in range(10):
            try:
                vpc_api.delete_private_network(
                    private_network_id=self.private_network.id,
                    region=self.private_network.region,
                )
                break
            except Exception:
                logger.debug("Failed deleting network, retrying...")
                time.sleep(3)
        if self.provider_conf.network_name not in [
            pn.name
            for pn in vpc_api.list_private_networks(
                region=self.provider_conf.region
            ).private_networks
        ]:
            logger.info("Deleted private network")
        else:
            logger.warning(
                f'Failed deleting private network "{self.private_network.name}" '
                f'with ID "{self.private_network.id}"'
            )

    def offset_walltime(self, offset: int):
        # No walltime
        pass
