package patchmanager

import (
	"archive/tar"
	"bufio"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"slices"

	"github.com/spf13/afero"
	"gopkg.in/yaml.v2"

	"github.com/go-logr/logr"
	"github.com/hashicorp/go-version"

	"edge-infra.dev/pkg/sds/patching/common"
)

var (
	upgradePhases = []string{
		"pre-reboot",
		"upgrade-mode",
		"post-reboot",
	}
)

const completeSuffix = "-complete"

type DirectivesConfig struct {
	Directives []struct {
		Regex         string `yaml:"regex"`
		ScriptOptions []struct {
			Name   string
			Skip   bool
			Weight int
		} `yaml:"script_options"`
	}
}

func (p *PatchManager) installPatchset() error {
	dlComplete, err := p.DownloadComplete()
	if err != nil {
		return err
	}
	if dlComplete {
		p.Log.Info("Download already complete")
		return nil
	}

	p.Log.Info("House keeping old version files")
	versionRetention := []string{p.CurrentVer, p.TargetVer}
	if err := housekeepVersions(versionRetention, p.Log, p.Cfg); err != nil {
		return err
	}

	if err := prepareLegacy(p.CurrentVer, p.Log, p.Cfg); err != nil {
		return err
	}

	p.Log.Info("Preparing directories")
	targetPath := GetArtefactsPath(p.TargetVer, p.Cfg)
	if err := prepareDirectories(targetPath, p.Cfg); err != nil {
		return err
	}

	p.Log.Info("Extracting files")
	if err := extractPatchset(p.TargetVer, p.Log, p.Cfg); err != nil {
		return fmt.Errorf("Failed to unpack artefacts from blob: %w", err)
	}

	p.Log.Info("Preparing upgrade scripts")
	if err := p.installPatchScripts(); err != nil {
		return fmt.Errorf("Failed to install upgrade scripts: %w", err)
	}

	// Clean up temp scripts dir
	if err := p.Fs.RemoveAll(p.Cfg.ScriptsTempPath); err != nil {
		return err
	}

	if err := addCasperSymlink(p.TargetVer, p.Cfg); err != nil {
		return err
	}

	if _, err := os.OpenFile(filepath.Join(targetPath, ".complete"), os.O_RDONLY|os.O_CREATE, 0655); err != nil {
		return err
	}

	//if err := ValidateFiles(targetPath, log); err != nil {
	//	log.Error(err, "Failed to validate checksum")
	//	return err
	//}
	return nil
}

func addCasperSymlink(version string, cfg common.Config) error {
	squashName := "filesystem.squashfs"
	versionPath := GetArtefactsPath(version, cfg)
	casper := path.Join(versionPath, "casper")
	if err := os.Mkdir(casper, 0755); err != nil {
		return err
	}
	return os.Symlink(path.Join(versionPath, squashName), path.Join(casper, squashName))
}

// Prepare patching directories for install
func prepareDirectories(targetPath string, cfg common.Config) error {
	if err := os.RemoveAll(targetPath); err != nil {
		return err
	}

	if err := os.RemoveAll(cfg.ArtefactsTempPath); err != nil {
		return err
	}

	if err := os.RemoveAll(cfg.ScriptsTempPath); err != nil {
		return err
	}

	return createScriptDirectories(cfg)
}

// Extract tar file to temp destination files dependent on root dir
func extractTarFile(source io.Reader, log logr.Logger, cfg common.Config) error {
	uncompressedStream, err := gzip.NewReader(source)
	if err != nil {
		return err
	}
	defer uncompressedStream.Close()

	tarReader := tar.NewReader(uncompressedStream)

	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}

		filePath, err := getDestPath(header.Name, cfg)
		if err != nil {
			return err
		}

		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(filePath, 0755); err != nil {
				return err
			}
		case tar.TypeReg:
			if err := writeFile(filePath, tarReader, header, log); err != nil {
				return err
			}
		default:
			return fmt.Errorf("Unknown header type: %x name: %s", header.Typeflag, header.Name)
		}
	}
	return nil
}

