package distros

import (
	"embed"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"text/template"

	"github.com/rs/zerolog"
	"libguestfs.org/guestfs"

	"git.csclub.uwaterloo.ca/cloud/cloudbuild/pkg/config"
)

//go:embed resources
var res embed.FS

func getResource(filename string) []byte {
	data, err := res.ReadFile("resources/" + filename)
	if err != nil {
		panic(err)
	}
	return data
}

func getTemplateResource(filename string) *template.Template {
	tmpl, err := template.ParseFS(res, "resources/"+filename)
	if err != nil {
		panic(err)
	}
	return tmpl
}

// A TemplateManager downloads and modifies VM templates for a distro.
type TemplateManager struct {
	cfg    *config.Config
	logger *zerolog.Logger
	impl   IDistroSpecificTemplateManager
}

type ITemplateManager interface {
	IDistroSpecificTemplateManager
	// GetLatestVersion returns the version number and codename of the
	// latest version of particular OS (e.g. version="22.04", codename="jammy")
	GetLatestVersion() (version string, codename string, err error)
	// DownloadTemplate downloads a VM template for a given OS version and
	// codename, and returns the path to where it was downloaded
	DownloadTemplate(version, codename string) (path string, err error)
	// ModifyTemplate makes custom modifications to the downloaded VM template
	ModifyTemplate(filename string) error
}

// An IDistroSpecificTemplateManager performs the distro-specific tasks
// when modifying a VM template. It is used by the generic TemplateManager.
type IDistroSpecificTemplateManager interface {
	// UsesNetworkManager returns whether the distro uses NetworkManager
	UsesNetworkManager() bool
	// PerformDistroSpecificModifications is called after
	// performDistroAgnosticModifications to modify a template in a
	// distro-specific way.
	PerformDistroSpecificModifications(handle *guestfs.Guestfs) error
	// HasSELinuxEnabled returns whether SELinux is enabled for this distro
	HasSELinuxEnabled() bool
	// CommandToUpdatePackageCache returns the command which updates
	// (but does not upgrade) the packages for this distro, e.g.
	// "sudo apt update".
	CommandToUpdatePackageCache() []string
}

func (mgr *TemplateManager) DownloadTemplateGeneric(filename, url string) (path string, err error) {
	mgr.logger.Debug().Str("url", url).Msg("Downloading template")
	resp, err := http.Get(url)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	// TODO: download to somewhere in /tmp
	path = filename
	out, err := os.Create(path)
	if err != nil {
		return
	}
	defer out.Close()
	_, err = io.Copy(out, resp.Body)
	return
}

func debianCommandToUpdatePackageCache() []string {
	return []string{"sudo", "apt", "update"}
}

func fedoraCommandToUpdatePackageCache() []string {
	return []string{"sudo", "dnf", "makecache"}
}

func (mgr *TemplateManager) performDistroAgnosticModifications(handle *guestfs.Guestfs) (err error) {
	if err = mgr.setupIpv6Scripts(handle); err != nil {
		return
	}
	if err = mgr.setResolvConf(handle); err != nil {
		return
	}
	if err = mgr.setMotd(handle); err != nil {
		return
	}
	if err = mgr.updateSshdConfig(handle); err != nil {
		return
	}
	return
}

func (mgr *TemplateManager) ModifyTemplate(filename string) (err error) {
	handle, err := mgr.getGuestfsMountedHandle(filename)
	if err != nil {
		return
	}
	defer mgr.unmountAndCloseGuestfsHandle(handle)
	if err = mgr.createAugeasHandle(handle); err != nil {
		return
	}
	defer mgr.closeAugeasHandle(handle)
	if err = mgr.performDistroAgnosticModifications(handle); err != nil {
		return
	}
	if err = mgr.impl.PerformDistroSpecificModifications(handle); err != nil {
		return
	}
	if mgr.impl.HasSELinuxEnabled() {
		if err = mgr.selinuxRelabelDirectories(handle); err != nil {
			return
		}
	}
	return mgr.saveAugeasValues(handle)
}

