package callerauth

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

type KubernetesTokenReviewer struct {
	endpoint          string
	audience          string
	reviewerTokenFile string
	client            *http.Client
}

func NewKubernetesTokenReviewer(endpoint, audience, reviewerTokenFile, caFile string) (*KubernetesTokenReviewer, error) {
	ca, err := os.ReadFile(caFile)
	if err != nil {
		return nil, fmt.Errorf("read Kubernetes CA: %w", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(ca) {
		return nil, fmt.Errorf("Kubernetes CA file contains no certificates")
	}
	return &KubernetesTokenReviewer{
		endpoint:          strings.TrimRight(endpoint, "/") + "/apis/authentication.k8s.io/v1/tokenreviews",
		audience:          audience,
		reviewerTokenFile: reviewerTokenFile,
		client: &http.Client{
			Timeout: 3 * time.Second,
			Transport: &http.Transport{TLSClientConfig: &tls.Config{
				MinVersion: tls.VersionTLS12,
				RootCAs:    pool,
			}},
		},
	}, nil
}

type tokenReview struct {
	APIVersion string            `json:"apiVersion"`
	Kind       string            `json:"kind"`
	Spec       tokenReviewSpec   `json:"spec"`
	Status     tokenReviewStatus `json:"status,omitempty"`
}

type tokenReviewSpec struct {
	Token     string   `json:"token"`
	Audiences []string `json:"audiences"`
}

type tokenReviewStatus struct {
	Authenticated bool     `json:"authenticated"`
	Audiences     []string `json:"audiences"`
	Error         string   `json:"error"`
	User          struct {
		Username string `json:"username"`
	} `json:"user"`
}

func (r *KubernetesTokenReviewer) Review(ctx context.Context, callerToken string) (Identity, error) {
	reviewerToken, err := os.ReadFile(r.reviewerTokenFile)
	if err != nil {
		return Identity{}, fmt.Errorf("read reviewer credential: %w", err)
	}
	payload, err := json.Marshal(tokenReview{
		APIVersion: "authentication.k8s.io/v1",
		Kind:       "TokenReview",
		Spec: tokenReviewSpec{
			Token:     callerToken,
			Audiences: []string{r.audience},
		},
	})
	if err != nil {
		return Identity{}, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.endpoint, bytes.NewReader(payload))
	if err != nil {
		return Identity{}, err
	}
	req.Header.Set("Authorization", "Bearer "+strings.TrimSpace(string(reviewerToken)))
	req.Header.Set("Content-Type", "application/json")
	resp, err := r.client.Do(req)
	if err != nil {
		return Identity{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		return Identity{}, fmt.Errorf("TokenReview returned %s: %s", resp.Status, strings.TrimSpace(string(body)))
	}
	var review tokenReview
	if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&review); err != nil {
		return Identity{}, fmt.Errorf("decode TokenReview: %w", err)
	}
	if review.Status.Error != "" {
		return Identity{}, fmt.Errorf("%w: TokenReview: %s", ErrUnauthenticated, review.Status.Error)
	}
	if !review.Status.Authenticated {
		return Identity{}, nil
	}
	identity := Identity{Username: review.Status.User.Username, Audiences: review.Status.Audiences}
	if exp, ok := tokenExpiry(callerToken); ok {
		identity.NotAfter = exp
	}
	return identity, nil
}

// tokenExpiry reads exp from a JWT payload. TokenReview already validated the
// token, including expiry; the claim is captured here because the review
// response does not return it.
func tokenExpiry(token string) (time.Time, bool) {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return time.Time{}, false
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		payload, err = base64.URLEncoding.DecodeString(parts[1])
		if err != nil {
			return time.Time{}, false
		}
	}
	var claims struct {
		Exp int64 `json:"exp"`
	}
	if err := json.Unmarshal(payload, &claims); err != nil || claims.Exp <= 0 {
		return time.Time{}, false
	}
	return time.Unix(claims.Exp, 0).UTC(), true
}