package security

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

const (
	// ClaimIPAddress is the login claim holding the client's IP address.
	ClaimIPAddress = "ip_address"
	// ClaimPreviousSession is the login claim holding the session token the
	// client presented when it logged in again.
	ClaimPreviousSession = "previous_session"
	// ClaimUserAgent is the login claim holding the client's User-Agent header.
	ClaimUserAgent = "user_agent"
	// ClaimServerIPAddress is the login claim holding the IP of the server
	// address that accepted the connection.
	ClaimServerIPAddress = "server_ip_address"
	// ClaimHost is the login claim holding the requested Host.
	ClaimHost = "host"
	// ClaimOrigin is the login claim holding the Origin header.
	ClaimOrigin = "origin"
	// ClaimReferer is the login claim holding the Referer header.
	ClaimReferer = "referer"
	// ClaimClientPort is the login claim holding the port of the peer that
	// connected (the proxy's port when behind one).
	ClaimClientPort = "client_port"
	// ClaimServerPort is the login claim holding the port of the server
	// address that accepted the connection.
	ClaimServerPort = "server_port"
	// ClaimIsPrivateIP is the login claim (bool) that is true when the client
	// IP is private, loopback or link-local.
	ClaimIsPrivateIP = "is_private_ip"
	// ClaimDeviceID is the login claim holding the client-supplied device id.
	ClaimDeviceID = "device_id"
	// ClaimClientID is the login claim holding the client-supplied client id
	// (X-Client-ID header).
	ClaimClientID = "client_id"
	// ClaimClientPort is the login claim holding the port of the peer that
)

type requestInfoKey struct{}

// RequestInfoOptions tunes what WithRequestInfo and RequestInfoMiddleware read.
// The zero value trusts proxy headers from any peer, reads the default session
// cookie and takes the device id from X-Device-ID or the device_id cookie.
type RequestInfoOptions struct {
	// Cookie is used to find the previous session token.
	Cookie SessionCookieOptions
	// TrustedProxies lists the CIDRs or IPs allowed to set X-Forwarded-For /
	// X-Real-IP. When set, those headers are only honoured if the connecting
	// peer is in the list, and X-Forwarded-For is read from the right, skipping
	// trusted proxies. When empty, the headers are trusted from any peer.
	TrustedProxies []string
	// IgnoreProxyHeaders always uses the connecting peer as the client IP.
	IgnoreProxyHeaders bool
	// DeviceIDHeader and DeviceIDCookie name where the device id comes from.
	// Defaults: "X-Device-ID" and "device_id".
	DeviceIDHeader string
	DeviceIDCookie string
	// ClientIDHeader names the header carrying the client id.
	// Default: "X-Client-ID".
	ClientIDHeader string
}

func (o RequestInfoOptions) clientIDHeader() string {
	if o.ClientIDHeader != "" {
		return o.ClientIDHeader
	}
	return "X-Client-ID"
}

func (o RequestInfoOptions) deviceHeader() string {
	if o.DeviceIDHeader != "" {
		return o.DeviceIDHeader
	}
	return "X-Device-ID"
}

func (o RequestInfoOptions) deviceCookie() string {
	if o.DeviceIDCookie != "" {
		return o.DeviceIDCookie
	}
	return "device_id"
}

// RequestInfo is what the login methods learn about the calling client.
type RequestInfo struct {
	IPAddress       string // client IP
	ClientPort      string // port of the connecting peer
	IsPrivateIP     bool   // IPAddress is private, loopback or link-local
	PreviousSession string // session token the client already held
	DeviceID        string
	ClientID        string // client-supplied application/client id
	UserAgent       string
	ServerIPAddress string // local address that accepted the connection
	ServerPort      string
	Host            string
	Origin          string
	Referer         string
}

// claims returns the non-empty fields keyed by their claim names.
func (i RequestInfo) claims() map[string]any {
	out := map[string]any{}
	for k, v := range map[string]string{
		ClaimIPAddress:       i.IPAddress,
		ClaimClientPort:      i.ClientPort,
		ClaimPreviousSession: i.PreviousSession,
		ClaimDeviceID:        i.DeviceID,
		ClaimClientID:        i.ClientID,
		ClaimUserAgent:       i.UserAgent,
		ClaimServerIPAddress: i.ServerIPAddress,
		ClaimServerPort:      i.ServerPort,
		ClaimHost:            i.Host,
		ClaimOrigin:          i.Origin,
		ClaimReferer:         i.Referer,
	} {
		if v != "" {
			out[k] = v
		}
	}
	if net.ParseIP(i.IPAddress) != nil {
		out[ClaimIsPrivateIP] = i.IsPrivateIP
	}
	return out
}

// splitHostPort returns host and port of addr; the port is "" when addr has none.
func splitHostPort(addr string) (host, port string) {
	if h, p, err := net.SplitHostPort(addr); err == nil {
		return h, p
	}
	return addr, ""
}

// serverAddr returns the IP and port of the local address that accepted r's
// connection, or "" when the server did not record it.
func serverAddr(r *http.Request) (ip, port string) {
	addr, ok := r.Context().Value(http.LocalAddrContextKey).(net.Addr)
	if !ok || addr == nil {
		return "", ""
	}
	return splitHostPort(addr.String())
}

func isPrivateIP(ip string) bool {
	p := net.ParseIP(ip)
	return p != nil && (p.IsPrivate() || p.IsLoopback() || p.IsLinkLocalUnicast())
}