func (mgr *TemplateManager) getGuestfsMountedHandle(filename string) (handle *guestfs.Guestfs, err error) {
	log := mgr.logger
	handle, err = guestfs.Create()
	if err != nil {
		return
	}
	log.Debug().Msg("Adding drive " + filename)
	if err = handle.Add_drive(filename, nil); err != nil {
		return
	}
	log.Debug().Msg("Launching VM")
	if err = handle.Launch(); err != nil {
		return
	}
	log.Debug().Msg("Inspecting OS")
	partitions, err := handle.Inspect_os()
	if err != nil {
		return
	}
	if len(partitions) != 1 {
		return nil, errors.New(fmt.Sprintf("Expected 1 root partition, found %d", len(partitions)))
	}
	rootPartition := partitions[0]
	log.Debug().Msg(fmt.Sprintf("Mounting root filesystem %s on /", rootPartition))
	if err = handle.Mount(rootPartition, "/"); err != nil {
		return
	}
	return
}

func (mgr *TemplateManager) unmountGuestfsDrive(handle *guestfs.Guestfs) {
	if err := handle.Umount("/", nil); err != nil {
		mgr.logger.Error().Err(err).Msg("")
	}
}

func (mgr *TemplateManager) unmountAndCloseGuestfsHandle(handle *guestfs.Guestfs) {
	mgr.unmountGuestfsDrive(handle)
	handle.Close()
}

func getSelinuxType(handle *guestfs.Guestfs) (selinuxType string, err error) {
	lines, err := handle.Grep("^SELINUXTYPE=", "/etc/selinux/config", nil)
	if err != nil {
		return
	}
	if len(lines) != 1 {
		err = errors.New(fmt.Sprintf("Expected 1 line containing SELINUXTYPE, found %d", len(lines)))
		return
	}
	selinuxType = strings.Split(lines[0], "=")[1]
	return
}

func getSelinuxDefaultSpecfile(handle *guestfs.Guestfs) (specfile string, err error) {
	selinuxType, err := getSelinuxType(handle)
	if err != nil {
		return
	}
	specfile = fmt.Sprintf("/etc/selinux/%s/contexts/files/file_contexts", selinuxType)
	return
}

func (mgr *TemplateManager) selinuxRelabel(handle *guestfs.Guestfs, specfile, dir string) (err error) {
	mgr.logger.Debug().
		Str("specfile", specfile).
		Str("path", dir).
		Msg("Relabeling SELinux security context")
	return handle.Selinux_relabel(specfile, dir, nil)
}

func (mgr *TemplateManager) selinuxRelabelDirectories(handle *guestfs.Guestfs) (err error) {
	specfile, err := getSelinuxDefaultSpecfile(handle)
	if err != nil {
		return
	}
	for _, dir := range []string{"/etc", "/var/lib/cloud"} {
		if err = mgr.selinuxRelabel(handle, specfile, dir); err != nil {
			return
		}
	}
	return
}

// setChronyOptions sets custom NTP server URLs in a chrony config file.
// It assumes that a line beginning with "pool" will already be present.
func (mgr *TemplateManager) setChronyOptions(handle *guestfs.Guestfs, path string) (err error) {
	oldLines, err := handle.Read_lines(path)
	if err != nil {
		return
	}
	snippet := string(getResource("chrony-snippet"))
	snippetLines := strings.Split(snippet, "\n")
	wroteSnippet := false
	newLines := make([]string, 0, len(oldLines)+len(snippetLines))
	for _, line := range oldLines {
		if strings.HasPrefix(line, "pool ") {
			newLines = append(newLines, "#"+line)
			if !wroteSnippet {
				newLines = append(newLines, snippetLines...)
				wroteSnippet = true
			}
		} else {
			newLines = append(newLines, line)
		}
	}
	newContent := strings.Join(newLines, "\n")
	mgr.logger.Debug().Msg("Writing new content to " + path)
	return handle.Write(path, []byte(newContent))
}

func (mgr *TemplateManager) setNetworkManagerOptions(handle *guestfs.Guestfs) (err error) {
	if err = handle.Mkdir_p("/etc/NetworkManager/conf.d"); err != nil {
		return
	}
	path := "/etc/NetworkManager/conf.d/99_csclub.conf"
	mgr.logger.Debug().Msg("Writing to " + path)
	return handle.Write(path, getResource("network-manager-snippet"))
}

