package middleware

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"edge-infra.dev/pkg/edge/api/bsl/types"
	"edge-infra.dev/pkg/edge/api/totp"
	"edge-infra.dev/pkg/edge/client"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

func TestAuthMiddlewareTotpToken(t *testing.T) {
	assert := assert.New(t)

	totpSecret := "totp-secret"
	router := gin.Default()

	db, _, err := sqlmock.New()
	assert.NoError(err)

	router.Use(AuthMiddleware("", totpSecret, db))

	router.GET("/", func(c *gin.Context) { c.String(http.StatusOK, "OK") })

	ts := httptest.NewServer(router)
	defer ts.Close()
	req, err := http.NewRequest("GET", ts.URL+"/", http.NoBody)
	assert.NoError(err)

	totpToken, err := totp.GenerateTotp(totpSecret)
	assert.NoError(err)

	req.Header.Set("Authorization", client.TotpToken+" "+totpToken.Code)

	resp, err := http.DefaultClient.Do(req)
	assert.NoError(err)

	respBody, err := io.ReadAll(resp.Body)
	assert.NoError(err)
	assert.Equal(http.StatusOK, resp.StatusCode)
	assert.Equal("OK", string(respBody))

	totpToken, err = totp.GenerateTotp("invalid-secret")
	assert.NoError(err)

	req.Header.Set("Authorization", client.TotpToken+" "+totpToken.Code)

	resp, err = http.DefaultClient.Do(req)
	assert.NoError(err)
	assert.Equal(http.StatusForbidden, resp.StatusCode)

	respBody, err = io.ReadAll(resp.Body)
	assert.NoError(err)
	assert.NotEqual("OK", string(respBody))
}

func TestGetEdgeRoles(t *testing.T) {
	assert := assert.New(t)
	c := NewContext(context.Background(), &types.AuthUser{
		Organization: "test-org",
		Roles:        []string{"EDGE_BANNER_ADMIN", "TEST_ROLE_1", "TEST_ROLE_2"},
	})
	edgeRoles, err := GetEdgeRoles(c)
	assert.NoError(err)
	assert.Equal(len(edgeRoles), 1)
	assert.Equal(edgeRoles[0], "EDGE_BANNER_ADMIN")

	c = NewContext(context.Background(), &types.AuthUser{
		Organization: "test-org",
		Roles:        []string{"TEST_ROLE_1", "TEST_ROLE_2"},
	})
	_, err = GetEdgeRoles(c)
	assert.Error(err)
	assert.Equal(err.Error(), "edge role not found for context user")

	c = NewContext(context.Background(), &types.AuthUser{
		Organization: "test-org",
		Roles:        []string{"EDGE_BANNER_ADMIN", "TEST_ROLE_1", "TEST_ROLE_2", "EDGE_ORG_ADMIN"},
	})
	edgeRoles, err = GetEdgeRoles(c)
	assert.NoError(err)
	assert.Equal(len(edgeRoles), 2)
	assert.Equal(edgeRoles[0], "EDGE_ORG_ADMIN")
	assert.Equal(edgeRoles[1], "EDGE_BANNER_ADMIN")
}