package magpie

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"edge-infra.dev/pkg/edge/edgeencrypt"
	"edge-infra.dev/pkg/lib/fog"

	"github.com/golang-jwt/jwt"
	"github.com/google/uuid"
	"gotest.tools/v3/assert"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

var (
	data = []byte("my-data")
	cfg  = &Config{
		Namespace: edgeencrypt.DecryptionName,
		JWTSecret: edgeencrypt.DecryptionJWTSecret,
		KmsKey: edgeencrypt.KmsKey{
			Location:  "us-central1",
			ProjectID: "my-project",
		},
	}
	server *Server
	cl     client.Client
)

func TestMain(m *testing.M) {
	cl = fake.NewFakeClient()
	server = NewDecryptionServer(cfg, cl, fog.New(), nil)
	os.Exit(m.Run())
}

func TestHealth(t *testing.T) {
	w := httptest.NewRecorder()
	req, err := http.NewRequest("GET", "/health", nil)
	assert.NilError(t, err)
	server.router.ServeHTTP(w, req)

	assert.Equal(t, 200, w.Code)
	assert.Equal(t, "ok", w.Body.String())
}

func TestDecryptUnAuthorize(t *testing.T) {
	w := httptest.NewRecorder()

	req, err := http.NewRequest("POST", "/decrypt", strings.NewReader(string(data)))
	assert.NilError(t, err)
	server.router.ServeHTTP(w, req)

	// no bearer token
	assert.Equal(t, 401, w.Code)
}

func TestEncryptionSuccess(t *testing.T) {
	privateKey, pem := createPublicPrivateKey(t)

	// encrypt data to decrypt for testing
	encryptedData, err := edgeencrypt.EncryptData(pem, data)
	assert.NilError(t, err)

	e := &edgeencrypt.EncryptedData{
		BannerEdgeID: "my-banner",
		Channel:      "my-channel",
		ChannelID:    uuid.NewString(),
		KeyVersion:   "1",
		Data:         encryptedData,
	}
	encryptedData, err = e.ToEncryptionResponse()
	assert.NilError(t, err)

	w := httptest.NewRecorder()
	req, err := http.NewRequest("POST", "/v1/decrypt/"+e.Channel, strings.NewReader(string(encryptedData)))
	assert.NilError(t, err)

	token, err := edgeencrypt.CreateToken(jwt.SigningMethodRS256, privateKey, edgeencrypt.DefaultDuration, e.ChannelID, e.Channel, edgeencrypt.Decryption)
	assert.NilError(t, err)

	req.Header.Set("Authorization", "Bearer "+token)

	server.decrypt = func(_ context.Context, _, _, _ string, aesKey []byte) ([]byte, error) {
		return rsa.DecryptOAEP(sha256.New(), rand.Reader, privateKey, aesKey, nil)
	}
	server.router.ServeHTTP(w, req)

	assert.Equal(t, 200, w.Code)
	body, err := io.ReadAll(w.Body)
	assert.NilError(t, err)

	assert.Equal(t, string(data), string(body))
}

func createPublicPrivateKey(t *testing.T) (*rsa.PrivateKey, string) {
	privateKey, err := rsa.GenerateKey(rand.Reader, edgeencrypt.RSA2048)
	assert.NilError(t, err)

	publicKey := &privateKey.PublicKey

	pem, err := edgeencrypt.ConvertRSAPublicKeyToPEM(publicKey)
	assert.NilError(t, err)

	ee := &edgeencrypt.PublicKey{Version: "1", PEM: pem}

	// public key to validate bearer token
	err = ee.Save(context.Background(), cl, edgeencrypt.DecryptionName, edgeencrypt.DecryptionJWTSecret)
	assert.NilError(t, err)

	// public key to encrypt data
	channelEncryptionSecret := fmt.Sprintf(edgeencrypt.EncryptionSecret, "my-channel")
	err = ee.Save(context.Background(), cl, edgeencrypt.DecryptionName, channelEncryptionSecret)
	assert.NilError(t, err)

	return privateKey, pem
}