package apptoken

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	jose "github.com/go-jose/go-jose/v4"
	"github.com/go-jose/go-jose/v4/jwt"
	"github.com/hashicorp/vault/sdk/logical"
)

const testIssuer = "https://vault.example/v1/apptoken"

// getTestBackend returns a backend backed by in-memory storage with sane
// default/max lease TTLs so issued tokens have a non-zero lifetime.
func getTestBackend(t *testing.T) (*appTokenBackend, logical.Storage) {
	t.Helper()

	sysView := logical.TestSystemView()
	sysView.DefaultLeaseTTLVal = time.Hour
	sysView.MaxLeaseTTLVal = 24 * time.Hour

	config := logical.TestBackendConfig()
	config.StorageView = &logical.InmemStorage{}
	config.System = sysView

	b, err := Factory(context.Background(), config)
	if err != nil {
		t.Fatalf("creating backend: %v", err)
	}
	return b.(*appTokenBackend), config.StorageView
}

// do runs a request against the backend and fails on transport or logical error.
func do(t *testing.T, b *appTokenBackend, s logical.Storage, op logical.Operation, path string, data map[string]interface{}) *logical.Response {
	t.Helper()
	resp, err := b.HandleRequest(context.Background(), &logical.Request{
		Operation:   op,
		Path:        path,
		Storage:     s,
		Data:        data,
		DisplayName: "test-caller",
	})
	if err != nil {
		t.Fatalf("%s %s: transport error: %v", op, path, err)
	}
	if resp != nil && resp.IsError() {
		t.Fatalf("%s %s: logical error: %v", op, path, resp.Error())
	}
	return resp
}

// doReq runs a request and returns both resp and err without failing, for
// negative-path assertions.
func doReq(b *appTokenBackend, s logical.Storage, req *logical.Request) (*logical.Response, error) {
	req.Storage = s
	return b.HandleRequest(context.Background(), req)
}

// configureIssuer writes the engine config with the test issuer.
func configureIssuer(t *testing.T, b *appTokenBackend, s logical.Storage, algorithm string) {
	t.Helper()
	data := map[string]interface{}{"issuer": testIssuer}
	if algorithm != "" {
		data["algorithm"] = algorithm
	}
	do(t, b, s, logical.CreateOperation, "config", data)
}

// readJWKS fetches the unauthenticated JWKS endpoint and parses it.
func readJWKS(t *testing.T, b *appTokenBackend, s logical.Storage) jose.JSONWebKeySet {
	t.Helper()
	resp := do(t, b, s, logical.ReadOperation, jwksPath, nil)
	raw, ok := resp.Data[logical.HTTPRawBody].([]byte)
	if !ok {
		t.Fatalf("jwks response missing raw body: %#v", resp.Data)
	}
	var jwks jose.JSONWebKeySet
	if err := json.Unmarshal(raw, &jwks); err != nil {
		t.Fatalf("parsing jwks: %v", err)
	}
	return jwks
}

// validateToken verifies a compact JWT against the given JWKS the way a
// consuming service would: find the key by kid, check the signature, then the
// registered claims. It returns the validated claims.
func validateToken(t *testing.T, jwks jose.JSONWebKeySet, token, expectedIssuer, expectedAudience string) jwt.Claims {
	t.Helper()

	parsed, err := jwt.ParseSigned(token, []jose.SignatureAlgorithm{jose.EdDSA, jose.RS256})
	if err != nil {
		t.Fatalf("parsing token: %v", err)
	}
	if len(parsed.Headers) == 0 || parsed.Headers[0].KeyID == "" {
		t.Fatalf("token has no kid header")
	}
	keys := jwks.Key(parsed.Headers[0].KeyID)
	if len(keys) == 0 {
		t.Fatalf("kid %q not present in JWKS", parsed.Headers[0].KeyID)
	}

	var claims jwt.Claims
	if err := parsed.Claims(keys[0].Key, &claims); err != nil {
		t.Fatalf("verifying token signature: %v", err)
	}
	if err := claims.Validate(jwt.Expected{
		Issuer:      expectedIssuer,
		AnyAudience: jwt.Audience{expectedAudience},
		Time:        time.Now(),
	}); err != nil {
		t.Fatalf("validating claims: %v", err)
	}
	return claims
}

// mintToken creates a role (if data provided) and reads creds, returning the token string.
func mintToken(t *testing.T, b *appTokenBackend, s logical.Storage, role string) *logical.Response {
	t.Helper()
	return do(t, b, s, logical.ReadOperation, "creds/"+role, nil)
}

// findAndVerify looks a token's kid up in the JWKS and checks its signature
// only (no claim validation). Returns an error when the kid is absent or the
// signature does not verify.
func findAndVerify(jwks jose.JSONWebKeySet, token string) error {
	parsed, err := jwt.ParseSigned(token, []jose.SignatureAlgorithm{jose.EdDSA, jose.RS256})
	if err != nil {
		return err
	}
	if len(parsed.Headers) == 0 {
		return errNoKID
	}
	keys := jwks.Key(parsed.Headers[0].KeyID)
	if len(keys) == 0 {
		return errNoKID
	}
	var claims jwt.Claims
	return parsed.Claims(keys[0].Key, &claims)
}

// verifyAudience verifies signature and requires the given audience.
func verifyAudience(jwks jose.JSONWebKeySet, token, audience string) error {
	parsed, err := jwt.ParseSigned(token, []jose.SignatureAlgorithm{jose.EdDSA, jose.RS256})
	if err != nil {
		return err
	}
	keys := jwks.Key(parsed.Headers[0].KeyID)
	if len(keys) == 0 {
		return errNoKID
	}
	var claims jwt.Claims
	if err := parsed.Claims(keys[0].Key, &claims); err != nil {
		return err
	}
	return claims.Validate(jwt.Expected{AnyAudience: jwt.Audience{audience}, Time: time.Now()})
}

func verifyWrongAudienceFails(jwks jose.JSONWebKeySet, token string) error {
	return verifyAudience(jwks, token, "some-other-app")
}

func decodeJSON(t *testing.T, raw []byte) map[string]interface{} {
	t.Helper()
	var m map[string]interface{}
	if err := json.Unmarshal(raw, &m); err != nil {
		t.Fatalf("decoding json: %v", err)
	}
	return m
}

var errNoKID = &kidError{}

type kidError struct{}

func (*kidError) Error() string { return "kid not found in JWKS" }