import importlib.resources
import math
from importlib.abc import Traversable
from pathlib import Path

import docker
import humanfriendly
import luigi
import netaddr
from docker.client import DockerClient
from docker.models.containers import Container
from docker.models.volumes import Volume
from docker.types import DeviceRequest
from jinja2 import Template

from exasol_integration_test_docker_environment.lib import PACKAGE_NAME
from exasol_integration_test_docker_environment.lib.base.docker_base_task import (
    DockerBaseTask,
)
from exasol_integration_test_docker_environment.lib.base.json_pickle_parameter import (
    JsonPickleParameter,
)
from exasol_integration_test_docker_environment.lib.base.ssh_access import (
    SshKey,
    SshKeyCache,
)
from exasol_integration_test_docker_environment.lib.base.still_running_logger import (
    StillRunningLogger,
)
from exasol_integration_test_docker_environment.lib.docker.images.create.utils.pull_log_handler import (
    PullLogHandler,
)
from exasol_integration_test_docker_environment.lib.docker.images.image_info import (
    ImageInfo,
)
from exasol_integration_test_docker_environment.lib.models.data.container_info import (
    ContainerInfo,
)
from exasol_integration_test_docker_environment.lib.models.data.database_info import (
    DatabaseInfo,
)
from exasol_integration_test_docker_environment.lib.models.data.docker_network_info import (
    DockerNetworkInfo,
)
from exasol_integration_test_docker_environment.lib.models.data.ssh_info import SshInfo
from exasol_integration_test_docker_environment.lib.test_environment.db_version import (
    DbVersion,
)
from exasol_integration_test_docker_environment.lib.test_environment.docker_container_copy import (
    DockerContainerCopy,
)
from exasol_integration_test_docker_environment.lib.test_environment.parameter.docker_db_test_environment_parameter import (
    DbOsAccess,
    DockerDBTestEnvironmentParameter,
)
from exasol_integration_test_docker_environment.lib.test_environment.ports import (
    Ports,
    find_free_ports,
)

CERTIFICATES_MOUNT_DIR = "/certificates"
CERTIFICATES_DEFAULT_DIR = "/exa/etc/ssl/"


def int_or_none(value: str | None) -> int | None:
    return None if value is None else int(value)