// parseTrusted turns CIDRs / plain IPs into networks; invalid entries are skipped.
func parseTrusted(entries []string) []*net.IPNet {
	var nets []*net.IPNet
	for _, e := range entries {
		e = strings.TrimSpace(e)
		if !strings.Contains(e, "/") {
			ip := net.ParseIP(e)
			if ip == nil {
				continue
			}
			bits := 32
			if ip.To4() == nil {
				bits = 128
			}
			e = e + "/" + strconv.Itoa(bits)
		}
		if _, n, err := net.ParseCIDR(e); err == nil {
			nets = append(nets, n)
		}
	}
	return nets
}

func inNets(nets []*net.IPNet, ip string) bool {
	p := net.ParseIP(ip)
	if p == nil {
		return false
	}
	for _, n := range nets {
		if n.Contains(p) {
			return true
		}
	}
	return false
}

// ClientIP returns the client IP, trusting X-Forwarded-For (first entry) and
// X-Real-IP from any peer, then RemoteAddr without its port. Use
// RequestInfoOptions to restrict which peers are trusted.
func ClientIP(r *http.Request) string {
	return clientIP(r, RequestInfoOptions{})
}

func clientIP(r *http.Request, o RequestInfoOptions) string {
	peer, _ := splitHostPort(r.RemoteAddr)
	if o.IgnoreProxyHeaders {
		return peer
	}
	trusted := parseTrusted(o.TrustedProxies)
	restrict := len(o.TrustedProxies) > 0
	if restrict && !inNets(trusted, peer) {
		return peer
	}
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		parts := strings.Split(xff, ",")
		if !restrict {
			return strings.TrimSpace(parts[0])
		}
		// Walk from the right: the first hop that is not a trusted proxy is the client.
		for i := len(parts) - 1; i >= 0; i-- {
			hop := strings.TrimSpace(parts[i])
			if !inNets(trusted, hop) {
				return hop
			}
		}
		return strings.TrimSpace(parts[0])
	}
	if xri := r.Header.Get("X-Real-IP"); xri != "" {
		return strings.TrimSpace(xri)
	}
	return peer
}

// previousSessionToken returns the first session token the request carries in
// the Authorization header or, failing that, the session cookie.
func previousSessionToken(r *http.Request, cookie SessionCookieOptions) string {
	if h := r.Header.Get("Authorization"); h != "" {
		first, _, _ := strings.Cut(h, ",")
		first = strings.TrimSpace(first)
		first = strings.TrimPrefix(first, "Bearer ")
		first = strings.TrimPrefix(first, "Token ")
		return strings.TrimSpace(first)
	}
	return GetSessionCookie(r, cookie)
}

func deviceID(r *http.Request, o RequestInfoOptions) string {
	if v := strings.TrimSpace(r.Header.Get(o.deviceHeader())); v != "" {
		return v
	}
	if c, err := r.Cookie(o.deviceCookie()); err == nil {
		return c.Value
	}
	return ""
}

// WithRequestInfo returns a context carrying details of r: client IP and port,
// private-IP flag, previous session token, device id, client id, user agent, server IP and
// port, host, origin and referer. Login, LoginWithAPIKey and LoginWithPasskey
// add them to the login claims (see the Claim* constants); empty values are
// left out.
func WithRequestInfo(ctx context.Context, r *http.Request, opts ...RequestInfoOptions) context.Context {
	var o RequestInfoOptions
	if len(opts) > 0 {
		o = opts[0]
	}
	ip := clientIP(r, o)
	_, clientPort := splitHostPort(r.RemoteAddr)
	serverIP, serverPort := serverAddr(r)
	return context.WithValue(ctx, requestInfoKey{}, RequestInfo{
		IPAddress:       ip,
		ClientPort:      clientPort,
		IsPrivateIP:     isPrivateIP(ip),
		PreviousSession: previousSessionToken(r, o.Cookie),
		DeviceID:        deviceID(r, o),
		ClientID:        strings.TrimSpace(r.Header.Get(o.clientIDHeader())),
		UserAgent:       r.UserAgent(),
		ServerIPAddress: serverIP,
		ServerPort:      serverPort,
		Host:            r.Host,
		Origin:          r.Header.Get("Origin"),
		Referer:         r.Referer(),
	})
}

// RequestInfoFromContext returns the info stored by WithRequestInfo.
func RequestInfoFromContext(ctx context.Context) (RequestInfo, bool) {
	info, ok := ctx.Value(requestInfoKey{}).(RequestInfo)
	return info, ok
}

// RequestInfoMiddleware stores the request details (see WithRequestInfo) in the
// request context for every route, including ones that skip authentication
// (such as login), so the login methods can send them as claims.
func RequestInfoMiddleware(opts ...RequestInfoOptions) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			next.ServeHTTP(w, r.WithContext(WithRequestInfo(r.Context(), r, opts...)))
		})
	}
}

// enrichClaims returns a copy of claims with the request info from ctx added.
// Values the caller already set are kept.
func enrichClaims(ctx context.Context, claims map[string]any) map[string]any {
	info, ok := RequestInfoFromContext(ctx)
	if !ok {
		return claims
	}
	extra := info.claims()
	out := make(map[string]any, len(claims)+len(extra))
	for k, v := range claims {
		out[k] = v
	}
	for k, v := range extra {
		if _, set := out[k]; !set {
			out[k] = v
		}
	}
	return out
}