package controller

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

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/service"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func setupDoubaoAssetProxyRouter(t *testing.T) *gin.Engine {
	t.Helper()

	oldMode := gin.Mode()
	gin.SetMode(gin.TestMode)
	t.Cleanup(func() {
		gin.SetMode(oldMode)
	})

	r := gin.New()
	r.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
	return r
}

func decodeDoubaoAssetErrorMessage(t *testing.T, body string) string {
	t.Helper()

	var payload struct {
		Error struct {
			Message string `json:"message"`
			Type    string `json:"type"`
		} `json:"error"`
	}
	require.NoError(t, common.Unmarshal([]byte(body), &payload))
	return payload.Error.Message
}

func setupDoubaoAssetProxyDB(t *testing.T) *gorm.DB {
	t.Helper()

	oldDB := model.DB
	oldSQLitePath := common.SQLitePath
	oldMemoryCacheEnabled := common.MemoryCacheEnabled
	oldIsMasterNode := common.IsMasterNode
	oldUsingSQLite := common.UsingSQLite
	oldUsingMySQL := common.UsingMySQL
	oldUsingPostgreSQL := common.UsingPostgreSQL
	oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")

	common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared"
	common.MemoryCacheEnabled = false
	common.IsMasterNode = false
	common.UsingSQLite = false
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	require.NoError(t, os.Setenv("SQL_DSN", "local"))
	require.NoError(t, model.InitDB())

	model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
	db := model.DB
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}, &model.UserAssetGroup{}))

	t.Cleanup(func() {
		_ = sqlDB.Close()
		model.DB = oldDB
		common.SQLitePath = oldSQLitePath
		common.MemoryCacheEnabled = oldMemoryCacheEnabled
		common.IsMasterNode = oldIsMasterNode
		common.UsingSQLite = oldUsingSQLite
		common.UsingMySQL = oldUsingMySQL
		common.UsingPostgreSQL = oldUsingPostgreSQL
		if hadSQLDSN {
			_ = os.Setenv("SQL_DSN", oldSQLDSN)
		} else {
			_ = os.Unsetenv("SQL_DSN")
		}
	})

	return db
}

func createDoubaoAssetProxyChannel(t *testing.T, db *gorm.DB, id int, channelType int, group string, key string, status int) {
	t.Helper()

	priority := int64(id)
	weight := uint(10)
	autoBan := 1
	require.NoError(t, db.Create(&model.Channel{
		Id:          id,
		Type:        channelType,
		Key:         key,
		Status:      status,
		Name:        fmt.Sprintf("channel-%d", id),
		Group:       group,
		Models:      "seedance-2",
		Priority:    &priority,
		Weight:      &weight,
		AutoBan:     &autoBan,
		CreatedTime: int64(id),
	}).Error)
}

func TestDoubaoAssetProxyMissingActionReturns400(t *testing.T) {
	router := setupDoubaoAssetProxyRouter(t)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Version=2024-01-01", nil)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Equal(t, "Action query parameter is required", decodeDoubaoAssetErrorMessage(t, w.Body.String()))
}

func TestDoubaoAssetProxyUnknownActionReturns400(t *testing.T) {
	router := setupDoubaoAssetProxyRouter(t)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=CreateRealPersonAuthSession", nil)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "unsupported asset Action")
}

