package docker

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/textproto"
	"net/url"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/reference"
	"github.com/containerd/containerd/remotes"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"golang.org/x/net/context/ctxhttp"
)

var (
	// ErrNoToken is returned if a request is successful but the body does not
	// contain an authorization token.
	ErrNoToken = errors.New("authorization server did not include a token in the response")

	// ErrInvalidAuthorization is used when credentials are passed to a server but
	// those credentials are rejected.
	ErrInvalidAuthorization = errors.New("authorization failed")
)

type dockerResolver struct {
	credentials func(string) (string, string, error)
	plainHTTP   bool
	client      *http.Client
}

// ResolverOptions are used to configured a new Docker register resolver
type ResolverOptions struct {
	// Credentials provides username and secret given a host.
	// If username is empty but a secret is given, that secret
	// is interpretted as a long lived token.
	Credentials func(string) (string, string, error)

	// PlainHTTP specifies to use plain http and not https
	PlainHTTP bool

	// Client is the http client to used when making registry requests
	Client *http.Client
}

// NewResolver returns a new resolver to a Docker registry
func NewResolver(options ResolverOptions) remotes.Resolver {
	return &dockerResolver{
		credentials: options.Credentials,
		plainHTTP:   options.PlainHTTP,
		client:      options.Client,
	}
}

var _ remotes.Resolver = &dockerResolver{}

func (r *dockerResolver) Resolve(ctx context.Context, ref string) (string, ocispec.Descriptor, remotes.Fetcher, error) {
	refspec, err := reference.Parse(ref)
	if err != nil {
		return "", ocispec.Descriptor{}, nil, err
	}

	var (
		base             url.URL
		username, secret string
	)

	host := refspec.Hostname()
	base.Scheme = "https"

	if host == "docker.io" {
		base.Host = "registry-1.docker.io"
	} else {
		base.Host = host

		if r.plainHTTP || strings.HasPrefix(host, "localhost:") {
			base.Scheme = "http"
		}
	}

	if r.credentials != nil {
		username, secret, err = r.credentials(base.Host)
		if err != nil {
			return "", ocispec.Descriptor{}, nil, err
		}
	}

	prefix := strings.TrimPrefix(refspec.Locator, host+"/")
	base.Path = path.Join("/v2", prefix)

	fetcher := &dockerFetcher{
		base:     base,
		client:   r.client,
		username: username,
		secret:   secret,
	}

	var (
		urls []string
		dgst = refspec.Digest()
	)

	if dgst != "" {
		if err := dgst.Validate(); err != nil {
			// need to fail here, since we can't actually resolve the invalid
			// digest.
			return "", ocispec.Descriptor{}, nil, err
		}

		// turns out, we have a valid digest, make a url.
		urls = append(urls, fetcher.url("manifests", dgst.String()))
	} else {
		urls = append(urls, fetcher.url("manifests", refspec.Object))
	}

	// fallback to blobs on not found.
	urls = append(urls, fetcher.url("blobs", dgst.String()))

	for _, u := range urls {
		req, err := http.NewRequest(http.MethodHead, u, nil)
		if err != nil {
			return "", ocispec.Descriptor{}, nil, err
		}

		// set headers for all the types we support for resolution.
		req.Header.Set("Accept", strings.Join([]string{
			images.MediaTypeDockerSchema2Manifest,
			images.MediaTypeDockerSchema2ManifestList,
			ocispec.MediaTypeImageManifest,
			ocispec.MediaTypeImageIndex, "*"}, ", "))

		log.G(ctx).Debug("resolving")
		resp, err := fetcher.doRequest(ctx, req)
		if err != nil {
			return "", ocispec.Descriptor{}, nil, err
		}
		resp.Body.Close() // don't care about body contents.

		if resp.StatusCode > 299 {
			if resp.StatusCode == http.StatusNotFound {
				continue
			}
			return "", ocispec.Descriptor{}, nil, errors.Errorf("unexpected status code %v: %v", u, resp.Status)
		}

		// this is the only point at which we trust the registry. we use the
		// content headers to assemble a descriptor for the name. when this becomes
		// more robust, we mostly get this information from a secure trust store.
		dgstHeader := digest.Digest(resp.Header.Get("Docker-Content-Digest"))

		if dgstHeader != "" {
			if err := dgstHeader.Validate(); err != nil {
				return "", ocispec.Descriptor{}, nil, errors.Wrapf(err, "%q in header not a valid digest", dgstHeader)
			}
			dgst = dgstHeader
		}

		if dgst == "" {
			return "", ocispec.Descriptor{}, nil, errors.Errorf("could not resolve digest for %v", ref)
		}

		var (
			size       int64
			sizeHeader = resp.Header.Get("Content-Length")
		)

		size, err = strconv.ParseInt(sizeHeader, 10, 64)
		if err != nil {

			return "", ocispec.Descriptor{}, nil, errors.Wrapf(err, "invalid size header: %q", sizeHeader)
		}
		if size < 0 {
			return "", ocispec.Descriptor{}, nil, errors.Errorf("%q in header not a valid size", sizeHeader)
		}

		desc := ocispec.Descriptor{
			Digest:    dgst,
			MediaType: resp.Header.Get("Content-Type"), // need to strip disposition?
			Size:      size,
		}

		log.G(ctx).WithField("desc.digest", desc.Digest).Debug("resolved")
		return ref, desc, fetcher, nil
	}

	return "", ocispec.Descriptor{}, nil, errors.Errorf("%v not found", ref)
}