// Generate output path based on file
func getDestPath(header string, cfg common.Config) (string, error) {
	var filePath string

	if strings.HasPrefix(header, common.BootFiles) {
		filePath = filepath.Join(cfg.ArtefactsTempPath, strings.Replace(header, common.BootFiles, "", 1))
	} else if strings.HasPrefix(header, common.ScriptFiles) {
		filePath = filepath.Join(cfg.ScriptsTempPath, strings.Replace(header, common.ScriptFiles, "", 1))
	} else {
		return "", fmt.Errorf("Unexpected filepath in header %s", header)
	}

	// #nosec G305
	// Validation for G305
	if !strings.HasPrefix(filePath, filepath.Clean(filePath)) {
		return "", fmt.Errorf("%s: %s", "content filepath is tainted", header)
	}
	return filePath, nil
}

// Write file from tar IOStream to disk
func writeFile(filepath string, tarReader *tar.Reader, header *tar.Header, log logr.Logger) error {
	const bufferSize int64 = 4 * common.KiB
	const stepSize int64 = 100 * common.MiB

	fsMode := fs.FileMode(header.Mode) /* #nosec G115 */
	outFile, err := os.OpenFile(filepath, os.O_WRONLY|os.O_CREATE, fsMode)
	if err != nil {
		return err
	}
	defer outFile.Close()

	for read := int64(0); read < header.Size; read += bufferSize {
		_, err := io.CopyN(outFile, tarReader, bufferSize)
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if read%(stepSize) == 0 {
			log.Info(fmt.Sprintf("Extracting %s %d Mib / %d Mib", filepath, read/common.MiB, header.Size/common.MiB))
		}
	}
	log.Info(fmt.Sprintf("Extracting %s done", filepath))

	stat, err := outFile.Stat()
	if err != nil {
		return err
	}
	if stat.Size() != header.Size {
		return fmt.Errorf("Extracted file %s size does not match header", header.Name)
	}
	return nil
}

func extractPatchset(targetVersion string, log logr.Logger, cfg common.Config) error {
	file, err := os.Open(cfg.PatchsetMount)
	if err != nil {
		return err
	}

	defer file.Close()

	if err := extractTarFile(file, log, cfg); err != nil {
		return err
	}

	return os.Rename(cfg.ArtefactsTempPath, GetArtefactsPath(targetVersion, cfg))
}

func GetArtefactsPath(version string, cfg common.Config) string {
	return filepath.Join(cfg.ArtefactsPath, version)
}

// During the move to IEN v1.10.0. The scripts directory location has changed.
// Rename existing and symlink
func prepareLegacy(currentVer string, log logr.Logger, cfg common.Config) error {
	maxVer, _ := version.NewVersion("v1.10.0")
	current, err := version.NewVersion(currentVer)
	if err != nil {
		return err
	}
	if current.GreaterThanOrEqual(maxVer) {
		log.Info("Removing BWC symlink")
		return os.RemoveAll(cfg.LegacyScriptsPath)
	}
	stat, err := os.Lstat(cfg.LegacyScriptsPath)
	if err == nil {
		log.Info("Legacy script dir/link exists")
		if stat.Mode()&fs.ModeSymlink != 0 {
			log.Info("Legacy upgrade scripts is already a symlink")
			return nil
		}
		// The new dir shouldn't exist outside of testing, clean it up either way
		log.Info("Deleting " + cfg.ScriptsPath)
		if err := os.RemoveAll(cfg.ScriptsPath); err != nil {
			return nil
		}
		log.Info("Legacy upgrade scripts dir exists. Renaming")
		if err := os.Rename(cfg.LegacyScriptsPath, cfg.ScriptsPath); err != nil {
			return err
		}
		if err := os.Chmod(cfg.ScriptsPath, 0744); err != nil {
			return err
		}
	}
	log.Info("Legacy script dir/link does not exist")
	return os.Symlink("patching", cfg.LegacyScriptsPath)
}