func (mgr *TemplateManager) setupIpv6Scripts(handle *guestfs.Guestfs) (err error) {
	log := mgr.logger
	scripts := []string{"99_csclub_ipv6_addr.sh"}
	if mgr.impl.UsesNetworkManager() {
		scripts = append(scripts, "98_csclub_disable_nm_ipv6.sh")
	} else {
		scripts = append(scripts, "98_csclub_disable_ipv6_ra.sh")
	}
	scriptDir := "/var/lib/cloud/scripts/per-boot"
	if err = handle.Mkdir_p(scriptDir); err != nil {
		return
	}
	for _, filename := range scripts {
		path := scriptDir + "/" + filename
		log.Debug().Msg("Writing to " + path)
		if err = handle.Write(path, getResource(filename)); err != nil {
			return
		}
		if err = handle.Chmod(0755, path); err != nil {
			return
		}
	}
	sysctlPath := "/etc/sysctl.d/csclub.conf"
	log.Debug().Msg("Writing to " + sysctlPath)
	return handle.Write(sysctlPath, getResource("sysctl.conf"))
}

func (mgr *TemplateManager) setResolvConf(handle *guestfs.Guestfs) (err error) {
	mgr.logger.Debug().Msg("Writing to /etc/resolv.conf")
	if err = handle.Rm_f("/etc/resolv.conf"); err != nil {
		return
	}
	return handle.Write("/etc/resolv.conf", getResource("resolv.conf"))
}

func (mgr *TemplateManager) createAugeasHandle(handle *guestfs.Guestfs) (err error) {
	mgr.logger.Debug().Msg("Creating a new Augeas handle")
	return handle.Aug_init("/", 0)
}

func (mgr *TemplateManager) closeAugeasHandle(handle *guestfs.Guestfs) {
	if err := handle.Aug_close(); err != nil {
		mgr.logger.Error().Err(err).Msg("")
	}
}

func (mgr *TemplateManager) saveAugeasValues(handle *guestfs.Guestfs) error {
	mgr.logger.Debug().Msg("Saving Augeas values")
	return handle.Aug_save()
}

// requires an Augeas handle to be open
func (mgr *TemplateManager) setDhclientOptions(handle *guestfs.Guestfs) (err error) {
	mgr.logger.Debug().Msg("Retrieving dhclient request options")
	dhclientRequestOptionNodes, err := handle.Aug_ls("/files/etc/dhcp/dhclient.conf/request")
	if err != nil {
		return
	}
	optionsToRemove := map[string]bool{
		"domain-name":          true,
		"domain-name-servers":  true,
		"domain-search":        true,
		"dhcp6.name-servers":   true,
		"dhcp6.domain-search":  true,
		"dhcp6.fqdn":           true,
		"dhcp6.sntp-servers":   true,
		"ntp-servers":          true,
		"netbios-name-servers": true,
	}
	for _, optionPath := range dhclientRequestOptionNodes {
		// optionPath will look something like /files/etc/dhcp/dhclient.conf/request/5
		var optionValue string
		optionValue, err = handle.Aug_get(optionPath)
		if err != nil {
			return
		}
		if _, ok := optionsToRemove[optionValue]; !ok {
			continue
		}
		mgr.logger.Debug().Msg("Removing dhclient request option " + optionValue)
		if _, err = handle.Aug_rm(optionPath); err != nil {
			return
		}
	}
	return
}

func getNamedRegexGroup(re *regexp.Regexp, submatches []string, groupName string) string {
	var value string
	for i, subexpName := range re.SubexpNames() {
		if subexpName == groupName {
			value = submatches[i]
			break
		}
	}
	if value == "" {
		panic("Could not find regex group " + groupName)
	}
	return value
}

func getNumAugeasComments(handle *guestfs.Guestfs, parentNode string) (numComments int, err error) {
	keyPaths, err := handle.Aug_ls(parentNode)
	if err != nil {
		return 0, fmt.Errorf("aug_ls(%s) failed: %w", parentNode, err)
	}
	for _, keyPath := range keyPaths {
		key, err := handle.Aug_label(keyPath)
		if err != nil {
			return 0, fmt.Errorf("aug_label(%s) failed: %w", keyPath, err)
		}
		if strings.HasPrefix(key, "#comment") {
			numComments += 1
		}
	}
	return
}

