package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"git.iamthefij.com/iamthefij/slog"
	dockerTypes "github.com/docker/docker/api/types"
	dockerClient "github.com/docker/docker/client"
)

var (
	// defaultRegistryBaseURL is the base URL of the docker registry
	defaultRegistryBaseURL = "https://hub.docker.com"
	// maxPages is the max number of pages to fetch from docker registry results
	maxPages = 10

	// Regexp used to extract tag information
	tagRegexp = regexp.MustCompile(`(.*):[vV]{0,1}([0-9.]+)(-(.*)){0,1}`)
	// version of tag checker
	version = "dev"
)

// ImageTag is wraps an image and tag values for a container
type ImageTag struct {
	ImageTag     string
	Registry     string
	Image        string
	TagDesc      string
	Version      string
	VersionParts []int
}

// IsComparable will return true if to images share the same base, description, and version resolution
func (thisTag ImageTag) IsComparable(otherTag ImageTag) bool {
	return thisTag.Image == otherTag.Image && thisTag.TagDesc == otherTag.TagDesc && len(thisTag.VersionParts) == len(otherTag.VersionParts)

}

// IsNewerThan will return true if two tags are comparable and this tag is newer than the one passed in
func (thisTag ImageTag) IsNewerThan(otherTag ImageTag) bool {
	return thisTag.CompareTo(otherTag) == 1
}

// CompareTo compares two ImageTags. It will return 0 if they are not-comparable or equal 1 if newer and -1 if less
func (thisTag ImageTag) CompareTo(otherTag ImageTag) int {
	if !thisTag.IsComparable(otherTag) {
		return 0
	}

	for i := range thisTag.VersionParts {
		if thisTag.VersionParts[i] > otherTag.VersionParts[i] {
			return 1
		} else if thisTag.VersionParts[i] < otherTag.VersionParts[i] {
			return -1
		}
	}

	// Everything must be equal
	return 0
}

// ParseImageTag parses an image and tag name to a struct
func ParseImageTag(imageTag string) (ImageTag, error) {
	results := tagRegexp.FindStringSubmatch(imageTag)
	if results == nil || results[0] == "" {
		return ImageTag{}, fmt.Errorf("could not recognize versions in %s", imageTag)
	}

	// Extract image name with repo
	image := results[1]
	registry := ""
	switch strings.Count(image, "/") {
	case 0:
		image = "library/" + image
	case 2:
		parts := strings.Split(image, "/")
		if parts[0] != "docker.io" {
			registry = parts[0]
		}
		image = strings.Join(parts[1:], "/")
	}

	// Extract version number
	version := results[2]
	versionParts := []int{}
	var verPart int
	var err error
	for _, v := range strings.Split(version, ".") {
		verPart, err = strconv.Atoi(v)
		if err != nil {
			return ImageTag{}, err
		}
		versionParts = append(versionParts, verPart)
	}

	return ImageTag{
		ImageTag:     imageTag,
		Image:        image,
		Registry:     registry,
		Version:      version,
		VersionParts: versionParts,
		TagDesc:      results[4],
	}, nil
}

func getJSON(url string, response interface{}) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	// defer func() { _ = resp.Body.Close() }()
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	err = decoder.Decode(response)
	if err != nil {
		return err
	}

	return nil
}

func listTags(current ImageTag) ([]ImageTag, error) {
	var err error
	results := []ImageTag{}

	type tagsResponse struct {
		Count    int
		Next     string
		Previous string
		Results  []struct {
			ID          int
			LastUpdated string `json:"last_updated"`
			Name        string
		}
	}

	registryBaseURL := defaultRegistryBaseURL
	if current.Registry != "" {
		registryBaseURL = fmt.Sprintf("https://%s", current.Registry)
	}
	url := fmt.Sprintf("%s/v2/repositories/%s/tags", registryBaseURL, current.Image)
	pageCount := 0
	var response tagsResponse
	var newTag ImageTag
	for url != "" && pageCount <= maxPages {
		err = getJSON(url, &response)
		if err != nil {
			return results, err
		}

		for _, tag := range response.Results {
			newTag, err = ParseImageTag(fmt.Sprintf("%s:%s", current.Image, tag.Name))
			if err == nil {
				results = append(results, newTag)
			}
		}
		url = response.Next
		pageCount++
	}

	return results, nil
}

func getNewerTags(current ImageTag) ([]ImageTag, error) {
	newerTags := []ImageTag{}
	tags, err := listTags(current)
	if err != nil {
		return newerTags, err
	}
	for _, tag := range tags {
		if tag.IsNewerThan(current) {
			newerTags = append(newerTags, tag)
		}
	}
	// Sort tags with newest first
	sort.Slice(newerTags, func(i, j int) bool { return newerTags[i].CompareTo(newerTags[j]) == 1 })
	return newerTags, nil
}

func main() {
	flag.StringVar(&defaultRegistryBaseURL, "registry-url", defaultRegistryBaseURL, "base url of the registry you want to check against")
	flag.IntVar(&maxPages, "max-pages", maxPages, "max number of pages to retrieve from registry")
	var showVersion = flag.Bool("version", false, "display the version and exit")
	flag.BoolVar(&slog.DebugLevel, "debug", false, "show debug logs")
	flag.Parse()

	// Print version if asked
	if *showVersion {
		fmt.Println("version:", version)
		os.Exit(0)
	}
	dockerClient, err := dockerClient.NewClientWithOpts(dockerClient.FromEnv)
	if err != nil {
		fmt.Println("Could not initialize docker client")
		panic(err)
	}

	containers, err := dockerClient.ContainerList(context.Background(), dockerTypes.ContainerListOptions{})
	if err != nil {
		fmt.Println("Could list container from docker client")
		panic(err)
	}

	hasUpdate := false
	images := map[string]bool{}
	for _, container := range containers {
		images[container.Image] = true
	}
	for image := range images {
		slog.Debug("[%s] Checking for updates...", image)
		it, err := ParseImageTag(image)
		if err != nil {
			slog.Debug("[%s] Can't parse tag: %v", image, err)
			continue
		}
		newerTags, err := getNewerTags(it)
		slog.PanicOnErr(err, "[%s] failed getting new tags", image)
		if len(newerTags) == 0 {
			slog.Debug("[%s] No newer versions found", image)
			continue
		}
		hasUpdate = true
		if slog.DebugLevel {
			slog.Info("[%s] Newer version found! Recommended update to %s", image, newerTags[0].ImageTag)
		} else {
			fmt.Printf("[%s] Newer version found! Recommended update to %s\n", image, newerTags[0].ImageTag)
		}
	}

	if hasUpdate {
		os.Exit(10)
	}
}