func createScriptDirectories(cfg common.Config) error {
	if err := os.MkdirAll(cfg.ScriptsPath, 0744); err != nil {
		return err
	}
	for _, phase := range upgradePhases {
		path := filepath.Join(cfg.ScriptsPath, phase)
		if err := os.RemoveAll(path); err != nil {
			return err
		}
		if err := os.MkdirAll(path, 0744); err != nil {
			return err
		}
		if err := os.MkdirAll(path, 0744); err != nil {
			return err
		}
		// create folder for completed scripts
		if err := os.MkdirAll(path+completeSuffix, 0744); err != nil {
			return err
		}
		if err := os.Chmod(path+completeSuffix, 0744); err != nil {
			return err
		}
	}
	return nil
}

// Delete versions files not included in retentionVers
func housekeepVersions(retentionVers []string, log logr.Logger, cfg common.Config) error {
	log.Info("House keeping old versions")
	versionsDir := cfg.ArtefactsPath
	artefacts, err := os.ReadDir(versionsDir)
	if err != nil {
		return err
	}

	for _, artefact := range artefacts {
		if !slices.Contains(retentionVers, artefact.Name()) {
			log.Info("Removing version " + artefact.Name())
			if err = os.RemoveAll(filepath.Join(versionsDir, artefact.Name())); err != nil {
				return err
			}
		} else {
			log.Info("Keeping version " + artefact.Name())
		}
	}
	return nil
}

func ValidateFiles(fs afero.Fs, targetPath string, log logr.Logger) error {
	file, err := fs.Open(filepath.Join(targetPath, "checksum"))
	if err != nil {
		log.Error(err, "Checksum file not found")
	}
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		var filePath string
		var sha256sum string
		n, err := fmt.Sscanf(scanner.Text(), "%s %s", &sha256sum, &filePath)
		if n != 2 || err != nil {
			return fmt.Errorf("Failed to parse checksum file")
		}
		if err := ValidateChecksum(filepath.Join(targetPath, filePath), sha256sum); err != nil {
			return err
		}
	}
	return nil
}

func ValidateChecksum(filePath, checksum string) error {
	fileChecksum, err := GetSha256(filePath)
	if err != nil {
		return err
	}

	if fileChecksum != checksum {
		return fmt.Errorf("File %s checksum is not %s", filePath, checksum)
	}
	return nil
}

func GetSha256(filePath string) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", err
	}

	defer file.Close()

	hasher := sha256.New()
	if _, err := io.Copy(hasher, file); err != nil {
		return "", err
	}
	return hex.EncodeToString(hasher.Sum(nil)), nil
}

func (p *PatchManager) installPatchScripts() error {
	controlPlane, err := p.getControlPlaneHostname()
	if err != nil {
		p.Log.Error(err, "Could not find control plane")
		return err
	}
	isCP := p.HostName == controlPlane

	currentVer, err := version.NewVersion(p.CurrentVer)
	if err != nil {
		return err
	}

	dirs, err := afero.ReadDir(p.Fs, p.Cfg.ScriptsTempPath)
	if err != nil {
		return err
	}

	versions, err := versionSort(dirs)
	if err != nil {
		return err
	}

	configs, err := loadDirectivesConfig(versions, p.Cfg)
	if err != nil {
		return err
	}

	hashMaps := make(map[string]map[string]string)
	for _, phase := range upgradePhases {
		hashMaps[phase] = make(map[string]string)
	}
	for _, version := range versions {
		if version.LessThanOrEqual(currentVer) && !strings.HasSuffix(version.Original(), "-dev") {
			p.Log.Info("Skipping upgrade scripts for version " + version.Original())
			continue
		}

		p.Log.Info("Installing upgrade scripts for version " + version.Original())
		if err := installScriptsVersion(currentVer, version, configs, isCP, p.Log, p.Cfg, hashMaps); err != nil {
			return err
		}
	}
	return nil
}

