package relocator

import (
	"debug/elf"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
)

// isDynamicELF checks if the given file is a dynamically linked ELF binary.
func isDynamicELF(filePath string) (bool, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return false, fmt.Errorf("failed to open file: %v", err)
	}
	defer file.Close()

	// Parse the file as an ELF binary using the debug/elf package
	elfFile, err := elf.NewFile(file)
	if err != nil {
		// If it's not an ELF file, return false
		return false, nil
	}

	// Check if the ELF binary has a PT_DYNAMIC segment (indicating dynamic linking)
	for _, prog := range elfFile.Progs {
		if prog.Type == elf.PT_DYNAMIC {
			return true, nil // Dynamically linked ELF
		}
	}

	return false, nil // Not dynamically linked
}

// copyRegularFile copies a file from src to dest, preserving ownership and permissions.
func copyRegularFile(dest, src string) error {
	fmt.Printf("Copying file from %s to %s\n", src, dest)

	// Open source file
	srcFile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	// Create destination file
	destFile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destFile.Close()

	// Copy file contents
	if _, err := io.Copy(destFile, srcFile); err != nil {
		return err
	}

	// Get file info to preserve permissions and ownership
	fileInfo, err := os.Lstat(src)
	if err != nil {
		return err
	}

	// Preserve file permissions
	if err := os.Chmod(dest, fileInfo.Mode()); err != nil {
		return err
	}

	// Preserve file ownership if possible
	if stat, ok := fileInfo.Sys().(*syscall.Stat_t); ok {
		_ = os.Lchown(dest, int(stat.Uid), int(stat.Gid)) // Ignore ownership error for non-root users
	}

	return nil
}

// relocateELFBinary copies the ELF binary and updates its shared library paths.
func relocateELFBinary(prefix, dest, src string) error {
	// Copy the ELF binary to the new destination
	if err := copyRegularFile(dest, src); err != nil {
		return err
	}

	// Get shared libraries using "ldd" command
	out, err := exec.Command("ldd", src).Output()
	if err != nil {
		return fmt.Errorf("failed to get linked libraries for %s: %v", src, err)
	}

	// Store unique parent directories of linked libraries
	dirSet := make(map[string]struct{})

	// Parse the output from ldd and relocate each library
	for _, line := range strings.Split(string(out), "\n") {
		parts := strings.Fields(line)
		if len(parts) >= 3 && parts[1] == "=>" {
			libraryPath := parts[2] // Linked library path

			// Get the directory of the library and add to the set
			libraryDir := filepath.Dir(libraryPath)
			dirSet[libraryDir] = struct{}{}

			// Recursively relocate the library
			if err := RelocateFile(prefix, libraryPath); err != nil {
				return fmt.Errorf("failed to relocate library %s: %v", libraryPath, err)
			}
		}
	}

	var interpreter string
	// Get the interpreter using "patchelf --print-interpreter"
	interpOut, err := exec.Command("patchelf", "--print-interpreter", src).Output()
	if err == nil {
		interpreter = strings.TrimSpace(string(interpOut)) // Trim any whitespace or newlines

		RelocateFile(prefix, interpreter)
	} else {
		interpreter = ""
	}

	// Aggregate the directories into a colon-separated rpath
	var dirList []string
	for dir := range dirSet {
		dirList = append(dirList, filepath.Join(prefix, dir))
	}

	// If an interpreter was found, relocate it
	if interpreter != "" {
		newInterpPath := filepath.Join(prefix, interpreter)
		if _, err := exec.Command("patchelf", "--set-interpreter", newInterpPath, dest).Output(); err != nil {
			return fmt.Errorf("failed to update interpreter for %s: %v", dest, err)
		}
		fmt.Printf("Relocated interpreter from %s to %s\n", interpreter, newInterpPath)
	}

	rpath := strings.Join(dirList, ":")

	if rpath != "" {
		// Update the ELF binary's rpath using patchelf
		if _, err := exec.Command("patchelf", "--set-rpath", rpath, dest).Output(); err != nil {
			return fmt.Errorf("failed to update rpath for %s: %v", dest, err)
		}

		fmt.Printf("Relocated binary %s to %s with rpath: %s\n", src, dest, rpath)
	}

	return nil
}

// RelocateFile relocates a file (regular file, directory, or symlink) to the target prefix.
func RelocateFile(prefix, src string) error {
	if src != "/" {
		parent := filepath.Dir(src)
		// Recursively relocate parent directories
		if err := RelocateFile(prefix, parent); err != nil {
			return err
		}
	}

	// Determine the destination path under the prefix
	dst := filepath.Join(prefix, src)

	// Check if the destination already exists
	if _, err := os.Lstat(dst); os.IsNotExist(err) {
		// Get source file information
		sinfo, err := os.Lstat(src)
		if err != nil {
			return fmt.Errorf("error processing %s: %v", src, err)
		}

		// Handle symbolic links
		if (sinfo.Mode() & os.ModeSymlink) != 0 {
			linkTarget, err := os.Readlink(src)
			if err != nil {
				return err
			}

			// Handle absolute and relative symlinks
			if filepath.IsAbs(linkTarget) {
				relocatedTarget := filepath.Join(prefix, linkTarget)
				fmt.Printf("Creating absolute symlink %s as %s\n", relocatedTarget, dst)
				if err := os.Symlink(relocatedTarget, dst); err != nil {
					return err
				}
				return RelocateFile(prefix, linkTarget)
			} else {
				parent := filepath.Dir(src)
				target := filepath.Join(parent, linkTarget)
				fmt.Printf("Creating relative symlink %s as %s\n", linkTarget, dst)
				if err := os.Symlink(linkTarget, dst); err != nil {
					return err
				}
				return RelocateFile(prefix, target)
			}
		}

		// Handle directories
		if sinfo.IsDir() {
			fmt.Printf("Creating directory %s\n", dst)
			return os.MkdirAll(dst, sinfo.Mode())
		}

		// Ignore non-regular files such as sockets, devices, etc.
		if !sinfo.Mode().IsRegular() {
			fmt.Printf("Skipping non-regular file: %s\n", src)
			return nil
		}

		// Handle regular files (ELF or otherwise)
		isElf, err := isDynamicELF(src)
		if err != nil {
			return err
		}
		if isElf {
			return relocateELFBinary(prefix, dst, src)
		} else {
			return copyRegularFile(dst, src)
		}
	}

	// Follow symlinks if the destination is a symlink
	if dinfo, err := os.Lstat(dst); err == nil && (dinfo.Mode()&os.ModeSymlink) != 0 {
		linkTarget, err := os.Readlink(src)
		if err != nil {
			return err
		}
		if filepath.IsAbs(linkTarget) {
			return RelocateFile(prefix, linkTarget)
		}
		parent := filepath.Dir(src)
		return RelocateFile(prefix, filepath.Join(parent, linkTarget))
	}

	return nil
}