package apis

import (
	"bytes"
	"clortho/lib/db"
	"clortho/lib/users"
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"net/http"
	"net/http/httptest"
	"testing"
)


func setupRouter() *gin.Engine {
	r := gin.Default()
	private := r.Group("/gui")
	InitKeysEndpoints(private)
	return r
}

func createTestUser(admin bool) *db.User {
	username := "testuser"
	if admin {
		username = "admin"
	}
	displayName := username
	user := db.User{
		Username:    username,
		DisplayName: &displayName,
		Admin:       admin,
	}
	db.Connection.Create(&user)
	return &user
}

func createTestKey(userID uint) *db.Key {
	key := db.Key{
		UserID:  userID,
		Content: "test-key-content",
	}
	db.Connection.Create(&key)
	return &key
}

func createTestSession(user *db.User) *db.UserSession {
	session := users.NewSession(*user)
	return session
}

func TestGetKeys(t *testing.T) {
	// Setup
	router := setupRouter()

	// Create admin user and session
	adminUser := createTestUser(true)
	adminSession := createTestSession(adminUser)

	// Create regular user and session
	regularUser := createTestUser(false)
	regularSession := createTestSession(regularUser)

	// Create keys for both users
	_ = createTestKey(adminUser.ID) // Create a key for admin to test that admin can see all keys
	regularKey := createTestKey(regularUser.ID)

	// Test admin user can get all keys
	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", "/gui/keys/", nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusUnauthorized, w.Code)

	// Set admin session in context
	router.Use(func(c *gin.Context) {
		c.Set("session", adminSession)
		c.Next()
	})

	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", "/gui/keys/", nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	var keys []db.Key
	err := json.Unmarshal(w.Body.Bytes(), &keys)
	assert.NoError(t, err)
	assert.Len(t, keys, 2)

	// Test regular user can only get their own keys
	router = setupRouter()
	router.Use(func(c *gin.Context) {
		c.Set("session", regularSession)
		c.Next()
	})

	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", "/gui/keys/", nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	err = json.Unmarshal(w.Body.Bytes(), &keys)
	assert.NoError(t, err)
	assert.Len(t, keys, 1)
	assert.Equal(t, regularKey.ID, keys[0].ID)
}

func TestGetKey(t *testing.T) {
	// Setup
	router := setupRouter()

	// Create admin user and session
	adminUser := createTestUser(true)
	adminSession := createTestSession(adminUser)

	// Create regular user and session
	regularUser := createTestUser(false)
	regularSession := createTestSession(regularUser)

	// Create keys for both users
	adminKey := createTestKey(adminUser.ID)
	regularKey := createTestKey(regularUser.ID)

	// Test admin user can get any key
	router.Use(func(c *gin.Context) {
		c.Set("session", adminSession)
		c.Next()
	})

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", fmt.Sprintf("/gui/keys/%d", regularKey.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	var key db.Key
	err := json.Unmarshal(w.Body.Bytes(), &key)
	assert.NoError(t, err)
	assert.Equal(t, regularKey.ID, key.ID)

	// Test regular user can only get their own key
	router = setupRouter()
	router.Use(func(c *gin.Context) {
		c.Set("session", regularSession)
		c.Next()
	})

	// Regular user can get their own key
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", fmt.Sprintf("/gui/keys/%d", regularKey.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	// Regular user cannot get admin's key
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", fmt.Sprintf("/gui/keys/%d", adminKey.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusForbidden, w.Code)
}

func TestCreateKey(t *testing.T) {
	// Setup
	router := setupRouter()

	// Create admin user and session
	adminUser := createTestUser(true)
	adminSession := createTestSession(adminUser)

	// Create regular user and session
	regularUser := createTestUser(false)
	regularSession := createTestSession(regularUser)

	// Test admin user can create a key
	router.Use(func(c *gin.Context) {
		c.Set("session", adminSession)
		c.Next()
	})

	keyRequest := keyRequest{
		UserID:  regularUser.ID,
		Content: "new-key-content",
	}

	requestBody, _ := json.Marshal(keyRequest)

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", "/gui/keys/", bytes.NewBuffer(requestBody))
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusCreated, w.Code)

	var key db.Key
	err := json.Unmarshal(w.Body.Bytes(), &key)
	assert.NoError(t, err)
	assert.Equal(t, regularUser.ID, key.UserID)
	assert.Equal(t, "new-key-content", key.Content)

	// Test regular user cannot create a key (only admins can)
	router = setupRouter()
	router.Use(func(c *gin.Context) {
		c.Set("session", regularSession)
		c.Next()
	})

	w = httptest.NewRecorder()
	req, _ = http.NewRequest("POST", "/gui/keys/", bytes.NewBuffer(requestBody))
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusForbidden, w.Code)
}

