/**
# Copyright (c) NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
**/

package nvcdi

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/sys/unix"

	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
	"github.com/NVIDIA/nvidia-container-toolkit/pkg/lookup"
)

// NewDriverDiscoverer creates a discoverer for the libraries and binaries associated with a driver installation.
// The supplied NVML Library is used to query the expected driver version.
func (l *nvmllib) NewDriverDiscoverer() (discover.Discover, error) {
	return (*nvcdilib)(l).newDriverVersionDiscoverer()
}

func (l *nvcdilib) newDriverVersionDiscoverer() (discover.Discover, error) {
	version, err := l.driver.Version()
	if err != nil || version == "" || version == "*.*" {
		return nil, fmt.Errorf("failed to determine driver version (%q): %w", version, err)
	}

	libraries, err := l.NewDriverLibraryDiscoverer(version)
	if err != nil {
		return nil, fmt.Errorf("failed to create discoverer for driver libraries: %v", err)
	}

	ipcs, err := l.newIPCDiscoverer()
	if err != nil {
		return nil, fmt.Errorf("failed to create discoverer for IPC sockets: %v", err)
	}

	firmwares, err := l.newDriverFirmwareDiscoverer(version)
	if err != nil {
		return nil, fmt.Errorf("failed to create discoverer for GSP firmware: %v", err)
	}

	binaries := l.newDriverBinariesDiscoverer()

	d := discover.Merge(
		libraries,
		ipcs,
		firmwares,
		binaries,
	)

	return d, nil
}

func (l *nvcdilib) newIPCDiscoverer() (discover.Discover, error) {
	if l.featureFlags[FeatureDisableIPCDiscoverer] {
		return nil, nil
	}
	return discover.NewIPCDiscoverer(l.logger, l.driver.Root)
}

// NewDriverLibraryDiscoverer creates a discoverer for the libraries associated with the specified driver version.
func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover, error) {
	versionSuffixLibraryMounts, err := l.getVersionSuffixDriverLibraryMounts(version)
	if err != nil {
		return nil, err
	}
	legacyNVVMLibraryMounts, err := l.getLegacyNVVMLibraryMounts()
	if err != nil {
		return nil, fmt.Errorf("failed to get legacy nvvm library mounts: %w", err)
	}
	explicitLibraryMounts, err := l.getExplicitDriverLibraryMounts()
	if err != nil {
		return nil, err
	}

	libraries := discover.Merge(
		versionSuffixLibraryMounts,
		legacyNVVMLibraryMounts,
		explicitLibraryMounts,
	)

	var discoverers []discover.Discover

	driverDotSoSymlinksDiscoverer := discover.WithDriverDotSoSymlinks(
		l.logger,
		libraries,
		// Since we don't only match version suffixes, we now need to match on wildcards.
		"",
		l.hookCreator,
	)
	discoverers = append(discoverers, driverDotSoSymlinksDiscoverer)

	cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, &discover.EnableCUDACompatHookOptions{HostDriverVersion: version})
	discoverers = append(discoverers, cudaCompatLibHookDiscoverer)

	updateLDCache, _ := discover.NewLDCacheUpdateHook(l.logger, libraries, l.hookCreator)
	discoverers = append(discoverers, updateLDCache)

	disableDeviceNodeModification := l.hookCreator.Create(DisableDeviceNodeModificationHook)
	discoverers = append(discoverers, disableDeviceNodeModification)

	driverLibDirectories, err := l.driver.GetDriverLibDirectories()
	if err != nil {
		return nil, fmt.Errorf("failed to get libcuda.so parent directory path: %w", err)
	}
	environmentVariable := &discover.EnvVar{
		Name:  "NVIDIA_CTK_LIBCUDA_DIR",
		Value: strings.Join(driverLibDirectories, string(filepath.ListSeparator)),
	}
	discoverers = append(discoverers, environmentVariable)

	d := discover.Merge(discoverers...)

	return d, nil
}

func (l *nvcdilib) getVersionSuffixDriverLibraryMounts(version string) (discover.Discover, error) {
	versionSuffixLibraryPaths, err := l.getVersionLibs(version)
	if err != nil {
		return nil, fmt.Errorf("failed to get libraries for driver version: %v", err)
	}

	mounts := discover.NewMounts(
		l.logger,
		lookup.NewFileLocator(
			lookup.WithLogger(l.logger),
			lookup.WithRoot(l.driver.Root),
		),
		l.driver.Root,
		versionSuffixLibraryPaths,
	)

	return mounts, nil
}

func (l *nvcdilib) getExplicitDriverLibraryMounts() (discover.Discover, error) {
	if !l.featureFlags[FeatureEnableExplicitDriverLibraries] {
		return nil, nil
	}

	// List of explicit libraries to locate
	// TODO(ArangoGutierrez): we should load the version of the libraries from
	// the sandboxutils-filelist or have a way to allow users to specify the
	// libraries to mount from the config file.
	explicitLibraries := []string{
		"libEGL.so",
		"libGL.so",
		"libGLESv1_CM.so",
		"libGLESv2.so",
		"libGLX.so",
		"libGLdispatch.so",
		"libOpenCL.so",
		"libOpenGL.so",
		"libnvidia-api.so",
		"libnvidia-egl-xcb.so",
		"libnvidia-egl-xlib.so",
	}

	driverLibraryLocator, err := l.driver.DriverLibraryLocator()
	if err != nil {
		return nil, fmt.Errorf("failed to get driver library locator: %w", err)
	}
	mounts := discover.NewMounts(
		l.logger,
		driverLibraryLocator,
		l.driver.Root,
		explicitLibraries,
	)

	return mounts, nil

}

