# Copyright (C) 2020 Linaro Limited
#
# Author: Antonio Terceiro <antonio.terceiro@linaro.org>
#
# SPDX-License-Identifier: GPL-2.0-or-later
from __future__ import annotations

import os
import random
import subprocess
import time
from pathlib import Path
from typing import TYPE_CHECKING

from lava_common.exceptions import InfrastructureError, JobError

if TYPE_CHECKING:
    from typing import Iterable, Optional

    from lava_dispatcher.action import Action


class DockerLogin:
    def __init__(self, registry: str, user: str, password: str):
        self.registry = registry
        self.user = user
        self.password = password

    def _build_docker_login_command(self) -> list[str]:
        return [
            "docker",
            "login",
            "--username",
            self.user,
            "--password-stdin",
            self.registry,
        ]

    def login(self, action: Action) -> str:
        docker_home = action.mkdtemp()
        action.logger.debug(f"Login in to remote registry {self.registry!r}")
        action.parsed_command(
            self._build_docker_login_command(),
            env={**os.environ, "HOME": docker_home},
            input=self.password,
        )
        action.logger.debug("Login successful")
        return docker_home


class DockerRun:
    def __init__(self, image: str):
        self.image = image
        self._is_local_image = False
        self._container_name: str | None = None
        self._network_name: str | None = None
        self._suffix = ""
        self._hostname: str | None = None
        self._workdir: str | None = None
        self._shared_devices: list[str] = []
        self._bind_mounts: list[tuple[str, str, bool]] = []
        self._environment_vars: list[tuple[str, str]] = []
        self._is_interactive = False
        self._is_tty = False
        self._has_init = True
        self._docker_options: list[str] = []
        self._docker_run_options: list[str] = []
        self._docker_login: DockerLogin | None = None

    @classmethod
    def from_parameters(cls, params, job):
        image: str = params["image"]
        run = cls(image)
        suffix = "-lava-" + str(job.job_id)
        if "container_name" in params:
            run.name(params["container_name"] + suffix)
        run.suffix(suffix)
        run.network(params.get("network_from", None))
        run.local(params.get("local", False))
        for device in params.get("devices", []):
            run.add_device(device)
        if login_data := params.get("login"):
            try:
                user = login_data["user"]
                password = login_data["password"]
            except KeyError:
                raise JobError(
                    "Incorrect docker login dictionary. "
                    "Expected 'user' and 'password' entries."
                )
            registry = login_data.get("registry", image.split("/", maxsplit=1)[0])

            run._docker_login = DockerLogin(
                registry=registry, user=user, password=password
            )
        return run

    def local(self, local):
        self._is_local_image = local

    def name(self, name, random_suffix=False):
        suffix = ""
        if random_suffix:
            CHARS = "01234567890abcdefghijklmnopqrtsuwxyz"
            suffix = "".join(random.SystemRandom().choice(CHARS) for i in range(10))
        self._container_name = name + suffix

    def network(self, network):
        self._network_name = network

    def suffix(self, suffix):
        self._suffix = suffix

    def hostname(self, hostname):
        self._hostname = hostname

    def workdir(self, workdir):
        self._workdir = workdir

    def init(self, init):
        self._has_init = init

    def add_device(self, device, skip_missing=False):
        if not Path(device).exists() and skip_missing:
            return
        if ":" in device:
            return
        self._shared_devices.append(device)

    def add_docker_options(self, *options):
        self._docker_options += options

    def add_docker_run_options(self, *options):
        self._docker_run_options += options

    def add_device_docker_method_options(
        self, docker_method_conf: dict[str, Iterable[None | list[str] | str]]
    ) -> None:
        # Preprocess docker option list, to better support partial
        # overriding of them via device dict:
        # 1. Filter out None, to make it easier to template
        # YAML syntactic lists with Jinja2:
        # '- {{ some_opt_from_device_dict }}'
        # (if not default, will be set to None).
        # 2. Flatten sublists, `- ['--opt1', '--opt2']`.
        def preproc_opts(opts: Iterable[None | list[str] | str]) -> list[str]:
            res = []
            for o in opts:
                if o is None:
                    continue
                elif isinstance(o, list):
                    res.extend(o)
                else:
                    res.append(o)
            return res

        if "global_options" in docker_method_conf:
            self.add_docker_options(*preproc_opts(docker_method_conf["global_options"]))
        if "options" in docker_method_conf:
            self.add_docker_run_options(*preproc_opts(docker_method_conf["options"]))

    def interactive(self):
        self._is_interactive = True

    def tty(self):
        self._is_tty = True

    def bind_mount(self, source, destination=None, read_only=False):
        if not destination:
            destination = source
        self._bind_mounts.append((source, destination, read_only))

    def environment(self, variable, value):
        self._environment_vars.append((variable, value))

    def cmdline(self, *args):
        cmd = ["docker"] + self._docker_options + ["run"] + self._docker_run_options
        cmd += self.interaction_options()
        cmd += self.start_options()
        cmd.append(self.image)
        cmd += args
        return cmd

    def interaction_options(self):
        cmd = []
        if self._is_interactive:
            cmd.append("--interactive")
        if self._is_tty:
            cmd.append("--tty")
        return cmd

    def start_options(self):
        cmd = ["--rm"]
        if self._has_init:
            cmd.append("--init")
        if self._container_name:
            cmd.append(f"--name={self._container_name}")
        if self._network_name:
            cmd.append(f"--network=container:{self._network_name}{self._suffix}")
        if self._hostname:
            cmd.append(f"--hostname={self._hostname}")
        if self._workdir:
            cmd.append(f"--workdir={self._workdir}")
        for dev in self._shared_devices:
            cmd.append(f"--device={dev}")
        for src, dest, read_only in self._bind_mounts:
            opt = f"--mount=type=bind,source={src},destination={dest}"
            if read_only:
                opt += ",readonly=true"
            cmd.append(opt)
        for variable, value in self._environment_vars:
            cmd.append(f"--env={variable}={value}")
        return cmd

    def run(self, *args, action, capture=False, error_msg=None):
        self.prepare(action)
        cmd = self.cmdline(*args)
        if capture:
            return subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode(
                "utf-8", errors="replace"
            )
        else:
            return action.run_cmd(cmd, error_msg=error_msg)

    def prepare(self, action: Action):
        pull = not self._is_local_image
        if self._is_local_image:
            if action.run_cmd(
                [
                    "docker",
                    *self._docker_options,
                    "image",
                    "inspect",
                    "--format",
                    f"Image {self.image} exists locally",
                    self.image,
                ],
                allow_fail=True,
            ):
                action.logger.warning(
                    "Unable to inspect docker image '%s'" % self.image
                )
                pull = True

        if pull:
            docker_pull_env: dict[str, str] = {**os.environ}
            if self._docker_login is not None:
                docker_pull_env["HOME"] = self._docker_login.login(action)

            action.run_cmd(
                ["docker", *self._docker_options, "pull", self.image],
                env=docker_pull_env,
            )
        self._check_image_arch(action)

    def wait(self, shell=None):
        delay = 1
        while True:
            try:
                # If possible, check that docker's shell command didn't exit
                # yet.
                if shell and not shell.isalive():
                    raise InfrastructureError("Docker container unexpectedly exited")
                subprocess.check_call(
                    [
                        "docker",
                        *self._docker_options,
                        "inspect",
                        "--format=.",
                        self._container_name,
                    ],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
                return
            except subprocess.CalledProcessError:
                time.sleep(delay)
                delay = delay * 2  # exponential backoff

    def wait_file(self, filename: str, timeout: Optional[int] = None) -> None:
        delay = 1
        start = time.monotonic()
        while True:
            try:
                subprocess.check_call(
                    [
                        "docker",
                        *self._docker_options,
                        "exec",
                        self._container_name,
                        "test",
                        "-e",
                        filename,
                    ],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
                return
            except subprocess.CalledProcessError:
                if timeout is not None and time.monotonic() - start > timeout:
                    raise
                time.sleep(delay)
                delay = delay * 2  # exponential backoff

    def destroy(self):
        if self._container_name:
            subprocess.call(
                ["docker", *self._docker_options, "rm", "-f", self._container_name],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )

    def _check_image_arch(self, action):
        host = subprocess.check_output(["arch"], text=True).strip()
        try:
            container = subprocess.check_output(
                [
                    "docker",
                    *self._docker_options,
                    "inspect",
                    "--format",
                    "{{.Architecture}}",
                    self.image,
                ],
                text=True,
            ).strip()
        except FileNotFoundError:
            raise InfrastructureError("'docker' command not available on the worker")
        # amd64 = x86_64
        if host == "amd64":
            host = "x86_64"
        if container == "amd64":
            container = "x86_64"
        if host == "arm64":
            host = "aarch64"
        if container == "arm64":
            container = "aarch64"
        if host != container:
            action.logger.warning(
                f"Architecture mismatch: host is {host}, container is {container}. This *might* work, but if it does, will probably be a lot slower than if the container image architecture matches the host."
            )


class DockerContainer(DockerRun):
    def __init__(self, image: str):
        super().__init__(image)
        self._started = False

    def run(self, args, action):
        self.start(action)
        cmd = ["docker", *self._docker_options, "exec"]
        cmd += self.interaction_options()
        cmd.append(self._container_name)
        cmd += args
        action.run_cmd(cmd)

    def get_output(self, args, action):
        self.start(action)
        cmd = ["docker", *self._docker_options, "exec"]
        cmd += self.interaction_options()
        cmd.append(self._container_name)
        cmd += args
        return action.parsed_command(cmd)

    def check_output(self, cmd):
        return subprocess.check_output(cmd).decode("utf-8")

    def start(self, action):
        if self._started:
            return

        cmd = [
            "docker",
            *self._docker_options,
            "run",
            *self._docker_run_options,
            "--detach",
        ]
        cmd += self.start_options()
        cmd.append(self.image)
        cmd += ["sleep", "infinity"]
        action.run_cmd(cmd)
        self.wait()
        self._started = True

    def stop(self, action):
        action.run_cmd(["docker", *self._docker_options, "stop", self._container_name])
