// Kubernetes TokenReview-backed workload-identity validator.
//
// The signing-service receives the workload's projected ServiceAccount
// token in the `X-Workload-Token` header. We submit that token to the
// cluster via `authentication.k8s.io/v1 TokenReview` using the pod's own
// in-cluster credentials, and return the authenticated
// `system:serviceaccount:<ns>:<sa>` as (namespace, serviceAccount).
//
// This file deliberately does not import k8s.io/client-go to keep the
// dependency surface narrow; a hand-rolled HTTPS POST against the API
// server is all that's required.
package token

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"
)

const (
	// WorkloadTokenHeader is the HTTP header carrying the workload's
	// projected SA token (not the Authorization header).
	WorkloadTokenHeader = "X-Workload-Token"

	defaultSATokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token"
	defaultCAPath      = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
)

// ErrNoWorkloadToken is returned when the X-Workload-Token header is
// missing. Caller MUST map this to a 401.
var ErrNoWorkloadToken = errors.New("tokenreview: missing X-Workload-Token header")

// ErrWorkloadAuthFailed is returned when TokenReview reports the token is
// invalid or unauthenticated. Caller MUST map this to a 401.
var ErrWorkloadAuthFailed = errors.New("tokenreview: workload token not authenticated")

// K8sIdentity is the subset of a TokenReview UserInfo we care about.
type K8sIdentity struct {
	Namespace      string
	ServiceAccount string
	// UID is the ServiceAccount's UID (useful for audit; may be empty for
	// non-SA callers which should be rejected before this type is built).
	UID string
}

// TokenReviewer validates a workload-presented token and returns the
// authenticated ServiceAccount identity.
type TokenReviewer interface {
	Review(ctx context.Context, workloadToken string) (K8sIdentity, error)
}

// FromWorkloadHeader extracts the workload token from the X-Workload-Token
// header; it does NOT accept Bearer prefixes because this header is
// documented as raw-token.
func FromWorkloadHeader(r *http.Request) (string, error) {
	t := strings.TrimSpace(r.Header.Get(WorkloadTokenHeader))
	if t == "" {
		return "", ErrNoWorkloadToken
	}
	return t, nil
}

// StaticTokenReviewer is a test seam: a fixed table from workload token
// → identity. All production paths use K8sTokenReviewer.
type StaticTokenReviewer map[string]K8sIdentity

func (s StaticTokenReviewer) Review(_ context.Context, workloadToken string) (K8sIdentity, error) {
	id, ok := s[workloadToken]
	if !ok {
		return K8sIdentity{}, ErrWorkloadAuthFailed
	}
	return id, nil
}

// K8sTokenReviewer talks to a Kubernetes API server.
type K8sTokenReviewer struct {
	APIServer   string // e.g. "https://kubernetes.default.svc"
	BearerToken string // pod's own SA token (not the workload token)
	HTTPClient  *http.Client
}

// NewInClusterTokenReviewer loads the pod's SA token and the cluster CA
// from the standard in-cluster paths and wires up a TLS-pinned client.
// The returned reviewer talks to `https://kubernetes.default.svc`.
func NewInClusterTokenReviewer() (*K8sTokenReviewer, error) {
	return NewTokenReviewerFromPaths(
		"https://kubernetes.default.svc",
		defaultSATokenPath,
		defaultCAPath,
	)
}

// NewTokenReviewerFromPaths loads SA token & CA from explicit paths; the
// in-cluster helper above is a thin wrapper.
func NewTokenReviewerFromPaths(apiServer, tokenPath, caPath string) (*K8sTokenReviewer, error) {
	tokenBytes, err := os.ReadFile(tokenPath)
	if err != nil {
		return nil, fmt.Errorf("tokenreview: read SA token: %w", err)
	}
	caBytes, err := os.ReadFile(caPath)
	if err != nil {
		return nil, fmt.Errorf("tokenreview: read CA: %w", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(caBytes) {
		return nil, errors.New("tokenreview: CA pem invalid")
	}
	hc := &http.Client{
		Timeout: 5 * time.Second,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12},
		},
	}
	return &K8sTokenReviewer{
		APIServer:   apiServer,
		BearerToken: strings.TrimSpace(string(tokenBytes)),
		HTTPClient:  hc,
	}, nil
}

// tokenReviewRequest / tokenReviewResponse mirror just the fields we
// read, so we don't pull in client-go's typed meta.
type tokenReviewRequest struct {
	APIVersion string                 `json:"apiVersion"`
	Kind       string                 `json:"kind"`
	Spec       tokenReviewRequestSpec `json:"spec"`
}

type tokenReviewRequestSpec struct {
	Token string `json:"token"`
}

type tokenReviewResponse struct {
	Status tokenReviewStatus `json:"status"`
}

type tokenReviewStatus struct {
	Authenticated bool                `json:"authenticated"`
	User          tokenReviewUserInfo `json:"user"`
	Error         string              `json:"error"`
}

type tokenReviewUserInfo struct {
	Username string `json:"username"`
	UID      string `json:"uid"`
}

func (r *K8sTokenReviewer) Review(ctx context.Context, workloadToken string) (K8sIdentity, error) {
	if workloadToken == "" {
		return K8sIdentity{}, ErrNoWorkloadToken
	}
	payload := tokenReviewRequest{
		APIVersion: "authentication.k8s.io/v1",
		Kind:       "TokenReview",
		Spec:       tokenReviewRequestSpec{Token: workloadToken},
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return K8sIdentity{}, fmt.Errorf("tokenreview: marshal: %w", err)
	}
	url := strings.TrimRight(r.APIServer, "/") + "/apis/authentication.k8s.io/v1/tokenreviews"
	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
	if err != nil {
		return K8sIdentity{}, fmt.Errorf("tokenreview: build request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+r.BearerToken)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	resp, err := r.HTTPClient.Do(req)
	if err != nil {
		return K8sIdentity{}, fmt.Errorf("tokenreview: POST: %w", err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode/100 != 2 {
		return K8sIdentity{}, fmt.Errorf("tokenreview: api server %d: %s", resp.StatusCode, truncate(string(raw), 256))
	}
	var out tokenReviewResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		return K8sIdentity{}, fmt.Errorf("tokenreview: decode: %w", err)
	}
	if !out.Status.Authenticated {
		if out.Status.Error != "" {
			return K8sIdentity{}, fmt.Errorf("%w: %s", ErrWorkloadAuthFailed, out.Status.Error)
		}
		return K8sIdentity{}, ErrWorkloadAuthFailed
	}
	ns, sa, ok := parseSAUsername(out.Status.User.Username)
	if !ok {
		return K8sIdentity{}, fmt.Errorf("%w: not a ServiceAccount: %q", ErrWorkloadAuthFailed, out.Status.User.Username)
	}
	return K8sIdentity{Namespace: ns, ServiceAccount: sa, UID: out.Status.User.UID}, nil
}

// parseSAUsername extracts (ns, sa) from `system:serviceaccount:<ns>:<sa>`.
func parseSAUsername(u string) (string, string, bool) {
	const prefix = "system:serviceaccount:"
	if !strings.HasPrefix(u, prefix) {
		return "", "", false
	}
	rest := u[len(prefix):]
	i := strings.IndexByte(rest, ':')
	if i <= 0 || i == len(rest)-1 {
		return "", "", false
	}
	return rest[:i], rest[i+1:], true
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}