// Package docker implements the OCI/Docker Registry v2 pull-through cache.
// The first path segment after /v2/ names the virtual repo:
//
//	/v2/<vrepo>/<name...>/manifests/<ref>
//	/v2/<vrepo>/<name...>/blobs/<digest>
//
// Multiple upstreams may be configured as an ordered fallback list, e.g. a
// chained tiny-proxy first and the registry itself second.
package docker

import (
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"tiny-proxy/internal/auth"
	"tiny-proxy/internal/config"
	"tiny-proxy/internal/metrics"
	"tiny-proxy/internal/proxycache"
	"tiny-proxy/internal/upstream"
)

// upstreamEntry pairs one upstream's config with its HTTP client.
type upstreamEntry struct {
	conf   config.UpstreamConf
	client *http.Client
}

type repo struct {
	cfg       config.VRepo
	upstreams []upstreamEntry
}

type Handler struct {
	repos map[string]*repo
	cache *proxycache.Cache
}

// NewHandler builds a docker handler. clients[vrepo] must have one client
// per entry of vrepo.AllUpstreams(), in the same order.
func NewHandler(repos []config.VRepo, clients map[string][]*http.Client, cache *proxycache.Cache) *Handler {
	h := &Handler{repos: map[string]*repo{}, cache: cache}
	for _, rc := range repos {
		if rc.Type != config.VRepoDocker {
			continue
		}
		rp := &repo{cfg: rc}
		for i, u := range rc.AllUpstreams() {
			rp.upstreams = append(rp.upstreams, upstreamEntry{conf: u, client: clients[rc.Name][i]})
		}
		h.repos[rc.Name] = rp
	}
	return h
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	rest := strings.TrimPrefix(r.URL.Path, "/v2/")
	if rest == "" {
		// API version ping; docker clients probe this before anything else.
		w.Header().Set("Docker-Distribution-Api-Version", "registry/2.0")
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte("{}\n"))
		return
	}
	if r.Method != http.MethodGet && r.Method != http.MethodHead {
		http.Error(w, "read-only registry", http.StatusMethodNotAllowed)
		return
	}

	vrepo, name, kind, ref, ok := parsePath(rest)
	if !ok {
		http.Error(w, "bad path", http.StatusBadRequest)
		return
	}
	rp, exists := h.repos[vrepo]
	if !exists {
		http.Error(w, "unknown repository", http.StatusNotFound)
		return
	}
	if !auth.Check(rp.cfg.Auth, r) {
		auth.Unauthorized(w)
		return
	}
	// Configured and authorized, so it is safe as a metrics label; the cache
	// reads it back from the context to attribute its hits and misses.
	metrics.MarkVRepo(r.Context(), vrepo)

	var key string
	var ttl = rp.cfg.MutableTTL()
	if kind == "manifests" {
		key = "docker/" + vrepo + "/manifests/" + name + "/" + ref
		if strings.Contains(ref, ":") { // digest refs are immutable
			ttl = 0
		}
	} else {
		key = "docker/" + vrepo + "/blobs/" + ref
		ttl = 0
	}

	h.cache.Serve(w, r, key, ttl, func(ctx context.Context) (proxycache.Result, error) {
		return rp.fetch(ctx, r, vrepo, name, kind, ref)
	})
}

// fetch walks the upstream list in configured order: 404, 401, 403, 429,
// 502, 503, 504 and transport errors fall through to the next upstream (a
// registry may not carry the image, reject anonymous pulls, throttle us, or
// — when the upstream is another tiny-proxy — be down behind its ingress);
// any other status is returned as-is, and so is whatever the last upstream
// answers. When no upstream served the object but one throttled us, the
// answer is a 429 rather than a 404, so the client retries instead of
// reporting the image as missing. Deliberately mirrors distro's fetch loop,
// as distro mirrors maven's. The one difference: when no upstream answered
// at all, the last transport error is returned, so a registry that is
// simply unreachable still reads as 502 rather than as a missing image. The
// result carries the answering upstream's cache policy and retry budget.
func (rp *repo) fetch(ctx context.Context, r *http.Request, vrepo, name, kind, ref string) (proxycache.Result, error) {
	limited, retryAfter := false, "" // whether an upstream throttled us, and its Retry-After
	answered := false                // whether any upstream produced a response
	var lastErr error
	for i, up := range rp.upstreams {
		upstreamURL := objectURL(up.conf.URL, normalizeName(up.conf.URL, name), kind, ref)
		req, err := http.NewRequestWithContext(ctx, r.Method, upstreamURL, nil)
		if err != nil {
			return proxycache.Result{}, err
		}
		if accept := r.Header.Get("Accept"); accept != "" {
			req.Header.Set("Accept", accept)
		}
		resp, err := up.client.Do(req)
		if err != nil {
			slog.WarnContext(ctx, "docker upstream error, trying next", "vrepo", vrepo, "name", name, "ref", ref, "url", up.conf.URL, "err", err)
			lastErr = err
			continue
		}
		answered = true
		if resp.StatusCode == http.StatusTooManyRequests && !limited {
			slog.WarnContext(ctx, "docker upstream rate limited", "vrepo", vrepo, "name", name, "ref", ref, "url", up.conf.URL)
			limited, retryAfter = true, resp.Header.Get("Retry-After")
		}
		if shouldFallthrough(resp.StatusCode) && i+1 < len(rp.upstreams) {
			resp.Body.Close()
			slog.DebugContext(ctx, "docker upstream miss, trying next", "vrepo", vrepo, "name", name, "ref", ref, "url", up.conf.URL, "status", resp.StatusCode)
			continue
		}
		if limited && resp.StatusCode == http.StatusNotFound {
			// An upstream was throttled, so "not found" is not
			// trustworthy: report the throttling instead.
			resp.Body.Close()
			return proxycache.Result{Resp: rateLimitedResponse(r, name, retryAfter)}, nil
		}
		res := proxycache.Result{Resp: resp, NoStore: !up.conf.CacheEnabled()}
		if kind == "blobs" {
			// Blobs are the large, content-addressed half of a pull and the
			// one worth resuming; manifests are small enough to just
			// refetch. A resume goes back to this same upstream.
			res.Resume = upstream.Resumer(up.client, r.Method, upstreamURL)
			res.ResumeAttempts = up.conf.RetryAttempts()
			res.ResumeMaxWait = up.conf.RetryMaxWait()
		}
		return res, nil
	}
	if limited {
		// Every upstream failed and at least one throttled us.
		return proxycache.Result{Resp: rateLimitedResponse(r, name, retryAfter)}, nil
	}
	if !answered && lastErr != nil {
		return proxycache.Result{}, lastErr
	}
	return proxycache.Result{Resp: notFoundResponse(r, name)}, nil
}

