package auth

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

// JobTokenDuration is the validity period for job tokens
const JobTokenDuration = 1 * time.Hour

// JobTokenClaims represents the claims in a job token
type JobTokenClaims struct {
	JobID    int64 `json:"job_id"`
	RunnerID int64 `json:"runner_id"`
	TaskID   int64 `json:"task_id"`
	Exp      int64 `json:"exp"` // Unix timestamp
}

// jobTokenSecret is the secret used to sign job tokens
// Generated once at startup and kept in memory
var jobTokenSecret []byte

func init() {
	// Generate a random secret for signing job tokens
	// This means tokens are invalidated on server restart, which is acceptable
	// for short-lived job tokens
	jobTokenSecret = make([]byte, 32)
	if _, err := rand.Read(jobTokenSecret); err != nil {
		panic(fmt.Sprintf("failed to generate job token secret: %v", err))
	}
}

// GenerateJobToken creates a new job token for a specific job/runner/task combination
func GenerateJobToken(jobID, runnerID, taskID int64) (string, error) {
	claims := JobTokenClaims{
		JobID:    jobID,
		RunnerID: runnerID,
		TaskID:   taskID,
		Exp:      time.Now().Add(JobTokenDuration).Unix(),
	}

	// Encode claims to JSON
	claimsJSON, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("failed to marshal claims: %w", err)
	}

	// Create HMAC signature
	h := hmac.New(sha256.New, jobTokenSecret)
	h.Write(claimsJSON)
	signature := h.Sum(nil)

	// Combine claims and signature: base64(claims).base64(signature)
	token := base64.RawURLEncoding.EncodeToString(claimsJSON) + "." +
		base64.RawURLEncoding.EncodeToString(signature)

	return token, nil
}

// ValidateJobToken validates a job token and returns the claims if valid
func ValidateJobToken(token string) (*JobTokenClaims, error) {
	// Split token into claims and signature
	var claimsB64, sigB64 string
	dotIdx := -1
	for i := len(token) - 1; i >= 0; i-- {
		if token[i] == '.' {
			dotIdx = i
			break
		}
	}
	if dotIdx == -1 {
		return nil, fmt.Errorf("invalid token format")
	}
	claimsB64 = token[:dotIdx]
	sigB64 = token[dotIdx+1:]

	// Decode claims
	claimsJSON, err := base64.RawURLEncoding.DecodeString(claimsB64)
	if err != nil {
		return nil, fmt.Errorf("invalid token encoding: %w", err)
	}

	// Decode signature
	signature, err := base64.RawURLEncoding.DecodeString(sigB64)
	if err != nil {
		return nil, fmt.Errorf("invalid signature encoding: %w", err)
	}

	// Verify signature
	h := hmac.New(sha256.New, jobTokenSecret)
	h.Write(claimsJSON)
	expectedSig := h.Sum(nil)
	if !hmac.Equal(signature, expectedSig) {
		return nil, fmt.Errorf("invalid signature")
	}

	// Parse claims
	var claims JobTokenClaims
	if err := json.Unmarshal(claimsJSON, &claims); err != nil {
		return nil, fmt.Errorf("invalid claims: %w", err)
	}

	// Check expiration
	if time.Now().Unix() > claims.Exp {
		return nil, fmt.Errorf("token expired")
	}

	return &claims, nil
}