func installScriptsVersion(currentVer, version *version.Version, configs []DirectivesConfig, isControlPlane bool, log logr.Logger, cfg common.Config, hashMaps map[string]map[string]string) error {
	path := filepath.Join(cfg.ScriptsTempPath, version.Original())

	for _, phase := range upgradePhases {
		phasePath := filepath.Join(path, phase)
		files, err := os.ReadDir(phasePath)
		if err != nil {
			log.Error(err, "Failed to read script dir. Skipping")
			return nil
		}
		for _, file := range files {
			fileName := file.Name()
			filePath := filepath.Join(phasePath, fileName)

			fileSha256, err := GetSha256(filePath)
			if err != nil {
				return err
			}

			phaseHashMap := hashMaps[phase]
			_, exists := phaseHashMap[fileSha256]
			if exists {
				log.Info(fmt.Sprintf("Skipping upgrade script (duplicate) %s %s %s", version.Original(), phase, fileName))
				continue
			}
			phaseHashMap[fileSha256] = fileName
			if stat, _ := os.Lstat(filepath.Join(cfg.ScriptsPath, phase+completeSuffix, fileSha256)); stat != nil {
				log.Info(fmt.Sprintf("Skipping upgrade script (exists) %s %s %s", version.Original(), phase, fileName))
				continue
			}
			if fileName == ".gitignore" {
				continue
			}
			if roleFilter(fileName, isControlPlane) {
				log.Info(fmt.Sprintf("Skipping upgrade script (role filter) %s %s %s", version.Original(), phase, fileName))
				continue
			}
			skip, weight := checkDirectives(configs, currentVer.Original(), fileName)
			if skip {
				log.Info(fmt.Sprintf("Skipping upgrade script (config.yaml) %s %s %s", version.Original(), phase, fileName))
				continue
			}
			if IsExecutable(filePath) {
				fileName = fmt.Sprintf("%d_%s_%s", weight, padVersion(version), fileName)
			}
			newPath := filepath.Join(cfg.ScriptsPath, phase, fileName)
			log.Info(fmt.Sprintf("Installing upgrade %s to %s", file.Name(), newPath))
			if err := os.Rename(filePath, newPath); err != nil {
				return err
			}
		}
	}

	return nil
}

func loadDirectivesConfig(versions []*version.Version, cfg common.Config) (configs []DirectivesConfig, err error) {
	for _, version := range versions {
		configPath := filepath.Join(cfg.ScriptsTempPath, version.Original(), "config.yaml")
		fh, err := os.OpenFile(configPath, os.O_RDONLY, 0644)
		if errors.Is(err, fs.ErrNotExist) {
			continue
		} else if err != nil {
			return configs, err
		}
		defer fh.Close()

		var fileContents []byte
		fileContents, err = io.ReadAll(fh)
		if err != nil {
			return configs, err
		}

		var config DirectivesConfig
		if err = yaml.Unmarshal(fileContents, &config); err != nil {
			return configs, err
		}

		configs = append(configs, config)
	}
	return configs, err
}

func versionSort(versionsRaw []fs.FileInfo) (versions []*version.Version, err error) {
	versions = make([]*version.Version, len(versionsRaw))
	for i, raw := range versionsRaw {
		var v *version.Version
		v, err = version.NewVersion(raw.Name())
		if err != nil {
			return
		}
		versions[i] = v
	}
	sort.Sort(version.Collection(versions))
	return
}

func checkDirectives(directivesConfigs []DirectivesConfig, currentVer, name string) (bool, int) {
	skip := false
	weight := 10

	for _, directivesConfig := range directivesConfigs {
		for _, directive := range directivesConfig.Directives {
			re, _ := regexp.Compile(directive.Regex)
			if !re.Match([]byte(currentVer)) {
				continue
			}
			for _, option := range directive.ScriptOptions {
				if option.Name == name {
					skip = option.Skip
					weight = option.Weight
				}
			}
		}
	}

	return skip, weight
}

func padVersion(version *version.Version) string {
	arr := version.Segments()
	out := fmt.Sprintf("v%02d.%02d.%02d", arr[0], arr[1], arr[2])
	if version.Prerelease() != "" {
		out += "-" + version.Prerelease()
	}
	return out
}

// Filters if script should be included by node role
// true if role mismatches filename's .tp or .cp
func roleFilter(fileName string, isControlPlane bool) bool {
	splits := strings.Split(fileName, ".")

	if (slices.Contains(splits, "cp") && !isControlPlane) ||
		(slices.Contains(splits, "tp") && isControlPlane) {
		return true
	}
	return false
}

func IsExecutable(filepath string) bool {
	stat, _ := os.Lstat(filepath)
	return stat.Mode()&0111 != 0
}