type dockerFetcher struct {
	base  url.URL
	token string

	client   *http.Client
	useBasic bool
	username string
	secret   string
}

func (r *dockerFetcher) Fetch(ctx context.Context, desc ocispec.Descriptor) (io.ReadCloser, error) {
	ctx = log.WithLogger(ctx, log.G(ctx).WithFields(
		logrus.Fields{
			"base":   r.base.String(),
			"digest": desc.Digest,
		},
	))

	paths, err := getV2URLPaths(desc)
	if err != nil {
		return nil, err
	}

	for _, path := range paths {
		u := r.url(path)

		req, err := http.NewRequest(http.MethodGet, u, nil)
		if err != nil {
			return nil, err
		}

		req.Header.Set("Accept", strings.Join([]string{desc.MediaType, `*`}, ", "))
		resp, err := r.doRequest(ctx, req)
		if err != nil {
			return nil, err
		}

		if resp.StatusCode > 299 {
			if resp.StatusCode == http.StatusNotFound {
				continue // try one of the other urls.
			}
			resp.Body.Close()
			return nil, errors.Errorf("unexpected status code %v: %v", u, resp.Status)
		}

		return resp.Body, nil
	}

	return nil, errors.New("not found")
}

func (r *dockerFetcher) url(ps ...string) string {
	url := r.base
	url.Path = path.Join(url.Path, path.Join(ps...))
	return url.String()
}

func (r *dockerFetcher) doRequest(ctx context.Context, req *http.Request) (*http.Response, error) {
	return r.doRequestWithRetries(ctx, req, nil)
}

func (r *dockerFetcher) doRequestWithRetries(ctx context.Context, req *http.Request, responses []*http.Response) (*http.Response, error) {
	ctx = log.WithLogger(ctx, log.G(ctx).WithField("url", req.URL.String()))
	log.G(ctx).WithField("request.headers", req.Header).Debug("fetch content")
	r.authorize(req)
	resp, err := ctxhttp.Do(ctx, r.client, req)
	if err != nil {
		return nil, err
	}
	log.G(ctx).WithFields(logrus.Fields{
		"status":           resp.Status,
		"response.headers": resp.Header,
	}).Debug("fetch response received")

	responses = append(responses, resp)
	req, err = r.retryRequest(ctx, req, responses)
	if err != nil {
		return nil, err
	}
	if req != nil {
		return r.doRequestWithRetries(ctx, req, responses)
	}
	return resp, err
}

func (r *dockerFetcher) authorize(req *http.Request) {
	if r.useBasic {
		req.SetBasicAuth(r.username, r.secret)
	} else if r.token != "" {
		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", r.token))
	}
}

func (r *dockerFetcher) retryRequest(ctx context.Context, req *http.Request, responses []*http.Response) (*http.Request, error) {
	if len(responses) > 5 {
		return nil, nil
	}
	last := responses[len(responses)-1]
	if last.StatusCode == http.StatusUnauthorized {
		log.G(ctx).WithField("header", last.Header.Get("WWW-Authenticate")).Debug("Unauthorized")
		for _, c := range parseAuthHeader(last.Header) {
			if c.scheme == bearerAuth {
				if errStr := c.parameters["error"]; errStr != "" {
					// TODO: handle expired case
					return nil, errors.Wrapf(ErrInvalidAuthorization, "server message: %s", errStr)
				}
				if err := r.setTokenAuth(ctx, c.parameters); err != nil {
					return nil, err
				}
				return req, nil
			} else if c.scheme == basicAuth {
				if r.username != "" && r.secret != "" {
					r.useBasic = true
				}
				return req, nil
			}
		}
		return nil, nil
	} else if last.StatusCode == http.StatusMethodNotAllowed && req.Method == http.MethodHead {
		// Support registries which have not properly implemented the HEAD method for
		// manifests endpoint
		if strings.Contains(req.URL.Path, "/manifests/") {
			// TODO: copy request?
			req.Method = http.MethodGet
			return req, nil
		}
	}

	// TODO: Handle 50x errors accounting for attempt history
	return nil, nil
}

func isManifestAccept(h http.Header) bool {
	for _, ah := range h[textproto.CanonicalMIMEHeaderKey("Accept")] {
		switch ah {
		case images.MediaTypeDockerSchema2Manifest:
			fallthrough
		case images.MediaTypeDockerSchema2ManifestList:
			fallthrough
		case ocispec.MediaTypeImageManifest:
			fallthrough
		case ocispec.MediaTypeImageIndex:
			return true
		}
	}
	return false
}