func TestDoubaoAssetProxyInvalidJSONReturns400(t *testing.T) {
	router := setupDoubaoAssetProxyRouter(t)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"PageNumber":`))
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadRequest, w.Code)
	assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "invalid JSON body")
}

func TestDoubaoAssetProxyRejectsChinaMobileAssetGroupActions(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	fake := &fakeDoubaoAssetAdapter{}
	t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
	createDoubaoAssetProxyChannel(t, db, 61, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
	router := gin.New()
	router.Use(func(c *gin.Context) {
		c.Set("id", 10)
		common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
		c.Next()
	})
	router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)

	for _, action := range []string{"CreateAssetGroup", "ListAssetGroups", "GetAssetGroup", "UpdateAssetGroup", "DeleteAssetGroup"} {
		t.Run(action, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, strings.NewReader(`{"Id":"group-1","GroupType":"AIGC","Name":"g"}`))
			w := httptest.NewRecorder()
			req.Header.Set("Content-Type", "application/json")
			router.ServeHTTP(w, req)

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

func TestDoubaoAssetProxyDoesNotBlockNonChinaMobileAssetGroupAction(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	fake := &fakeDoubaoAssetAdapter{}
	t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeDoubaoVideoCompatibleAiping, fake))
	createDoubaoAssetProxyChannel(t, db, 63, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "video-generation-key", common.ChannelStatusEnabled)
	router := newDoubaoAssetProxyTestRouter()

	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssetGroups", strings.NewReader(`{"PageNumber":1,"PageSize":10}`))
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, service.AssetOperationAssetGroupList, fake.operation)
}

func TestDoubaoAssetProxyAutoGroupContinuesAfterMissingGroup(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	fake := &fakeDoubaoAssetAdapter{}
	t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
	createDoubaoAssetProxyChannel(t, db, 62, constant.ChannelTypeChinaMobileSeedance, "vip", "video-generation-key", common.ChannelStatusEnabled)
	router := gin.New()
	router.Use(func(c *gin.Context) {
		c.Set("id", 10)
		common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
		common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
		c.Next()
	})
	router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
	oldAutoGroups := doubaoAssetAutoGroups
	doubaoAssetAutoGroups = func(string) []string {
		return []string{"default", "vip"}
	}
	t.Cleanup(func() {
		doubaoAssetAutoGroups = oldAutoGroups
	})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
	w := httptest.NewRecorder()
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, service.AssetOperationAssetList, fake.operation)
}

func TestDoubaoAssetProxyForwardsAdapterResponseHeaders(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	fake := &fakeDoubaoAssetAdapter{responseHeader: http.Header{"X-Upstream-Request-Id": []string{"upstream-1"}, "Content-Length": []string{"999"}}}
	t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
	createDoubaoAssetProxyChannel(t, db, 62, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
	router := gin.New()
	router.Use(func(c *gin.Context) {
		c.Set("id", 10)
		common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
		c.Next()
	})
	router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
	w := httptest.NewRecorder()
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, "upstream-1", w.Header().Get("X-Upstream-Request-Id"))
	assert.NotEqual(t, "999", w.Header().Get("Content-Length"))
}

func TestDoubaoAssetProxyScopesChinaMobileCreateAndListRequests(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	fake := &fakeDoubaoAssetAdapter{}
	t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
	createDoubaoAssetProxyChannel(t, db, 71, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
	require.NoError(t, model.CreateUserAssetGroup(10, 71, "owned"))
	router := newDoubaoAssetProxyTestRouter()

	for _, tc := range []struct {
		action string
		body   string
	}{
		{"CreateAsset", `{"GroupId":"forged","Name":"n","URL":"https://example.com/a.png","AssetType":"Image"}`},
		{"ListAssets", `{"Filter":{"GroupType":"AIGC","GroupIds":["forged"]},"PageNumber":1,"PageSize":10}`},
	} {
		t.Run(tc.action, func(t *testing.T) {
			fake.requests = nil
			req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+tc.action, strings.NewReader(tc.body))
			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)
			require.Equal(t, http.StatusOK, w.Code)
			require.Len(t, fake.requests, 1)
			if tc.action == "CreateAsset" {
				assert.Equal(t, "owned", fake.requests[0].Body["GroupId"])
			} else {
				assert.Equal(t, []string{"owned"}, fake.requests[0].Body["Filter"].(map[string]any)["GroupIds"])
			}
		})
	}
}

func TestDoubaoAssetProxyRejectsChinaMobileListAssetsWithoutFilter(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	createDoubaoAssetProxyChannel(t, db, 73, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
	require.NoError(t, model.CreateUserAssetGroup(10, 73, "owned"))
	router := newDoubaoAssetProxyTestRouter()

	for _, body := range []string{`{"PageNumber":1,"PageSize":10}`, `{"Filter":null,"PageNumber":1,"PageSize":10}`} {
		w := httptest.NewRecorder()
		req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(body))
		router.ServeHTTP(w, req)

		assert.Equal(t, http.StatusBadRequest, w.Code)
		assert.Contains(t, decodeDoubaoAssetErrorMessage(t, w.Body.String()), "Filter.GroupType is required")
	}
}

func TestDoubaoAssetProxyHidesChinaMobileAssetsOwnedByAnotherUser(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	fake := &fakeDoubaoAssetAdapter{getAssetGroupID: "other"}
	t.Cleanup(service.OverrideAssetAdapterForTest(constant.ChannelTypeChinaMobileSeedance, fake))
	createDoubaoAssetProxyChannel(t, db, 72, constant.ChannelTypeChinaMobileSeedance, "default", "video-generation-key", common.ChannelStatusEnabled)
	require.NoError(t, model.CreateUserAssetGroup(10, 72, "owned"))
	router := newDoubaoAssetProxyTestRouter()

	for _, action := range []string{"GetAsset", "UpdateAsset", "DeleteAsset"} {
		t.Run(action, func(t *testing.T) {
			fake.requests = nil
			req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action="+action, strings.NewReader(`{"Id":"asset-other","Name":"n"}`))
			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)

			assert.Equal(t, http.StatusNotFound, w.Code)
			require.Len(t, fake.requests, 1)
			assert.Equal(t, service.AssetOperationAssetGet, fake.requests[0].Action.Operation)
		})
	}
}

func newDoubaoAssetProxyTestRouter() *gin.Engine {
	router := gin.New()
	router.Use(func(c *gin.Context) {
		c.Set("id", 10)
		common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
		c.Next()
	})
	router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
	return router
}

func TestDoubaoAssetProxyAutoGroupDoesNotBypassInvalidBinding(t *testing.T) {
	db := setupDoubaoAssetProxyDB(t)
	createDoubaoAssetProxyChannel(t, db, 1, constant.ChannelTypeChinaMobileSeedance, "other", "bad-key", common.ChannelStatusEnabled)
	createDoubaoAssetProxyChannel(t, db, 2, constant.ChannelTypeChinaMobileSeedance, "vip", "video-generation-key", common.ChannelStatusEnabled)
	require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeChinaMobileSeedance, "default", 1))
	router := gin.New()
	router.Use(func(c *gin.Context) {
		c.Set("id", 10)
		common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
		common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
		c.Next()
	})
	router.POST("/api/v1/volcengine/asset", DoubaoAssetProxy)
	oldAutoGroups := doubaoAssetAutoGroups
	doubaoAssetAutoGroups = func(string) []string {
		return []string{"default", "vip"}
	}
	t.Cleanup(func() {
		doubaoAssetAutoGroups = oldAutoGroups
	})

	req := httptest.NewRequest(http.MethodPost, "/api/v1/volcengine/asset?Action=ListAssets", strings.NewReader(`{"Filter":{"GroupType":"AIGC"},"PageNumber":1,"PageSize":10}`))
	w := httptest.NewRecorder()
	req.Header.Set("Content-Type", "application/json")
	router.ServeHTTP(w, req)

	assert.Equal(t, http.StatusBadGateway, w.Code)
	assert.Contains(t, w.Body.String(), service.AssetErrorBindingInvalid)
}

func TestEffectiveDoubaoAssetGroupUsesUsingGroupBeforeBlankTokenGroup(t *testing.T) {
	c, _ := gin.CreateTestContext(httptest.NewRecorder())
	common.SetContextKey(c, constant.ContextKeyTokenGroup, "")
	common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")

	assert.Equal(t, "default", effectiveDoubaoAssetGroup(c))
}

func TestConcreteDoubaoAssetGroupsForAutoUsesUserAutoGroups(t *testing.T) {
	c, _ := gin.CreateTestContext(httptest.NewRecorder())
	common.SetContextKey(c, constant.ContextKeyUsingGroup, "auto")
	common.SetContextKey(c, constant.ContextKeyTokenGroup, "")

	groups := concreteDoubaoAssetGroupsForRequest(c, func(string) []string {
		return []string{"default", "vip"}
	})

	assert.Equal(t, []string{"default", "vip"}, groups)
}

type fakeDoubaoAssetAdapter struct {
	operation       service.AssetOperation
	responseHeader  http.Header
	requests        []service.AssetRequest
	getAssetGroupID string
	calls           int
}

func (a *fakeDoubaoAssetAdapter) Name() string {
	return "fake_asset"
}

func (a *fakeDoubaoAssetAdapter) Supports(operation service.AssetOperation) bool {
	return true
}

func (a *fakeDoubaoAssetAdapter) DoAssetRequest(_ context.Context, _ *model.Channel, req service.AssetRequest) (*service.AssetUpstreamResponse, *service.AssetError) {
	a.operation = req.Action.Operation
	a.requests = append(a.requests, req)
	a.calls++
	if req.Action.Operation == service.AssetOperationAssetGroupCreate {
		body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{"GroupId": "generated-group"})
		if err != nil {
			return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
		}
		return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
	}
	if req.Action.Operation == service.AssetOperationAssetGet {
		body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{"Id": "asset-1", "GroupId": a.getAssetGroupID})
		if err != nil {
			return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
		}
		return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
	}
	body, err := service.BuildAssetSuccessResponse(req.Action.Action, req.Version, map[string]any{
		"Items": []map[string]any{
			{"Id": "group-1", "Name": "g", "GroupType": "AIGC"},
		},
		"TotalCount": 1,
	})
	if err != nil {
		return nil, &service.AssetError{Type: service.AssetErrorServer, Message: err.Error(), HTTPStatus: http.StatusInternalServerError}
	}
	return &service.AssetUpstreamResponse{StatusCode: http.StatusOK, Header: a.responseHeader, Body: body}, nil
}