func (l *nvcdilib) getLegacyNVVMLibraryMounts() (discover.Discover, error) {
	legacyNVMMLibrary := []string{
		"libnvidia-nvvm70.so.*",
	}

	driverLibraryLocator, err := l.driver.DriverLibraryLocator()
	if err != nil {
		return nil, fmt.Errorf("failed to get driver library locator: %w", err)
	}
	mounts := discover.NewMounts(
		l.logger,
		driverLibraryLocator,
		l.driver.Root,
		legacyNVMMLibrary,
	)

	return mounts, nil
}

func getUTSRelease() (string, error) {
	utsname := &unix.Utsname{}
	if err := unix.Uname(utsname); err != nil {
		return "", err
	}
	return unix.ByteSliceToString(utsname.Release[:]), nil
}

func getFirmwareSearchPaths(logger logger.Interface) ([]string, error) {

	var firmwarePaths []string
	if p := getCustomFirmwareClassPath(logger); p != "" {
		logger.Debugf("using custom firmware class path: %s", p)
		firmwarePaths = append(firmwarePaths, p)
	}

	utsRelease, err := getUTSRelease()
	if err != nil {
		return nil, fmt.Errorf("failed to get UTS_RELEASE: %v", err)
	}

	standardPaths := []string{
		filepath.Join("/lib/firmware/updates/", utsRelease),
		"/lib/firmware/updates/",
		filepath.Join("/lib/firmware/", utsRelease),
		"/lib/firmware/",
	}

	return append(firmwarePaths, standardPaths...), nil
}

// getCustomFirmwareClassPath returns the custom firmware class path if it exists.
func getCustomFirmwareClassPath(logger logger.Interface) string {
	customFirmwareClassPath, err := os.ReadFile("/sys/module/firmware_class/parameters/path")
	if err != nil {
		logger.Warningf("failed to get custom firmware class path: %v", err)
		return ""
	}

	return strings.TrimSpace(string(customFirmwareClassPath))
}

// newDriverFirmwareDiscoverer creates a discoverer for GSP firmware associated with the specified driver version.
func (l *nvcdilib) newDriverFirmwareDiscoverer(version string) (discover.Discover, error) {
	gspFirmwareSearchPaths, err := getFirmwareSearchPaths(l.logger)
	if err != nil {
		return nil, fmt.Errorf("failed to get firmware search paths: %v", err)
	}
	gspFirmwarePaths := filepath.Join("nvidia", version, "gsp*.bin")
	return discover.NewMounts(
		l.logger,
		lookup.NewFileLocator(
			lookup.WithLogger(l.logger),
			lookup.WithRoot(l.driver.Root),
			lookup.WithSearchPaths(gspFirmwareSearchPaths...),
		),
		l.driver.Root,
		[]string{gspFirmwarePaths},
	), nil
}

// newDriverBinariesDiscoverer creates a discoverer for GSP firmware associated with the GPU driver.
func (l *nvcdilib) newDriverBinariesDiscoverer() discover.Discover {
	return discover.NewMounts(
		l.logger,
		lookup.NewExecutableLocator(l.logger, l.driver.Root),
		l.driver.Root,
		[]string{
			"nvidia-smi",              /* System management interface */
			"nvidia-debugdump",        /* GPU coredump utility */
			"nvidia-persistenced",     /* Persistence mode utility */
			"nvidia-cuda-mps-control", /* Multi process service CLI */
			"nvidia-cuda-mps-server",  /* Multi process service server */
			"nvidia-imex",             /* NVIDIA IMEX Daemon */
			"nvidia-imex-ctl",         /* NVIDIA IMEX control */
		},
	)
}

// getVersionLibs checks the LDCache for libraries ending in the specified driver version.
// Although the ldcache at the specified driverRoot is queried, the paths are returned relative to this driverRoot.
// This allows the standard mount location logic to be used for resolving the mounts.
func (l *nvcdilib) getVersionLibs(version string) ([]string, error) {
	l.logger.Infof("Using driver version %v", version)

	libraries, err := l.driver.DriverLibraryLocator("vdpau")
	if err != nil {
		return nil, fmt.Errorf("failed to get driver library locator: %w", err)
	}

	libs, err := libraries.Locate("*.so." + version)
	if err != nil {
		return nil, fmt.Errorf("failed to locate libraries for driver version %v: %v", version, err)
	}

	if l.driver.Root == "/" || l.driver.Root == "" {
		return libs, nil
	}

	var relative []string
	for _, lib := range libs {
		relative = append(relative, strings.TrimPrefix(lib, l.driver.Root))
	}

	return relative, nil
}
