package docker

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
)

var TooManyRequests = fmt.Errorf("toomanyrequests: You have reached your pull rate limit")
var ManifestUnknown = fmt.Errorf("manifest unknown")

var (
	manifestAccept     = "application/vnd.docker.distribution.manifest.v2+json"
	manifestAcceptList = "application/vnd.docker.distribution.manifest.list.v2+json"
	manifestAccepts    = []string{
		manifestAccept,
		manifestAcceptList,
		// Some registries like GitHub have a strict verification on Accept header.
		// See https://github.com/opencontainers/image-spec/blob/083f635f2b04f7f340685a1e0c4393a0feec3679/specs-go/v1/mediatype.go#L28.
		"application/vnd.oci.image.index.v1+json",
	}
)

type ImageChecker struct {
	apiURL string
	token  string
	image  string
}

func NewImageCheckerFromImage(image string) (*ImageChecker, error) {
	var err error
	image, err = completeImage(image)
	if err != nil {
		return nil, err
	}
	hub := strings.Split(image, "/")[0]
	repo := getRepoFromImage(image)
	var requestURL string

	switch hub {
	case "docker.io":
		requestURL = "registry-1.docker.io"
	default:
		requestURL = hub
	}
	url := fmt.Sprintf("https://%s/v2", requestURL)
	response, err := http.Get(url + "/")
	if err != nil {
		return nil, fmt.Errorf("request %s failed: %v", url, err)
	}
	auth := response.Header.Get("Www-Authenticate")
	if auth == "" {
		return &ImageChecker{
			apiURL: url,
			image:  image,
		}, nil
	}
	tokenURL := extractBearerRealm(auth)
	if tokenURL == "" {
		return nil, fmt.Errorf("unable to get the token URL")
	}
	service := extractService(auth)
	if service == "" {
		return nil, fmt.Errorf("unable to get the service")
	}
	token, err := getToken(tokenURL, service, repo)
	if err != nil {
		return nil, err
	}
	return &ImageChecker{
		apiURL: url,
		token:  token,
		image:  image,
	}, nil
}

