package controller

import (
	"fmt"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/setting/operation_setting"
	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidwall/gjson"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
)

func TestBotBalanceQuery(t *testing.T) {
	key := strings.Repeat("b", 48)
	t.Setenv("BOT_BALANCE_KEY", key)
	oldUnit, oldRate := common.QuotaPerUnit, operation_setting.USDExchangeRate
	common.QuotaPerUnit, operation_setting.USDExchangeRate = 500000, 7.3
	t.Cleanup(func() { common.QuotaPerUnit, operation_setting.USDExchangeRate = oldUnit, oldRate })
	for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
		t.Run(dialect, func(t *testing.T) {
			var driver gorm.Dialector = sqlite.Open(":memory:")
			if dialect == "mysql" {
				require.NotEmpty(t, os.Getenv("TEST_MYSQL_DSN"))
				driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
			}
			if dialect == "postgres" {
				require.NotEmpty(t, os.Getenv("TEST_POSTGRES_DSN"))
				driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
			}
			db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("bot_%d_", time.Now().UnixNano())}})
			require.NoError(t, err)
			conn, err := db.DB()
			require.NoError(t, err)
			conn.SetMaxOpenConns(1)
			oldDB := model.DB
			model.DB = db
			t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(&model.User{})); model.DB = oldDB; _ = conn.Close() })
			require.NoError(t, db.AutoMigrate(&model.User{}))
			rows := []model.User{
				{Id: 1, Username: "bot1", Email: "a@example.com", Status: 1, Quota: 250000},
				{Id: 2, Username: "bot2", Email: "dup@example.com", Status: 1, Quota: 0},
				{Id: 3, Username: "bot3", Email: "dup@example.com", Status: 1, Quota: -500000},
				{Id: 4, Username: "bot4", Email: "disabled@example.com", Status: 2, Quota: 900000},
				{Id: 5, Username: "bot5", Email: "gone@example.com", Status: 1},
			}
			for i := range rows {
				rows[i].AffCode = fmt.Sprintf("bot%d", i)
			}
			require.NoError(t, db.Create(&rows).Error)
			require.NoError(t, db.Delete(&rows[4]).Error)
			r := gin.New()
			r.POST("/api/bot/balances", BotBalanceQuery)
			for _, tc := range []struct {
				name, body, token string
				status            int
				code              string
			}{
				{"missing key", `{}`, "", 401, "unauthorized"},
				{"wrong key", `{}`, "Bearer wrong", 401, "unauthorized"},
				{"empty", `{}`, "Bearer " + key, 400, "invalid_request"},
				{"both", `{"email":"a@example.com","ids":[1]}`, "Bearer " + key, 400, "invalid_request"},
				{"bad id", `{"ids":[0]}`, "Bearer " + key, 400, "invalid_request"},
				{"bad email", `{"email":"invalid"}`, "Bearer " + key, 400, "invalid_email"},
				{"deleted", `{"email":"gone@example.com"}`, "Bearer " + key, 409, "not_found"},
				{"duplicate", `{"email":"dup@example.com"}`, "Bearer " + key, 409, "ambiguous_email"},
				{"exact not substring", `{"email":"aa@example.com"}`, "Bearer " + key, 409, "not_found"},
				{"normalized", `{"email":" A@EXAMPLE.COM "}`, "Bearer " + key, 200, ""},
				{"batch", `{"ids":[1,2,3,4,5,999]}`, "Bearer " + key, 200, ""},
			} {
				t.Run(tc.name, func(t *testing.T) {
					w := httptest.NewRecorder()
					req := httptest.NewRequest("POST", "/api/bot/balances", strings.NewReader(tc.body))
					req.Header.Set("Authorization", tc.token)
					req.Header.Set("Content-Type", "application/json")
					r.ServeHTTP(w, req)
					require.Equal(t, tc.status, w.Code, w.Body.String())
					assert.Equal(t, "no-store", w.Header().Get("Cache-Control"))
					assert.Equal(t, tc.code, gjson.Get(w.Body.String(), "code").String())
					if tc.status == 200 {
						assert.NotContains(t, w.Body.String(), "email")
						assert.NotContains(t, w.Body.String(), "quota")
						assert.NotContains(t, w.Body.String(), "password")
						assert.Equal(t, "3.65", gjson.Get(w.Body.String(), "data.0.balance_cny").String())
						if tc.name == "batch" {
							assert.Equal(t, int64(4), gjson.Get(w.Body.String(), "data.#").Int())
							assert.Equal(t, "0", gjson.Get(w.Body.String(), "data.1.balance_cny").String())
							assert.Equal(t, "-7.3", gjson.Get(w.Body.String(), "data.2.balance_cny").String())
							assert.False(t, gjson.Get(w.Body.String(), "data.3.active").Bool())
							assert.False(t, gjson.Get(w.Body.String(), "data.3.balance_cny").Exists())
						}
					}
				})
			}
			t.Setenv("BOT_BALANCE_KEY", "")
			w := httptest.NewRecorder()
			r.ServeHTTP(w, httptest.NewRequest("POST", "/api/bot/balances", strings.NewReader(`{}`)))
			assert.Equal(t, 401, w.Code)
			t.Setenv("BOT_BALANCE_KEY", key)
		})
	}
}