// shouldFallthrough reports whether an upstream status means "try the next
// upstream": not found, anonymous access rejected, rate limited (the
// client already waited out retry.attempts), or a gateway error — what an
// ingress answers for a chained tiny-proxy that is down, and what
// tiny-proxy itself answers when its own upstream fails. A plain 500 is
// the upstream's own answer and is returned as-is.
func shouldFallthrough(code int) bool {
	switch code {
	case http.StatusNotFound, http.StatusUnauthorized, http.StatusForbidden,
		http.StatusTooManyRequests, http.StatusBadGateway,
		http.StatusServiceUnavailable, http.StatusGatewayTimeout:
		return true
	}
	return false
}

// notFoundResponse synthesizes a 404 so proxycache relays it uncached when no
// upstream could serve the object.
func notFoundResponse(r *http.Request, name string) *http.Response {
	return statusResponse(r, http.StatusNotFound, "not found: "+name+"\n", nil)
}

// rateLimitedResponse synthesizes the 429 returned when throttling, not
// absence, is why the object could not be served, carrying over the
// throttling upstream's Retry-After if it sent one.
func rateLimitedResponse(r *http.Request, name, retryAfter string) *http.Response {
	header := http.Header{}
	if retryAfter != "" {
		header.Set("Retry-After", retryAfter)
	}
	return statusResponse(r, http.StatusTooManyRequests,
		"rate limited by upstream: "+name+"\n", header)
}

// statusResponse builds a synthetic response for proxycache to relay; only
// 200s are cached, so these are always passed through uncached.
func statusResponse(r *http.Request, code int, body string, header http.Header) *http.Response {
	if header == nil {
		header = http.Header{}
	}
	header.Set("Content-Type", "text/plain; charset=utf-8")
	return &http.Response{
		Status:        strconv.Itoa(code) + " " + http.StatusText(code),
		StatusCode:    code,
		Proto:         "HTTP/1.1",
		ProtoMajor:    1,
		ProtoMinor:    1,
		Header:        header,
		Body:          io.NopCloser(strings.NewReader(body)),
		ContentLength: int64(len(body)),
		Request:       r,
	}
}

// objectURL builds the upstream URL of one manifest or blob. An upstream
// URL whose path starts with /v2/ already names a repository root —
// typically another tiny-proxy's vrepo, https://host/v2/<vrepo> — so the
// image path is appended to it as-is; any other URL is a registry base and
// gets the /v2/ prefix.
func objectURL(base, name, kind, ref string) string {
	root := strings.TrimSuffix(base, "/")
	if u, err := url.Parse(base); err != nil || !(u.Path == "/v2" || strings.HasPrefix(u.Path, "/v2/")) {
		root += "/v2"
	}
	return root + "/" + name + "/" + kind + "/" + ref
}

// parsePath splits "<vrepo>/<name...>/<manifests|blobs>/<ref>".
func parsePath(rest string) (vrepo, name, kind, ref string, ok bool) {
	segs := strings.Split(strings.Trim(rest, "/"), "/")
	if len(segs) < 4 {
		return "", "", "", "", false
	}
	vrepo = segs[0]
	kind = segs[len(segs)-2]
	if kind != "manifests" && kind != "blobs" {
		return "", "", "", "", false
	}
	ref = segs[len(segs)-1]
	name = strings.Join(segs[1:len(segs)-2], "/")
	if name == "" || ref == "" {
		return "", "", "", "", false
	}
	return vrepo, name, kind, ref, true
}

// normalizeName applies the Docker Hub "library/" convention for
// single-component image names when the upstream is Docker Hub itself. Only
// the host is checked: a chained tiny-proxy (https://host/v2/dockerhub)
// normalizes on its own side, and a "docker.io" in its path must not add a
// second library/.
func normalizeName(upstreamURL, name string) string {
	if strings.Contains(name, "/") {
		return name
	}
	u, err := url.Parse(upstreamURL)
	if err != nil {
		return name
	}
	host := u.Hostname()
	for _, hub := range []string{"docker.io", "docker.com"} {
		if host == hub || strings.HasSuffix(host, "."+hub) {
			return "library/" + name
		}
	}
	return name
}