func TestUpdateKey(t *testing.T) {
	// Setup
	router := setupRouter()

	// Create admin user and session
	adminUser := createTestUser(true)
	adminSession := createTestSession(adminUser)

	// Create regular user and session
	regularUser := createTestUser(false)
	regularSession := createTestSession(regularUser)

	// Create keys for both users
	_ = createTestKey(adminUser.ID) // Create a key for admin but we'll only test updating the regular user's key
	regularKey := createTestKey(regularUser.ID)

	// Test admin user can update any key
	router.Use(func(c *gin.Context) {
		c.Set("session", adminSession)
		c.Next()
	})

	keyRequest := keyRequest{
		UserID:  regularUser.ID,
		Content: "updated-key-content",
	}

	requestBody, _ := json.Marshal(keyRequest)

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("PUT", fmt.Sprintf("/gui/keys/%d", regularKey.ID), bytes.NewBuffer(requestBody))
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	var key db.Key
	err := json.Unmarshal(w.Body.Bytes(), &key)
	assert.NoError(t, err)
	assert.Equal(t, regularKey.ID, key.ID)
	assert.Equal(t, "updated-key-content", key.Content)

	// Test regular user cannot update a key (only admins can)
	router = setupRouter()
	router.Use(func(c *gin.Context) {
		c.Set("session", regularSession)
		c.Next()
	})

	w = httptest.NewRecorder()
	req, _ = http.NewRequest("PUT", fmt.Sprintf("/gui/keys/%d", regularKey.ID), bytes.NewBuffer(requestBody))
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusForbidden, w.Code)
}

func TestDeleteKey(t *testing.T) {
	// Setup
	router := setupRouter()

	// Create admin user and session
	adminUser := createTestUser(true)
	adminSession := createTestSession(adminUser)

	// Create regular user and session
	regularUser := createTestUser(false)
	regularSession := createTestSession(regularUser)

	// Create keys for both users
	adminKey := createTestKey(adminUser.ID)
	regularKey := createTestKey(regularUser.ID)

	// Test admin user can delete any key
	router.Use(func(c *gin.Context) {
		c.Set("session", adminSession)
		c.Next()
	})

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("DELETE", fmt.Sprintf("/gui/keys/%d", regularKey.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	// Test regular user cannot delete a key (only admins can)
	router = setupRouter()
	router.Use(func(c *gin.Context) {
		c.Set("session", regularSession)
		c.Next()
	})

	w = httptest.NewRecorder()
	req, _ = http.NewRequest("DELETE", fmt.Sprintf("/gui/keys/%d", adminKey.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusForbidden, w.Code)
}

func TestGetUserKeys(t *testing.T) {
	// Setup
	router := setupRouter()

	// Create admin user and session
	adminUser := createTestUser(true)
	adminSession := createTestSession(adminUser)

	// Create regular user and session
	regularUser := createTestUser(false)
	regularSession := createTestSession(regularUser)

	// Create keys for both users
	_ = createTestKey(adminUser.ID) // Create a key for admin to test that regular users can't access admin keys
	regularKey := createTestKey(regularUser.ID)

	// Test admin user can get keys for any user
	router.Use(func(c *gin.Context) {
		c.Set("session", adminSession)
		c.Next()
	})

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("GET", fmt.Sprintf("/gui/users/%d/keys", regularUser.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	var keys []db.Key
	err := json.Unmarshal(w.Body.Bytes(), &keys)
	assert.NoError(t, err)
	assert.Len(t, keys, 1)
	assert.Equal(t, regularKey.ID, keys[0].ID)

	// Test regular user can only get their own keys
	router = setupRouter()
	router.Use(func(c *gin.Context) {
		c.Set("session", regularSession)
		c.Next()
	})

	// Regular user can get their own keys
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", fmt.Sprintf("/gui/users/%d/keys", regularUser.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)

	// Regular user cannot get admin's keys
	w = httptest.NewRecorder()
	req, _ = http.NewRequest("GET", fmt.Sprintf("/gui/users/%d/keys", adminUser.ID), nil)
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusForbidden, w.Code)
}

func TestCreateUserKey(t *testing.T) {
	// Setup
	router := setupRouter()

	// Create admin user and session
	adminUser := createTestUser(true)
	adminSession := createTestSession(adminUser)

	// Create regular user and session
	regularUser := createTestUser(false)
	regularSession := createTestSession(regularUser)

	// Test admin user can create a key for any user
	router.Use(func(c *gin.Context) {
		c.Set("session", adminSession)
		c.Next()
	})

	keyContent := "new-user-key-content"
	requestBody, _ := json.Marshal(map[string]string{"content": keyContent})

	w := httptest.NewRecorder()
	req, _ := http.NewRequest("POST", fmt.Sprintf("/gui/users/%d/keys", regularUser.ID), bytes.NewBuffer(requestBody))
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusCreated, w.Code)

	var key db.Key
	err := json.Unmarshal(w.Body.Bytes(), &key)
	assert.NoError(t, err)
	assert.Equal(t, regularUser.ID, key.UserID)
	assert.Equal(t, keyContent, key.Content)

	// Test regular user cannot create a key (only admins can)
	router = setupRouter()
	router.Use(func(c *gin.Context) {
		c.Set("session", regularSession)
		c.Next()
	})

	w = httptest.NewRecorder()
	req, _ = http.NewRequest("POST", fmt.Sprintf("/gui/users/%d/keys", regularUser.ID), bytes.NewBuffer(requestBody))
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusForbidden, w.Code)
}