class SpawnTestDockerDatabase(DockerBaseTask, DockerDBTestEnvironmentParameter):
    environment_name: str = luigi.Parameter()
    db_container_name: str = luigi.Parameter()
    attempt: int = luigi.IntParameter(default=1)
    network_info: DockerNetworkInfo = JsonPickleParameter(DockerNetworkInfo, significant=False)  # type: ignore
    ip_address_index_in_subnet: int = luigi.IntParameter(significant=False)
    docker_runtime: str | None = luigi.OptionalParameter(
        default=None, significant=False
    )
    certificate_volume_name: str | None = luigi.OptionalParameter(
        default=None, significant=False
    )
    additional_db_parameter: tuple[str, ...] = luigi.ListParameter()
    docker_environment_variables: tuple[str, ...] = luigi.ListParameter()
    accelerator: tuple[str, ...] = luigi.ListParameter()
    ssh_user: str = luigi.Parameter(default="root")
    ssh_key_file: str | Path | None = luigi.OptionalParameter(
        default=None, significant=False
    )

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        if self.ip_address_index_in_subnet < 0:
            raise Exception(
                "ip_address_index_in_subnet needs to be greater than 0 got %s"
                % self.ip_address_index_in_subnet
            )

        self.db_version = DbVersion.from_db_version_str(self.docker_db_image_version)
        self.docker_db_config_resource_name = f"docker_db_config/{self.db_version}"
        self.internal_ports = Ports.default_ports
        if self.ssh_port_forward is None:
            self.ssh_port_forward = str(find_free_ports(1)[0])

        def first_of(*args):
            return next(x for x in args if x is not None)

        bucketfs_http_port: int = int(
            first_of(
                self.bucketfs_http_port_forward,
                self.bucketfs_port_forward,
                Ports.default_ports.bucketfs_http,
            )
        )
        bucketfs_https_port: int = int(
            first_of(
                self.bucketfs_https_port_forward, Ports.default_ports.bucketfs_https
            )
        )

        self.forwarded_ports = Ports(
            database=int_or_none(self.database_port_forward),
            bucketfs=bucketfs_http_port,
            ssh=int_or_none(self.ssh_port_forward),
            bucketfs_https=bucketfs_https_port,
        )

    def run_task(self) -> None:
        subnet = netaddr.IPNetwork(self.network_info.subnet)
        db_ip_address = str(subnet[2 + self.ip_address_index_in_subnet])
        db_private_network = f"{db_ip_address}/{subnet.prefixlen}"
        database_info = None
        if self.network_info.reused:
            database_info = self._try_to_reuse_database(db_ip_address)
        if database_info is None:
            database_info = self._create_database_container(
                db_ip_address, db_private_network
            )
        self.return_object(database_info)

    def _try_to_reuse_database(self, db_ip_address: str) -> DatabaseInfo | None:
        self.logger.info("Try to reuse database container %s", self.db_container_name)
        try:
            database_info = self._create_database_info(
                db_ip_address=db_ip_address, reused=True
            )
            return database_info
        except Exception as e:
            self.logger.warning(
                "Tried to reuse database container %s, but got Exeception %s. "
                "Fallback to create new database.",
                self.db_container_name,
                e,
            )
        return None

    def _get_ssh_key(self) -> SshKey:
        if self.ssh_key_file:
            return SshKey.read_from(self.ssh_key_file)
        self.ssh_key_file = SshKeyCache().private_key
        return SshKey.from_cache()

    def _handle_output(self, output_generator, image_info: ImageInfo) -> None:
        log_file_path = self.get_log_path().joinpath("pull_docker_db_image.log")
        with PullLogHandler(log_file_path, self.logger, image_info) as log_handler:
            still_running_logger = StillRunningLogger(
                self.logger, "pull image %s" % image_info.get_source_complete_name()
            )
            for log_line in output_generator:
                still_running_logger.log()
                log_handler.handle_log_lines(log_line)

    def _get_network_aliases(self) -> list[str]:
        network_aliases = [
            "exasol_test_database",
            "exasol-test-database",
            self.db_container_name,
        ]
        return network_aliases

    def _connect_docker_network(
        self,
        docker_client: DockerClient,
        container: Container,
        ip_address: str,
    ) -> None:
        network = docker_client.networks.get(self.network_info.network_name)
        aliases = self._get_network_aliases()
        network.connect(container, ipv4_address=ip_address, aliases=aliases)

    def _port_mapping(
        self, internal_ports: Ports, forwarded_ports: Ports
    ) -> dict[str, tuple[str, str]]:
        result = {}
        for name, internal in internal_ports.__dict__.items():
            forward = forwarded_ports.__getattribute__(name)
            if forward:
                result[f"{internal}/tcp"] = ("0.0.0.0", forward)
        return result

    def _create_database_container(
        self, db_ip_address: str, db_private_network: str
    ) -> DatabaseInfo:
        def get_authorized_keys(ssh_key) -> str:
            """
            Multiple authorized_keys can be comma-separated.
            """
            if self.db_os_access != DbOsAccess.SSH:
                return ""
            return ssh_key.public_key_as_string("itde-ssh-access")

        def enable_ssh_access(container: Container, authorized_keys: str):
            copy = DockerContainerCopy(container)
            copy.add_string_to_file(".ssh/authorized_keys", authorized_keys)
            copy.copy("/root/")

        self.logger.info("Starting database container %s", self.db_container_name)
        ssh_key = self._get_ssh_key()
        authorized_keys = get_authorized_keys(ssh_key)
        with self._get_docker_client() as docker_client:
            try:
                docker_client.containers.get(self.db_container_name).remove(
                    force=True, v=True
                )
            except:
                pass
            docker_db_image_info = self._pull_docker_db_images_if_necessary()
            db_volume = self._prepare_db_volume(
                docker_client,
                db_private_network,
                authorized_keys,
                docker_db_image_info,
            )
            port_mapping = self._port_mapping(self.internal_ports, self.forwarded_ports)
            volumes = {db_volume.name: {"bind": "/exa", "mode": "rw"}}
            if self.certificate_volume_name is not None:
                volumes[self.certificate_volume_name] = {
                    "bind": CERTIFICATES_MOUNT_DIR,
                    "mode": "rw",
                }
            device_requests = self._build_device_requests()
            db_container = docker_client.containers.create(
                image="%s" % (docker_db_image_info.get_source_complete_name()),
                name=self.db_container_name,
                detach=True,
                privileged=True,
                volumes=volumes,
                network_mode=None,
                ports=port_mapping,
                runtime=self.docker_runtime,
                environment=list(self.docker_environment_variables),
                device_requests=device_requests,
            )
            enable_ssh_access(db_container, authorized_keys)
            self._connect_docker_network(docker_client, db_container, db_ip_address)
            db_container.start()
            database_info = self._create_database_info(
                db_ip_address=db_ip_address, reused=False
            )
            return database_info

    def _create_database_info(self, db_ip_address: str, reused: bool) -> DatabaseInfo:
        with self._get_docker_client() as docker_client:
            db_container = docker_client.containers.get(self.db_container_name)
            if db_container.status != "running":
                raise Exception(f"Container {self.db_container_name} not running")
            network_aliases = self._get_network_aliases()
            container_info = ContainerInfo(
                container_name=self.db_container_name,
                ip_address=db_ip_address,
                network_aliases=network_aliases,
                network_info=self.network_info,
                volume_name=self._get_db_volume_name(),
            )
            ssh_info = SshInfo(self.ssh_user, str(self.ssh_key_file or ""))
            database_info = DatabaseInfo(
                host=db_ip_address,
                ports=self.internal_ports,
                reused=reused,
                container_info=container_info,
                ssh_info=ssh_info,
                forwarded_ports=self.forwarded_ports,
            )
            return database_info

    def _pull_docker_db_images_if_necessary(self) -> ImageInfo:
        image_name = "exasol/docker-db"
        assert self.docker_db_image_version is not None
        docker_db_image_info = ImageInfo(
            target_repository_name=image_name,
            source_repository_name=image_name,
            source_tag=self.docker_db_image_version,
            target_tag=self.docker_db_image_version,
            hash_value="",
            commit="",
            image_description=None,
        )
        with self._get_docker_client() as docker_client:
            try:
                docker_client.images.get(
                    docker_db_image_info.get_source_complete_name()
                )
            except docker.errors.ImageNotFound as e:
                self.logger.info(
                    "Pulling docker-db image %s",
                    docker_db_image_info.get_source_complete_name(),
                )
                output_generator = docker_client.api.pull(
                    docker_db_image_info.source_repository_name,
                    tag=docker_db_image_info.source_tag,
                    stream=True,
                )
                self._handle_output(output_generator, docker_db_image_info)
        return docker_db_image_info

    def _prepare_db_volume(
        self,
        docker_client,
        db_private_network: str,
        authorized_keys: str,
        docker_db_image_info: ImageInfo,
    ) -> Volume:
        volume, container = self._prepare_volume(
            docker_client,
            self._get_db_volume_name(),
            self._get_db_volume_preparation_container_name(),
            remove_old_instances=True,
        )
        try:
            self._upload_init_db_files(container, db_private_network, authorized_keys)
            self._execute_init_db(volume, container)
            return volume
        finally:
            container.remove(force=True)

    def _get_db_volume_preparation_container_name(self):
        return f"""{self.db_container_name}_preparation"""

    def _get_db_volume_name(self):
        return f"""{self.db_container_name}_volume"""

    def _remove_container(self, container_name: str) -> None:
        try:
            with self._get_docker_client() as docker_client:
                docker_client.containers.get(container_name).remove(force=True)
                self.logger.info("Removed container %s", container_name)
        except docker.errors.NotFound:
            pass

    def _remove_volume(self, volume_name: str) -> None:
        try:
            with self._get_docker_client() as docker_client:
                docker_client.volumes.get(volume_name).remove(force=True)
                self.logger.info("Removed volume %s", volume_name)
        except docker.errors.NotFound:
            pass

    def _prepare_volume(
        self,
        docker_client: docker.api.APIClient,
        volume_name,
        container_name,
        remove_old_instances: bool = False,
    ) -> tuple[Volume, Container]:
        """
        Create an intermediate Docker Container containing a volume that
        can be mounted into another Docker Container.
        """
        if remove_old_instances:
            self._remove_container(container_name)
            self._remove_volume(volume_name)
        volume = docker_client.volumes.create(volume_name)
        container = docker_client.containers.run(
            image="ubuntu:22.04",
            name=container_name,
            auto_remove=True,
            command="sleep infinity",
            detach=True,
            volumes={volume.name: {"bind": "/exa", "mode": "rw"}},
            labels={
                "test_environment_name": self.environment_name,
                "container_type": "db_volume_preparation_container",
            },
        )
        return volume, container

    def _db_file(self, filename: str) -> Traversable:
        return (
            importlib.resources.files(PACKAGE_NAME)
            / self.docker_db_config_resource_name
            / filename
        )

    def _upload_init_db_files(
        self,
        container: Container,
        db_private_network: str,
        authorized_keys: str,
    ) -> None:
        copy = DockerContainerCopy(container)
        init_script = self._db_file("init_db.sh")
        copy.add_string_to_file("init_db.sh", init_script.read_text())
        self._add_exa_conf(copy, db_private_network, authorized_keys)
        copy.copy("/")

    def _add_exa_conf(
        self,
        copy: DockerContainerCopy,
        db_private_network: str,
        authorized_keys: str,
    ) -> None:
        """
        Multiple authorized_keys can be comma-separated.
        """
        certificate_dir = (
            CERTIFICATES_MOUNT_DIR
            if self.certificate_volume_name is not None
            else CERTIFICATES_DEFAULT_DIR
        )
        template_file = self._db_file("EXAConf")
        template = Template(template_file.read_text())
        additional_db_parameter_str = " ".join(self.additional_db_parameter)
        rendered_template = template.render(
            private_network=db_private_network,
            db_version=str(self.db_version),
            db_port=self.internal_ports.database,
            ssh_port=self.internal_ports.ssh,
            bucketfs_http_port=self.internal_ports.bucketfs_http,
            bucketfs_https_port=self.internal_ports.bucketfs_https,
            image_version=self.docker_db_image_version,
            mem_size=self.mem_size,
            disk_size=self.disk_size,
            name_servers=",".join(self.nameservers),
            certificate_dir=certificate_dir,
            additional_db_parameters=additional_db_parameter_str,
            authorized_keys=authorized_keys,
        )
        copy.add_string_to_file("EXAConf", rendered_template)

    def _execute_init_db(
        self, db_volume: Volume, preparation_container: Container
    ) -> None:
        disk_size_in_bytes = humanfriendly.parse_size(self.disk_size)
        min_overhead_in_gigabyte = 2  # Exasol needs at least a 2 GB larger device than the configured disk size
        overhead_factor = max(
            0.01, (min_overhead_in_gigabyte * 1024 * 1024 * 1024) / disk_size_in_bytes
        )  # and in general 1% larger
        device_size_in_bytes = disk_size_in_bytes * (1 + overhead_factor)
        device_size_in_megabytes = math.ceil(
            device_size_in_bytes / (1024 * 1024)
        )  # The init_db.sh script works with MB, because its faster
        self.logger.info(
            f"Creating database volume of size {device_size_in_megabytes / 1024} GB using and overhead factor of {overhead_factor}"
        )
        (exit_code, output) = preparation_container.exec_run(
            cmd=f"bash /init_db.sh {device_size_in_megabytes}"
        )
        if exit_code != 0:
            raise Exception(
                "Error during preparation of docker-db volume {} got following output {}".format(
                    db_volume.name, output
                )
            )

    def cleanup_task(self, success) -> None:
        if (success and not self.no_database_cleanup_after_success) or (
            not success and not self.no_database_cleanup_after_failure
        ):
            volume_container = self._get_db_volume_preparation_container_name()
            try:
                self.logger.info(f"Cleaning up container %s", volume_container)
                self._remove_container(volume_container)
            except Exception as e:
                self.logger.error(
                    f"Error during removing container %s: %s", volume_container, e
                )

            try:
                self.logger.info(f"Cleaning up container %s", self.db_container_name)
                self._remove_container(self.db_container_name)
            except Exception as e:
                self.logger.error(
                    f"Error during removing container %s: %s", self.db_container_name, e
                )

            db_volume_name = self._get_db_volume_name()
            try:
                self.logger.info(f"Cleaning up docker volume %s", db_volume_name)
                self._remove_volume(db_volume_name)
            except Exception as e:
                self.logger.error(
                    f"Error during removing docker volume %s: %s", db_volume_name, e
                )

    def _build_device_requests(self) -> list[DeviceRequest]:
        if "nvidia=all" in self.accelerator:
            return [
                DeviceRequest(
                    driver="nvidia",
                    count=-1,
                    capabilities=[["gpu"]],
                )
            ]
        else:
            return []
