package middleware

import (
	"context"
	"net/http"
	"net/netip"
	"strings"
)

// ========== RealIP - ClientInfo ==========

// ClientInfo describes the client as seen after proxy processing.
type ClientInfo struct {
	IP       netip.Addr // best-effort real client IP
	Peer     netip.Addr // immediate TCP peer (after PROXY protocol, if used)
	Scheme   string     // "http" or "https", honouring trusted X-Forwarded-Proto
	ViaProxy bool       // true when IP came from a trusted forwarding header
}

// Client returns the ClientInfo stored by RealIP. Without RealIP it falls back
// to the connection itself (see peerInfo) and ignores forwarding headers.
func Client(r *http.Request) ClientInfo {
	if info, ok := r.Context().Value(clientInfoKey).(ClientInfo); ok {
		return info
	}
	return peerInfo(r)
}

// ClientIP returns the client IP as a string ("" if unknown).
func ClientIP(r *http.Request) string {
	ip := Client(r).IP
	if !ip.IsValid() {
		return ""
	}
	return ip.String()
}

// IsSecure reports whether the request reached the edge over HTTPS. Without
// RealIP it only reports whether the connection to this server is TLS.
func IsSecure(r *http.Request) bool {
	return Client(r).Scheme == "https"
}

// peerInfo builds a ClientInfo from the connection alone: the TCP peer (after
// PROXY protocol, if used) and whether this server terminated TLS.
func peerInfo(r *http.Request) ClientInfo {
	peer := addrFromRemote(r.RemoteAddr)

	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}

	return ClientInfo{
		IP:     peer,
		Peer:   peer,
		Scheme: scheme,
	}
}

// ========== RealIP ==========

// RealIPOptions configures RealIP.
type RealIPOptions struct {
	// Trusted are the proxies whose forwarding headers are believed.
	Trusted []netip.Prefix
	// Headers are the headers the client IP is read from, in priority order:
	// the first of them that the request has and that gives an address is
	// used. E.g. X-Forwarded-For, Forwarded, X-Real-IP.
	Headers []string
	// StripUntrusted removes the forwarding headers from requests that are
	// not from a trusted proxy.
	StripUntrusted bool
}

type ctxKey int

const clientInfoKey ctxKey = iota

var forwardingHeaders = []string{
	"X-Forwarded-For",
	"X-Forwarded-Proto",
	"X-Forwarded-Host",
	"X-Forwarded-Port",
	"X-Forwarded-Ssl",
	"Forwarded",
	"X-Real-IP",
}

// RealIP determines the client IP and scheme. Forwarding headers are only
// honoured when the TCP peer is a trusted proxy; the X-Forwarded-For and
// Forwarded chains are walked right-to-left skipping trusted hops, so a
// client cannot spoof its address by prepending entries.
//
// PROXY protocol is handled at the listener (see server package), so by the
// time this runs r.RemoteAddr already holds the address from the PROXY header.
func RealIP(opts RealIPOptions) func(http.Handler) http.Handler {
	trusted := func(addr netip.Addr) bool {
		for _, prefix := range opts.Trusted {
			if prefix.Contains(addr) {
				return true
			}
		}
		return false
	}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			info := peerInfo(r)

			if info.Peer.IsValid() && trusted(info.Peer) {
				for _, name := range opts.Headers {
					if ip, ok := clientFromHeader(r.Header, name, trusted); ok {
						info.IP = ip
						info.ViaProxy = true
						break
					}
				}
				proto := forwardedProto(r.Header)
				if proto == "http" || proto == "https" {
					info.Scheme = proto
				}
			} else if opts.StripUntrusted {
				for _, name := range forwardingHeaders {
					r.Header.Del(name)
				}
			}

			ctx := context.WithValue(r.Context(), clientInfoKey, info)
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

func clientFromHeader(
	header http.Header, name string, trusted func(netip.Addr) bool,
) (netip.Addr, bool) {
	switch http.CanonicalHeaderKey(name) {
	case "X-Forwarded-For":
		return walkChain(splitList(header.Values("X-Forwarded-For")), trusted)

	case "Forwarded":
		var chain []string
		for _, element := range splitList(header.Values("Forwarded")) {
			if value, ok := forwardedParam(element, "for"); ok {
				chain = append(chain, value)
			} else {
				chain = append(chain, "") // element without for= breaks the chain
			}
		}
		return walkChain(chain, trusted)

	case "X-Real-Ip":
		if value := strings.TrimSpace(header.Get("X-Real-IP")); value != "" {
			return parseIPToken(value)
		}
	}

	return netip.Addr{}, false
}

// walkChain returns the right-most address that is not a trusted proxy. If
// every hop is trusted, the left-most is returned. Any unparsable hop
// ("unknown", obfuscated identifiers, garbage) stops the walk.
func walkChain(
	chain []string, trusted func(netip.Addr) bool,
) (netip.Addr, bool) {
	var last netip.Addr
	for i := len(chain) - 1; i >= 0; i-- {
		ip, ok := parseIPToken(chain[i])
		if !ok {
			return netip.Addr{}, false
		}
		if !trusted(ip) {
			return ip, true
		}
		last = ip
	}
	return last, last.IsValid()
}

func forwardedProto(header http.Header) string {
	if value := header.Get("X-Forwarded-Proto"); value != "" {
		first, _, _ := strings.Cut(value, ",")
		return strings.ToLower(strings.TrimSpace(first))
	}

	if elements := splitList(header.Values("Forwarded")); len(elements) > 0 {
		if value, ok := forwardedParam(elements[0], "proto"); ok {
			return strings.ToLower(value)
		}
	}

	return ""
}

// forwardedParam extracts a parameter from one Forwarded HTTP Extension
// element (RFC 7239).
func forwardedParam(element, key string) (string, bool) {
	for _, pair := range strings.Split(element, ";") {
		name, value, ok := strings.Cut(strings.TrimSpace(pair), "=")
		if ok && strings.EqualFold(name, key) {
			return strings.Trim(strings.TrimSpace(value), `"`), true
		}
	}
	return "", false
}

func splitList(values []string) []string {
	var out []string
	for _, value := range values {
		for _, part := range strings.Split(value, ",") {
			if part = strings.TrimSpace(part); part != "" {
				out = append(out, part)
			}
		}
	}
	return out
}

// parseIPToken accepts "1.2.3.4", "1.2.3.4:80", "::1", "[::1]" and "[::1]:80".
func parseIPToken(token string) (netip.Addr, bool) {
	token = strings.Trim(strings.TrimSpace(token), `"`)
	if addrPort, err := netip.ParseAddrPort(token); err == nil {
		return addrPort.Addr().Unmap(), true
	}

	token = strings.TrimSuffix(strings.TrimPrefix(token, "["), "]")
	if addr, err := netip.ParseAddr(token); err == nil {
		return addr.Unmap().WithZone(""), true
	}

	return netip.Addr{}, false
}

func addrFromRemote(remote string) netip.Addr {
	if addr, ok := parseIPToken(remote); ok {
		return addr
	}
	return netip.Addr{}
}