func addAugeasComment(
	handle *guestfs.Guestfs, numExistingComments int, parentNode string, comment string,
) error {
	return handle.Aug_set(
		parentNode+fmt.Sprintf("/#comment[%d]", numExistingComments),
		comment,
	)
}

func (mgr *TemplateManager) commentOutAugeasPath(
	handle *guestfs.Guestfs, numExistingComments int, augeasPath string, commentedValue string,
) error {
	lastSlashIndex := strings.LastIndex(augeasPath, "/")
	if lastSlashIndex == -1 {
		// sanity check
		panic("augeasPath must have a slash")
	}
	parentNode := augeasPath[:lastSlashIndex]
	mgr.logger.Debug().Msg("Commenting out " + augeasPath)
	if _, err := handle.Aug_rm(augeasPath); err != nil {
		return fmt.Errorf("Could not remove %s: %w", augeasPath, err)
	}
	if err := addAugeasComment(handle, numExistingComments, parentNode, commentedValue); err != nil {
		return fmt.Errorf("Could not insert comment in %s: %w", parentNode, err)
	}
	return nil
}

// replaceYumRepoMirrorUrls comments out the metalink and mirrorlist URLs and
// replaces the baseurl URLs for each repo in /etc/yum.repos.d.
// It assumes that the baseurl will be present or commented.
//
// transformBaseUrl accepts a repo base URL and replaces the host (and
// optionally parts of the path) as appropriate. It should return the input
// if the URL should not be modified.
func (mgr *TemplateManager) replaceYumMirrorUrls(
	handle *guestfs.Guestfs, transformBaseurl func(string) string,
) error {
	repoPaths, err := handle.Aug_ls("/files/etc/yum.repos.d")
	if err != nil {
		return fmt.Errorf("Could not enumerate yum repos: %w", err)
	}
	// A repoPath looks like e.g. /files/etc/yum.repos.d/almalinux-appstream.repo
	for _, repoPath := range repoPaths {
		subrepoPaths, err := handle.Aug_ls(repoPath)
		if err != nil {
			return fmt.Errorf("Could not enumerate subrepos for %s: %w", repoPath, err)
		}
		// A subrepoPath looks like e.g. /files/etc/yum.repos.d/almalinux-appstream.repo/appstream
		for _, subrepoPath := range subrepoPaths {
			if err = mgr.replaceYumMirrorUrlsForSingleSubrepo(handle, transformBaseurl, subrepoPath); err != nil {
				return fmt.Errorf("Could not replace mirror URLs for %s: %w", subrepoPath, err)
			}
		}
	}
	return nil
}

