// llm-token-vendor mints ZITADEL access tokens for a caller that cannot sign.
//
// Some callers are vendor images: paperless-ai can only read a bearer token from a
// file. It cannot sign an assertion, so it cannot use a machine key, and a client
// secret would rotate with a window.
//
// This sidecar closes the gap. It signs with a machine key, exchanges the assertion
// for an access token, and writes the token to the file the vendor image already
// reads. The file is replaced atomically and refreshed at 80% of the token lifetime,
// so the vendor never reads a partial or expired value.
//
// The key is re-read on every mint. A rotation therefore lands without a restart:
// publish the new key, the next mint uses it, retire the old one.
package main

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"time"
)

type config struct {
	tokenURL string
	issuer   string
	userID   string
	keyID    string
	keyPath  string
	scope    string
	outPath  string
}

func main() {
	var c config
	flag.StringVar(&c.tokenURL, "token-url", env("TOKEN_URL", ""), "ZITADEL token endpoint")
	flag.StringVar(&c.issuer, "issuer", env("ISSUER", ""), "ZITADEL issuer, the assertion audience")
	flag.StringVar(&c.userID, "user-id", env("USER_ID", ""), "machine user id")
	flag.StringVar(&c.keyID, "key-id", env("KEY_ID", ""), "machine key id")
	flag.StringVar(&c.keyPath, "key-file", env("KEY_FILE", ""), "PEM private key, re-read on every mint")
	flag.StringVar(&c.scope, "scope", env("SCOPE", "openid"), "requested scopes")
	flag.StringVar(&c.outPath, "out", env("OUT", ""), "file the vendor image reads")
	flag.Parse()

	if c.tokenURL == "" || c.issuer == "" || c.userID == "" || c.keyPath == "" || c.outPath == "" {
		log.Fatal("token-url, issuer, user-id, key-file and out are required")
	}

	for {
		lifetime, err := mintAndWrite(c)
		if err != nil {
			// Exit rather than spin: a pod that cannot mint must be visible, and the
			// vendor image must not keep serving the caller a stale token.
			log.Fatalf("cannot mint a token: %v", err)
		}
		wait := time.Duration(float64(lifetime) * 0.8)
		log.Printf("token written to %s, next mint in %s", c.outPath, wait.Round(time.Second))
		time.Sleep(wait)
	}
}

func mintAndWrite(c config) (time.Duration, error) {
	pemBytes, err := os.ReadFile(c.keyPath)
	if err != nil {
		return 0, fmt.Errorf("cannot read the key %s: %w", c.keyPath, err)
	}
	assertion, err := assertion(c, time.Now(), pemBytes)
	if err != nil {
		return 0, err
	}

	form := url.Values{}
	form.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
	form.Set("assertion", assertion)
	form.Set("scope", c.scope)

	resp, err := http.PostForm(c.tokenURL, form)
	if err != nil {
		return 0, fmt.Errorf("token request failed: %w", err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
	if resp.StatusCode != http.StatusOK {
		return 0, fmt.Errorf("token endpoint returned %d: %s", resp.StatusCode, string(body))
	}
	var out struct {
		AccessToken string `json:"access_token"`
		ExpiresIn   int    `json:"expires_in"`
	}
	if err := json.Unmarshal(body, &out); err != nil {
		return 0, fmt.Errorf("cannot decode the token response: %w", err)
	}
	if out.AccessToken == "" {
		return 0, fmt.Errorf("the token response holds no access_token")
	}

	// Write beside the target, then rename. A rename is atomic, so the vendor either
	// reads the previous token or the new one, never a half-written file.
	tmp, err := os.CreateTemp(filepath.Dir(c.outPath), ".token-*")
	if err != nil {
		return 0, fmt.Errorf("cannot create a temporary file: %w", err)
	}
	defer os.Remove(tmp.Name())
	if _, err := tmp.WriteString(out.AccessToken); err != nil {
		tmp.Close()
		return 0, fmt.Errorf("cannot write the token: %w", err)
	}
	tmp.Close()
	if err := os.Chmod(tmp.Name(), 0o644); err != nil {
		return 0, fmt.Errorf("cannot set the file mode: %w", err)
	}
	if err := os.Rename(tmp.Name(), c.outPath); err != nil {
		return 0, fmt.Errorf("cannot replace %s: %w", c.outPath, err)
	}

	lifetime := time.Duration(out.ExpiresIn) * time.Second
	if lifetime <= 0 {
		lifetime = 5 * time.Minute
	}
	return lifetime, nil
}

func assertion(c config, now time.Time, pemBytes []byte) (string, error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return "", fmt.Errorf("the key file is not PEM")
	}
	var key *rsa.PrivateKey
	if k, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
		key = k
	} else {
		parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return "", fmt.Errorf("the key is not an RSA private key: %w", err)
		}
		rk, ok := parsed.(*rsa.PrivateKey)
		if !ok {
			return "", fmt.Errorf("the key holds %T, not RSA", parsed)
		}
		key = rk
	}

	enc := func(v any) (string, error) {
		raw, err := json.Marshal(v)
		if err != nil {
			return "", err
		}
		return base64.RawURLEncoding.EncodeToString(raw), nil
	}
	header := map[string]any{"alg": "RS256", "typ": "JWT"}
	if c.keyID != "" {
		header["kid"] = c.keyID
	}
	claims := map[string]any{
		"iss": c.userID,
		"sub": c.userID,
		"aud": c.issuer,
		// Backdated, so a clock that is a few seconds fast cannot make the assertion
		// "not yet valid".
		"iat": now.Add(-5 * time.Second).Unix(),
		"exp": now.Add(5 * time.Minute).Unix(),
	}
	h, err := enc(header)
	if err != nil {
		return "", err
	}
	cl, err := enc(claims)
	if err != nil {
		return "", err
	}
	signing := h + "." + cl
	digest := sha256.Sum256([]byte(signing))
	sig, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
	if err != nil {
		return "", fmt.Errorf("cannot sign the assertion: %w", err)
	}
	return signing + "." + base64.RawURLEncoding.EncodeToString(sig), nil
}

func env(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}