func (i *ImageChecker) IsImageExist() (bool, error) {
	_, err := i.getManifestFromImage(i.image)
	if err != nil {
		if err == ManifestUnknown {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

func (i *ImageChecker) IsPlatformSupport(arch string) (bool, error) {
	repo := getRepoFromImage(i.image)
	manifest, err := i.getManifestFromImage(i.image)
	if err != nil {
		return false, err
	}
	switch manifest["mediaType"] {
	case manifestAccept:
		config, ok := manifest["config"].(map[string]interface{})
		if !ok {
			return false, fmt.Errorf("unknown config: %+v", manifest["config"])
		}
		digest, ok := config["digest"].(string)
		if !ok {
			return false, fmt.Errorf("unknown digest: %+v", digest)
		}
		blobs, err := i.getBlobs(repo, digest)
		if err != nil {
			return false, err
		}
		return blobs["architecture"].(string) == arch, nil
	case manifestAcceptList:
		manifests := manifest["manifests"].([]interface{})
		for _, manifest := range manifests {
			platform, ok := manifest.(map[string]interface{})["platform"].(map[string]interface{})
			if !ok {
				return false, fmt.Errorf("unknown platform structure")
			}
			if platform["architecture"] == arch {
				return true, nil
			}
		}
	default:
		return false, fmt.Errorf("unknown media type: %s", manifest["mediaType"])
	}
	return false, nil
}

func (i *ImageChecker) getManifestFromImage(image string) (map[string]interface{}, error) {
	repo := getRepoFromImage(image)
	tag := getRepoTagFromImage(image)

	client := &http.Client{}
	requestURL := fmt.Sprintf("%s/%s/manifests/%s", i.apiURL, repo, tag)
	req, _ := http.NewRequest(http.MethodGet, requestURL, nil)

	if i.token != "" {
		req.Header.Set("Authorization", "Bearer "+i.token)
	}
	for _, accept := range manifestAccepts {
		req.Header.Add("Accept", accept)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode == 404 {
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
		type Error struct {
			Code string `json:"code"`
		}
		type ManifestStatus struct {
			Errors []Error `json:"errors"`
		}
		var status ManifestStatus
		if err := json.Unmarshal(b, &status); err != nil {
			return nil, err
		}
		if len(status.Errors) == 1 {
			switch status.Errors[0].Code {
			case "MANIFEST_UNKNOWN", "NOT_FOUND":
				return nil, ManifestUnknown
			}
			return nil, ManifestUnknown
		}
		return nil, fmt.Errorf("404 Not Found")
	}
	if resp.StatusCode == http.StatusTooManyRequests {
		return nil, TooManyRequests
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("expect 200 OK, got %s", resp.Status)
	}
	resBody := make(map[string]interface{})
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if err := json.Unmarshal(b, &resBody); err != nil {
		return nil, err
	}
	return resBody, nil
}

func (i *ImageChecker) getBlobs(repo string, sha256 string) (map[string]interface{}, error) {
	client := &http.Client{}
	requestURL := fmt.Sprintf("%s/%s/blobs/%s", i.apiURL, repo, sha256)
	req, _ := http.NewRequest(http.MethodGet, requestURL, nil)
	if i.token != "" {
		req.Header.Set("Authorization", "Bearer "+i.token)
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode == 404 {
		return nil, fmt.Errorf("404 Not Found")
	}
	if resp.StatusCode == http.StatusTooManyRequests {
		return nil, TooManyRequests
	}
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("expect 200 OK, got %s", resp.Status)
	}
	resBody := make(map[string]interface{})
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if err := json.Unmarshal(b, &resBody); err != nil {
		return nil, err
	}
	return resBody, nil
}

// https://hub.pingcap.net/service/token?account=admin&scope=repository:qa/go-tpc:push,pull&service=harbor-registry
func getToken(url string, service string, repo string) (string, error) {
	resp, err := http.Get(fmt.Sprintf("%s?account=admin&scope=repository:%s:push,pull&service=%s", url, repo, service))
	if err != nil {
		return "", err
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	type Resp struct {
		Token string `json:"token"`
	}
	var r Resp
	if err := json.Unmarshal(body, &r); err != nil {
		return "", fmt.Errorf("parse token error: %v", err)
	}
	return r.Token, nil
}

// docker.io/pingcap/tidb -> pingcap/tidb
// hub.pingcap.net/pingcap/tidb:tag -> pingcap/tidb
func getRepoFromImage(image string) string {
	// TODO: polish
	repo := strings.Split(strings.SplitN(image, "/", 2)[1], ":")[0]
	return repo
}

func getRepoTagFromImage(image string) string {
	parts := strings.Split(image, ":")
	if len(parts) == 1 {
		return "latest"
	}
	return parts[1]
}

// centos
// pingcap/tidb
// docker.io/pingcap/tidb
// hub.pingcap.net/pingcap/tidb
// ghcr.io/homebrew/core/redis
func completeImage(image string) (string, error) {
	parts := strings.Split(image, "/")
	switch len(parts) {
	case 1:
		return fmt.Sprintf("docker.io/library/%s", parts[0]), nil
	case 2:
		return fmt.Sprintf("docker.io/%s/%s", parts[0], parts[1]), nil
	case 3, 4:
		return image, nil
	}
	return "", fmt.Errorf("cannot process the image: %s", image)
}

func extractBearerRealm(auth string) string {
	var tokenURL string
	// Bearer realm="https://hub.pingcap.net/service/token",service="harbor-registry",error="invalid_token"
	for _, item := range strings.Split(auth, ",") {
		if strings.HasPrefix(item, "Bearer realm=") {
			tokenURL = strings.Split(item, `"`)[1]
			break
		}
	}
	return tokenURL
}

func extractService(auth string) string {
	var service string
	// Bearer realm="https://hub.pingcap.net/service/token",service="harbor-registry",error="invalid_token"
	for _, item := range strings.Split(auth, ",") {
		if strings.HasPrefix(item, "service=") {
			service = strings.Split(item, `"`)[1]
			break
		}
	}
	return service
}