package main

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/go-jose/go-jose/v4"
	"github.com/go-jose/go-jose/v4/jwt"
	"github.com/grafana/grafana-plugin-sdk-go/config"
)

type authFixture struct {
	key       *ecdsa.PrivateKey
	server    *httptest.Server
	admin     atomic.Bool
	status    atomic.Int32
	badJSON   atomic.Bool
	keyCalls  atomic.Int32
	userCalls atomic.Int32
	bearer    atomic.Value
}

func newAuthFixture(t *testing.T) *authFixture {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	f := &authFixture{key: key}
	f.admin.Store(true)
	f.status.Store(http.StatusOK)
	f.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		switch req.URL.Path {
		case "/api/signing-keys/keys":
			f.keyCalls.Add(1)
			_ = json.NewEncoder(w).Encode(jose.JSONWebKeySet{Keys: []jose.JSONWebKey{{Key: &key.PublicKey, KeyID: "synthetic"}}})
		case "/api/users/17":
			f.userCalls.Add(1)
			f.bearer.Store(req.Header.Get("Authorization"))
			w.WriteHeader(int(f.status.Load()))
			if f.badJSON.Load() {
				_, _ = w.Write([]byte("invalid JSON"))
				return
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"isGrafanaAdmin": f.admin.Load(), "role": "Admin"})
		default:
			http.NotFound(w, req)
		}
	}))
	t.Cleanup(f.server.Close)
	return f
}

func signedID(t *testing.T, key *ecdsa.PrivateKey, subject string, expires time.Time) string {
	t.Helper()
	signer, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.ES256, Key: jose.JSONWebKey{Key: key, KeyID: "synthetic"}}, new(jose.SignerOptions).WithType("jwt"))
	if err != nil {
		t.Fatal(err)
	}
	token, err := jwt.Signed(signer).Claims(jwt.Claims{Subject: subject, Expiry: jwt.NewNumericDate(expires)}).Serialize()
	if err != nil {
		t.Fatal(err)
	}
	return token
}

func authRequest(identity, secret string) *http.Request {
	ctx := config.WithGrafanaConfig(context.Background(), config.NewGrafanaCfg(map[string]string{config.AppClientSecret: secret}))
	req := httptest.NewRequest("POST", "/command", strings.NewReader(`{"command":"list"}`)).WithContext(ctx)
	req.Header.Set("X-Grafana-Id", identity)
	req.Header.Set("X-Grafana-User", "forged-admin")
	req.Header.Set("X-Grafana-Role", "Admin")
	return req
}

func TestRoleRevocation(t *testing.T) {
	f := newAuthFixture(t)
	auth := newAuthorizer(f.server.URL)
	identity := signedID(t, f.key, "user:17", time.Now().Add(time.Minute))
	calls := 0
	c := &console{allowed: auth.canSubmit, execute: func(context.Context, string) (string, error) { calls++; return "ok", nil }}
	for _, step := range []struct {
		admin  bool
		secret string
		status int
	}{
		{true, "synthetic-first", 200},
		{false, "synthetic-second", 403},
		{false, "synthetic-second", 403},
		{true, "synthetic-second", 200},
	} {
		f.admin.Store(step.admin)
		w := httptest.NewRecorder()
		c.routes().ServeHTTP(w, authRequest(identity, step.secret))
		if w.Code != step.status {
			t.Fatalf("status=%d want=%d", w.Code, step.status)
		}
	}
	if calls != 2 || f.keyCalls.Load() != 1 || f.userCalls.Load() != 4 || f.bearer.Load() != "Bearer synthetic-second" {
		t.Fatalf("commands=%d keys=%d users=%d bearer=%v", calls, f.keyCalls.Load(), f.userCalls.Load(), f.bearer.Load())
	}
}

func TestIdentityFailures(t *testing.T) {
	for _, name := range []string{"missing ID", "forged ID", "expired ID", "wrong signature", "service account", "empty user", "missing secret", "user API failure", "invalid profile", "organization admin"} {
		t.Run(name, func(t *testing.T) {
			f := newAuthFixture(t)
			identity := signedID(t, f.key, "user:17", time.Now().Add(time.Minute))
			secret := "synthetic-service"
			status := http.StatusServiceUnavailable
			switch name {
			case "missing ID":
				identity = ""
			case "forged ID":
				identity = "admin"
			case "expired ID":
				identity = signedID(t, f.key, "user:17", time.Now().Add(-time.Hour))
			case "wrong signature":
				key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
				if err != nil {
					t.Fatal(err)
				}
				identity = signedID(t, key, "user:17", time.Now().Add(time.Minute))
			case "service account":
				identity = signedID(t, f.key, "service-account:17", time.Now().Add(time.Minute))
				status = http.StatusForbidden
			case "empty user":
				identity = signedID(t, f.key, "user:", time.Now().Add(time.Minute))
				status = http.StatusForbidden
			case "missing secret":
				secret = ""
			case "user API failure":
				f.status.Store(http.StatusInternalServerError)
			case "invalid profile":
				f.badJSON.Store(true)
			case "organization admin":
				f.admin.Store(false)
				status = http.StatusForbidden
			}
			calls := 0
			c := &console{allowed: newAuthorizer(f.server.URL).canSubmit, execute: func(context.Context, string) (string, error) { calls++; return "ok", nil }}
			w := httptest.NewRecorder()
			c.routes().ServeHTTP(w, authRequest(identity, secret))
			if w.Code != status || calls != 0 {
				t.Fatalf("status=%d want=%d commands=%d", w.Code, status, calls)
			}
		})
	}
}