package slurm

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"time"
)

type JWTClaims struct {
	Sun string `json:"sun"`
	IAT int64  `json:"iat"`
	EXP int64  `json:"exp"`
}

func SignJWT(key []byte, username string, lifespan time.Duration) (string, error) {
	if username == "" {
		return "", fmt.Errorf("username must not be empty")
	}

	now := time.Now()
	header := map[string]string{"alg": "HS256", "typ": "JWT"}
	claims := JWTClaims{
		Sun: username,
		IAT: now.Unix(),
		EXP: now.Add(lifespan).Unix(),
	}

	headerJSON, err := json.Marshal(header)
	if err != nil {
		return "", fmt.Errorf("marshal header: %w", err)
	}
	claimsJSON, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("marshal claims: %w", err)
	}

	enc := base64.RawURLEncoding
	headerEnc := enc.EncodeToString(headerJSON)
	claimsEnc := enc.EncodeToString(claimsJSON)

	signingInput := headerEnc + "." + claimsEnc
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte(signingInput))
	sig := enc.EncodeToString(mac.Sum(nil))

	return signingInput + "." + sig, nil
}

func ReadJWTKey(path string) ([]byte, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	if len(data) < 16 {
		return nil, fmt.Errorf("key must be at least 16 bytes, got %d", len(data))
	}
	return data, nil
}