func (mgr *TemplateManager) replaceYumMirrorUrlsForSingleSubrepo(
	handle *guestfs.Guestfs, transformBaseurl func(string) string, subrepoPath string,
) (err error) {
	log := mgr.logger
	// subrepoPath looks like e.g. /files/etc/yum.repos.d/almalinux-appstream.repo/appstream
	keyPaths, err := handle.Aug_ls(subrepoPath)
	if err != nil {
		return fmt.Errorf("Could not enumerate keys for %s: %w", subrepoPath, err)
	}
	var (
		numComments        int
		commentedBaseurl   string
		mirrorlist         string
		metalink           string
		uncommentedBaseurl string
	)
	// A keyPath looks like e.g. /files/etc/yum.repos.d/almalinux-appstream.repo/appstream/mirrorlist
	for _, keyPath := range keyPaths {
		// extract the last part of the node path
		key, err := handle.Aug_label(keyPath)
		if err != nil {
			return fmt.Errorf("Could not get label of %s: %w", keyPath, err)
		}
		value, err := handle.Aug_get(keyPath)
		if err != nil {
			return fmt.Errorf("Could not get %s: %w", keyPath, err)
		}
		if strings.HasPrefix(key, "#comment") {
			numComments += 1
			if strings.HasPrefix(value, "baseurl=") {
				commentedBaseurl = strings.Split(value, "=")[1]
			}
		} else if key == "mirrorlist" {
			mirrorlist = value
		} else if key == "metalink" {
			metalink = value
		} else if key == "baseurl" {
			uncommentedBaseurl = value
		}
	}
	var baseurl string
	if uncommentedBaseurl != "" {
		baseurl = transformBaseurl(uncommentedBaseurl)
	} else if commentedBaseurl != "" {
		baseurl = transformBaseurl(commentedBaseurl)
	}
	if baseurl == "" {
		return
	}
	log.Debug().Msg(fmt.Sprintf("Setting %s to %s", subrepoPath+"/baseurl", baseurl))
	if err = handle.Aug_set(subrepoPath+"/baseurl", baseurl); err != nil {
		return fmt.Errorf("Could not set baseurl for %s: %w", subrepoPath, err)
	}
	if mirrorlist != "" {
		if err = mgr.commentOutAugeasPath(handle, numComments, subrepoPath+"/mirrorlist", "mirrorlist="+mirrorlist); err != nil {
			return
		}
		numComments += 1
	}
	if metalink != "" {
		if err = mgr.commentOutAugeasPath(handle, numComments, subrepoPath+"/metalink", "metalink="+metalink); err != nil {
			return
		}
		numComments += 1
	}
	return
}

// requires an Augeas handle to be open
func (mgr *TemplateManager) replaceDebianMirrorUrls(handle *guestfs.Guestfs) (err error) {
	log := mgr.logger
	// Some Augeas nodes under /files/etc/apt/sources.list are comments,
	// so we use /*/uri to make sure that we only get the actual entries
	sourcesListEntries, err := handle.Aug_match("/files/etc/apt/sources.list/*/uri")
	if err != nil {
		return
	}
	for _, uriPath := range sourcesListEntries {
		var (
			uriValue  string
			parsedUrl *url.URL
			distValue string
			typeValue string
		)
		if uriValue, err = handle.Aug_get(uriPath); err != nil {
			return
		}
		parsedUrl, err = url.Parse(uriValue)
		if err != nil {
			return
		}
		parsedUrl.Host = mgr.cfg.MirrorHost
		newUriValue := parsedUrl.String()
		// strip off the "/uri" from the node path
		entryPath := uriPath[:len(uriPath)-4]
		typePath := entryPath + "/type"
		if typeValue, err = handle.Aug_get(typePath); err != nil {
			return
		}
		distPath := entryPath + "/distribution"
		if distValue, err = handle.Aug_get(distPath); err != nil {
			return
		}
		if typeValue == "deb-src" {
			log.Debug().
				Str("URL", uriValue).
				Str("distribution", distValue).
				Msg("Removing deb-src entry")
			if _, err = handle.Aug_rm(entryPath); err != nil {
				return
			}
			continue
		}
		if uriValue == newUriValue {
			continue
		}
		log.Debug().
			Str("distribution", distValue).
			Str("oldURL", uriValue).
			Str("newURL", newUriValue).
			Msg("Replacing URL in sources.list")
		if err = handle.Aug_set(uriPath, newUriValue); err != nil {
			return
		}
	}
	return
}

// requires an Augeas handle to be open
func (mgr *TemplateManager) updateSshdConfig(handle *guestfs.Guestfs) error {
	mgr.logger.Debug().Msg("Setting PrintLastLog=no in sshd_config")
	return handle.Aug_set("/files/etc/ssh/sshd_config/PrintLastLog", "no")
}

func (mgr *TemplateManager) setTimesyncdConf(handle *guestfs.Guestfs) (err error) {
	mgr.logger.Debug().Msg("Writing custom timesyncd.conf")
	if err = handle.Mkdir_p("/etc/systemd/timesyncd.conf.d"); err != nil {
		return
	}
	return handle.Write("/etc/systemd/timesyncd.conf.d/csclub.conf", getResource("timesyncd.conf"))
}

func (mgr *TemplateManager) setMotd(handle *guestfs.Guestfs) error {
	mgr.logger.Debug().Msg("Writing to /etc/motd")
	return handle.Write("/etc/motd", getResource("motd"))
}