func (r *dockerFetcher) setTokenAuth(ctx context.Context, params map[string]string) error {
	realm, ok := params["realm"]
	if !ok {
		return errors.New("no realm specified for token auth challenge")
	}

	realmURL, err := url.Parse(realm)
	if err != nil {
		return fmt.Errorf("invalid token auth challenge realm: %s", err)
	}

	to := tokenOptions{
		realm:   realmURL.String(),
		service: params["service"],
	}

	scope, ok := params["scope"]
	if !ok {
		return errors.Errorf("no scope specified for token auth challenge")
	}

	// TODO: Get added scopes from context
	to.scopes = []string{scope}

	if r.secret != "" {
		// Credential information is provided, use oauth POST endpoint
		r.token, err = r.fetchTokenWithOAuth(ctx, to)
		if err != nil {
			return errors.Wrap(err, "failed to fetch oauth token")
		}
	} else {
		// Do request anonymously
		r.token, err = r.getToken(ctx, to)
		if err != nil {
			return errors.Wrap(err, "failed to fetch anonymous token")
		}
	}

	return nil
}

// getV2URLPaths generates the candidate urls paths for the object based on the
// set of hints and the provided object id. URLs are returned in the order of
// most to least likely succeed.
func getV2URLPaths(desc ocispec.Descriptor) ([]string, error) {
	var urls []string

	switch desc.MediaType {
	case images.MediaTypeDockerSchema2Manifest, images.MediaTypeDockerSchema2ManifestList,
		ocispec.MediaTypeImageManifest, ocispec.MediaTypeImageIndex:
		urls = append(urls, path.Join("manifests", desc.Digest.String()))
	}

	// always fallback to attempting to get the object out of the blobs store.
	urls = append(urls, path.Join("blobs", desc.Digest.String()))

	return urls, nil
}

type tokenOptions struct {
	realm   string
	service string
	scopes  []string
}

type postTokenResponse struct {
	AccessToken  string    `json:"access_token"`
	RefreshToken string    `json:"refresh_token"`
	ExpiresIn    int       `json:"expires_in"`
	IssuedAt     time.Time `json:"issued_at"`
	Scope        string    `json:"scope"`
}

func (r *dockerFetcher) fetchTokenWithOAuth(ctx context.Context, to tokenOptions) (string, error) {
	form := url.Values{}
	form.Set("scope", strings.Join(to.scopes, " "))
	form.Set("service", to.service)
	// TODO: Allow setting client_id
	form.Set("client_id", "containerd-dist-tool")

	if r.username == "" {
		form.Set("grant_type", "refresh_token")
		form.Set("refresh_token", r.secret)
	} else {
		form.Set("grant_type", "password")
		form.Set("username", r.username)
		form.Set("password", r.secret)
	}

	resp, err := ctxhttp.PostForm(ctx, r.client, to.realm, form)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode == 405 && r.username != "" {
		// It would be nice if registries would implement the specifications
		return r.getToken(ctx, to)
	} else if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		b, _ := ioutil.ReadAll(resp.Body)
		log.G(ctx).WithFields(logrus.Fields{
			"status": resp.Status,
			"body":   string(b),
		}).Debugf("token request failed")
		// TODO: handle error body and write debug output
		return "", errors.Errorf("unexpected status: %s", resp.Status)
	}

	decoder := json.NewDecoder(resp.Body)

	var tr postTokenResponse
	if err = decoder.Decode(&tr); err != nil {
		return "", fmt.Errorf("unable to decode token response: %s", err)
	}

	return tr.AccessToken, nil
}

type getTokenResponse struct {
	Token        string    `json:"token"`
	AccessToken  string    `json:"access_token"`
	ExpiresIn    int       `json:"expires_in"`
	IssuedAt     time.Time `json:"issued_at"`
	RefreshToken string    `json:"refresh_token"`
}

// getToken fetches a token using a GET request
func (r *dockerFetcher) getToken(ctx context.Context, to tokenOptions) (string, error) {
	req, err := http.NewRequest("GET", to.realm, nil)
	if err != nil {
		return "", err
	}

	reqParams := req.URL.Query()

	if to.service != "" {
		reqParams.Add("service", to.service)
	}

	for _, scope := range to.scopes {
		reqParams.Add("scope", scope)
	}

	if r.secret != "" {
		req.SetBasicAuth(r.username, r.secret)
	}

	req.URL.RawQuery = reqParams.Encode()

	resp, err := ctxhttp.Do(ctx, r.client, req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		// TODO: handle error body and write debug output
		return "", errors.Errorf("unexpected status: %s", resp.Status)
	}

	decoder := json.NewDecoder(resp.Body)

	var tr getTokenResponse
	if err = decoder.Decode(&tr); err != nil {
		return "", fmt.Errorf("unable to decode token response: %s", err)
	}

	// `access_token` is equivalent to `token` and if both are specified
	// the choice is undefined.  Canonicalize `access_token` by sticking
	// things in `token`.
	if tr.AccessToken != "" {
		tr.Token = tr.AccessToken
	}

	if tr.Token == "" {
		return "", ErrNoToken
	}

	return tr.Token, nil
}