# SPDX-FileCopyrightText: 2021-2023 Blender Authors
#
# SPDX-License-Identifier: Apache-2.0

import base64
import glob
import inspect
import multiprocessing
import os
import pathlib
import platform
import pickle
import subprocess
import sys

from collections.abc import (
    Callable,
)

from .common import normalize_device_id
from .config import TestConfig
from .device import TestMachine


class TestFailure(Exception):
    def __init__(self, *args, message, output_lines=[], **kwargs):
        super().__init__(message, *args)
        self.message = message
        self.output_lines = output_lines

    def __str__(self):
        msg = self.message
        if self.output_lines:
            msg += f":\n{'': <10} | "
            msg += f"\n{'': <10} | ".join(l.rstrip(' \r\n\t') for l in self.output_lines)
        return msg


class TestEnvironment:
    def __init__(self, blender_git_dir: pathlib.Path, base_dir: pathlib.Path):
        self.blender_git_dir = blender_git_dir
        self.base_dir = base_dir
        self.blender_dir = base_dir / 'blender'
        self.build_dir = base_dir / 'build'
        self.install_dir = self.build_dir / "bin"
        self.benchmarks_dir = self.blender_git_dir / 'tests' / 'benchmarks'
        self.git_executable = 'git'
        self.cmake_executable = 'cmake'
        self.cmake_options = ['-DWITH_INTERNATIONAL=OFF', '-DWITH_BUILDINFO=OFF']
        self.log_file = None
        self.machine = None
        self._title_cache = {}
        self._init_default_blender_executable()
        self.set_default_blender_executable()

    def get_machine(self, need_gpus: bool = True) -> None:
        if not self.machine or (need_gpus and not self.machine.has_gpus):
            self.machine = TestMachine(self, need_gpus)

        return self.machine

    def init(self, build: bool, blender_bin: str) -> None:
        if not self.benchmarks_dir.exists():
            sys.stderr.write(f'Error: benchmark files directory not found at {self.benchmarks_dir}')
            sys.exit(1)

        # Create benchmarks folder contents.
        print(f'Init {self.base_dir}')
        self.base_dir.mkdir(parents=True, exist_ok=True)

        if len(self.get_config_names()) == 0:
            config_dir = self.base_dir / 'default'
            print(f'Creating default configuration in {config_dir}')
            TestConfig.write_default_config(self, config_dir, blender_bin)

        if build:
            if not self.blender_dir.exists():
                print(f'Init git worktree in {self.blender_dir}')
                self.call([self.git_executable, 'worktree', 'add', '--detach',
                          self.blender_dir, 'HEAD'], self.blender_git_dir)
            else:
                print(f'Exists {self.blender_dir}')

            if not self.build_dir.exists():
                print(f'Init build in {self.build_dir}')
                self.build_dir.mkdir()
                # No translation to avoid dealing with submodules
                self.call([self.cmake_executable, self.blender_dir, '.'] + self.cmake_options, self.build_dir)
            else:
                print(f'Exists {self.build_dir}')

            print("Building")
            git_hash = self.resolve_git_hash('HEAD')
            self.build(git_hash, self.install_dir)

        print('Done')

    def checkout(self, git_hash: str, update_submodules: bool = True) -> None:
        # Checkout Blender revision
        if not self.blender_dir.exists():
            sys.stderr.write('\n\nError: no build set up, run `./benchmark init --build` first\n')
            sys.exit(1)

        self.call([self.git_executable, 'clean', '-f', '-d'], self.blender_dir)
        self.call([self.git_executable, 'reset', '--hard', 'HEAD'], self.blender_dir)
        self.call([self.git_executable, 'checkout', '--detach', git_hash], self.blender_dir)

        if update_submodules:
            self.call([self.git_executable, 'submodule', 'update', '--recursive', '--force'], self.blender_dir)

    def _submodule_key(self) -> str:
        """
        Return a key to identify the current submodule checkout. The key only includes checked out
        submodules.

        The key is used to determine if submodules have changed between builds. In that case the
        CMakeCache should be reevaluated as it can point to different libraries/versions.
        """
        log_lines = self.call([self.git_executable, 'submodule', 'status'], self.blender_dir, silent=True)
        # Only add submodules that are checked out
        log_lines = [line for line in log_lines if not line.startswith("-")]
        submodule_key = ",".join([l.split()[0] for l in log_lines])
        return submodule_key

    def build(self, git_hash: str, install_dir: pathlib.Path, update_submodules: bool = True) -> bool:
        # Build Blender revision
        if not self.build_dir.exists():
            sys.stderr.write('\n\nError: no build set up, run `./benchmark init --build` first\n')
            sys.exit(1)

        # Skip if build with same hash is already done.
        if install_dir.resolve() != self.install_dir.resolve():
            complete_txt = pathlib.Path(install_dir) / "complete.txt"
            if complete_txt.is_file():
                if complete_txt.read_text().strip() == git_hash:
                    self._init_default_blender_executable()
                    return True
                # Different hash, build again.
                complete_txt.unlink()
        else:
            complete_txt = None

        old_submodule_key = self._submodule_key()
        self.checkout(git_hash, update_submodules)
        new_submodule_key = self._submodule_key()
        if old_submodule_key != new_submodule_key:
            cmake_cache = self.build_dir / "CMakeCache.txt"
            if cmake_cache.exists():
                cmake_cache.unlink()

        jobs = str(multiprocessing.cpu_count())
        cmake_options = list(self.cmake_options)
        cmake_options += [f"-DCMAKE_INSTALL_PREFIX={install_dir}"]
        try:
            self.call([self.cmake_executable, self.blender_dir, '.'] + cmake_options, self.build_dir)
            self.call([self.cmake_executable, '--build', '.', '-j', jobs, '--target', 'install'], self.build_dir)
            if complete_txt:
                complete_txt.write_text(git_hash)
        except KeyboardInterrupt as e:
            raise e
        except:
            return False

        self._init_default_blender_executable()
        return True

    def set_blender_executable(self, executable_path: pathlib.Path, environment: dict = {}) -> None:
        if executable_path.is_dir():
            executable_path = self._blender_executable_from_path(executable_path)

        # Run all Blender commands with this executable.
        self.blender_executable = executable_path
        self.blender_executable_environment = environment

    def _blender_executable_name(self) -> pathlib.Path:
        if platform.system() == "Windows":
            return pathlib.Path('blender.exe')
        elif platform.system() == "Darwin":
            return pathlib.Path('Blender.app') / 'Contents' / 'MacOS' / 'Blender'
        else:
            return pathlib.Path('blender')

    def _has_install_files(self, executable_path: pathlib.Path) -> bool:
        """
        Check if the given executable is a full installation of blender by checking
        the availability of the 'license' directory.
        """
        # Follow symlink to actual install location.
        executable_path = executable_path.resolve()
        if platform.system() == "Darwin":
            license_path = executable_path.parent.parent / 'Resources' / 'text' / 'license'
        else:
            license_path = executable_path.parent / 'license'
        return license_path.is_dir()

    def _blender_executable_from_path(self, executable: pathlib.Path) -> pathlib.Path:
        if executable.is_dir():
            # Directory
            executable = executable / self._blender_executable_name()
        elif not executable.is_file() and executable.name == 'blender':
            # Executable path without proper path on Windows or macOS.
            executable = executable.parent / self._blender_executable_name()

        if executable.is_file() and self._has_install_files(executable):
            return executable

        return None

    def _init_default_blender_executable(self) -> None:
        # Find a default executable to run commands independent of testing a specific build.
        # Try own built executable.
        built_executable = self._blender_executable_from_path(self.install_dir)
        if built_executable:
            self.default_blender_executable = built_executable
            return

        # Try find an executable in the configs.
        for config_name in self.get_config_names():
            for executable in TestConfig.read_blender_executables(self, config_name):
                executable = self._blender_executable_from_path(executable)
                if executable:
                    self.default_blender_executable = executable
                    return

        # Fallback to a "blender" command in the hope it's available.
        self.default_blender_executable = pathlib.Path("blender")

    def set_default_blender_executable(self) -> None:
        self.blender_executable = self.default_blender_executable
        self.blender_executable_environment = {}

    def set_log_file(self, filepath: pathlib.Path, clear=True) -> None:
        # Log all commands and output to this file.
        self.log_file = filepath

        if clear:
            self.log_file.unlink(missing_ok=True)

    def unset_log_file(self) -> None:
        self.log_file = None

    def call(self, args: list[str], cwd: pathlib.Path, silent: bool = False, environment: dict = {}) -> list[str]:
        # Execute command with arguments in specified directory,
        # and return combined stdout and stderr output.

        # Open log file for writing
        f = None
        if self.log_file:
            if not self.log_file.exists():
                self.log_file.parent.mkdir(parents=True, exist_ok=True)
            f = open(self.log_file, 'a', encoding='utf-8', errors='ignore')
            f.write('\n' + ' '.join([str(arg) for arg in args]) + '\n\n')

        env = os.environ
        if len(environment):
            env = env.copy()
            for key, value in environment.items():
                env[key] = value

        proc = subprocess.Popen(args, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env)

        # Read line by line
        lines = []
        try:
            while proc.poll() is None:
                line = proc.stdout.readline()
                if line:
                    line_str = line.decode('utf-8', 'ignore')
                    lines.append(line_str)
                    if f:
                        f.write(line_str)
        except KeyboardInterrupt as e:
            # Avoid processes that keep running when interrupting.
            proc.terminate()
            raise e

        # Raise error on failure
        if proc.returncode != 0 and not silent:
            raise TestFailure(message="Error executing command", output_lines=lines)

        return lines

    def call_blender(self, args: list[str], foreground=False) -> list[str]:
        # Execute Blender command with arguments.
        common_args = ['--factory-startup', '-noaudio', '--enable-autoexec', '--python-exit-code', '1']
        if sys.platform == 'win32':
            # Set HighQoS level on Windows to avoid reduced performance when the window is out of focus.
            # See: https://learn.microsoft.com/en-us/windows/win32/procthread/quality-of-service
            common_args += ['--qos', 'high']
        if foreground:
            common_args += ['--no-window-focus', '--window-geometry', '0', '0', '1024', '768', '--gpu-vsync', 'off']
        else:
            common_args += ['--background']

        return self.call([self.blender_executable] + common_args + args, cwd=self.base_dir,
                         environment=self.blender_executable_environment)

    def run_in_blender(self,
                       function: Callable[[dict], dict],
                       args: dict,
                       blender_args: list = [],
                       foreground=False) -> dict:
        # Run function in a Blender instance. Arguments and return values are
        # passed as a Python object that must be serializable with pickle.

        # Get information to call this function from Blender.
        package_path = pathlib.Path(__file__).parent.parent
        functionname = function.__name__
        modulename = inspect.getmodule(function).__name__

        # Serialize arguments in base64, to avoid having to escape it.
        args = base64.b64encode(pickle.dumps(args))
        output_prefix = 'TEST_OUTPUT: '

        expression = (f'import sys, pickle, base64;'
                      f'sys.path.append(r"{package_path}");'
                      f'import {modulename};'
                      f'args = pickle.loads(base64.b64decode({args}));'
                      f'result = {modulename}.{functionname}(args);'
                      f'result = base64.b64encode(pickle.dumps(result));'
                      f'print("\\n{output_prefix}" + result.decode() + "\\n")')

        expr_args = blender_args + ['--python-expr', expression]
        lines = self.call_blender(expr_args, foreground=foreground)

        # Parse output.
        for line in lines:
            if line.startswith(output_prefix):
                output = line[len(output_prefix):].strip()
                result = pickle.loads(base64.b64decode(output))
                return result, lines

        return {}, lines

    def find_blend_files(self, dirpath: pathlib.Path) -> list:
        # Find .blend files in subdirectories of the given directory in the
        # lib/benchmarks directory.
        dirpath = self.benchmarks_dir / dirpath
        filepaths = []
        for filename in glob.iglob(str(dirpath / '*.blend'), recursive=True):
            filepaths.append(pathlib.Path(filename))
        return filepaths

    def get_config_names(self) -> list:
        names = []

        if self.base_dir.exists():
            for dirname in os.listdir(self.base_dir):
                dirpath = self.base_dir / dirname / 'config.py'
                if dirpath.exists():
                    names.append(dirname)

        return names

    def get_configs(self, name: str = None, names_only: bool = False) -> list:
        # Get list of configurations in the benchmarks directory.
        configs = []

        for config_name in self.get_config_names():
            if not name or config_name == name:
                if names_only:
                    configs.append(config_name)
                else:
                    configs.append(TestConfig(self, config_name))

        return configs

    def resolve_git_hash(self, revision):
        # Get git hash for a tag or branch.
        lines = self.call([self.git_executable, 'rev-parse', revision], self.blender_git_dir)
        return lines[0].strip() if len(lines) else revision

    def git_hash_date(self, git_hash):
        # Get commit data for a git hash.
        lines = self.call([self.git_executable, 'log', '-n1', git_hash, '--format=%at'], self.blender_git_dir)
        return int(lines[0].strip()) if len(lines) else 0

    def commits_in_window(self, after_ts: int, before_ts: int) -> list[tuple[str, int]]:
        """List commits in a time window, oldest first.

        Returns a list of ``(commit_hash, unix_timestamp)`` tuples
        for commits reachable from ``HEAD`` whose commit date falls
        between ``after_ts`` and ``before_ts``.
        """
        try:
            lines = self.call(
                [self.git_executable, 'log', '--first-parent', '--reverse',
                 '--after=' + str(after_ts - 1), '--before=' + str(before_ts),
                 '--format=%H %ct', 'HEAD'],
                self.blender_git_dir, silent=True)
        except:
            return []
        result: list[tuple[str, int]] = []
        for line in lines:
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            if len(parts) >= 2:
                try:
                    result.append((parts[0][:12], int(parts[1])))
                except:
                    pass
        return result

    def commit_title(self, git_hash: str) -> str:
        """Return the one-line subject of a commit, with caching."""
        if git_hash in self._title_cache:
            return self._title_cache[git_hash]
        try:
            lines = self.call(
                [self.git_executable, 'log', '-n1', '--format=%s', git_hash],
                self.blender_git_dir, silent=True)
            title = lines[0].strip() if lines else ''
        except:
            title = ''
        self._title_cache[git_hash] = title
        return title

    def resolve_device(self, device_str: str) -> tuple[str, str]:
        """Resolve a device string to a device_id and gpu_backend pair."""
        machine = self.get_machine(need_gpus=True)
        device_id = device_str
        gpu_backend = 'default'

        sanitized_str = normalize_device_id(device_str)
        for device in machine.devices:
            if normalize_device_id(device.id) == sanitized_str or device.type == device_str:
                device_id = device.id
                gpu_backend = {
                    'VULKAN': 'vulkan',
                    'METAL': 'metal',
                    'OPENGL': 'opengl'
                }.get(device.type, 'default')
                break

        return device_id, gpu_backend