package agent

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"

	"security-center-backend/internal/config"
	"security-center-backend/internal/model"
	"security-center-backend/internal/repository"
	"security-center-backend/internal/service/jwt"
)

// rotateFakeRepo is an in-memory APIKeyRepository for the Rotate tests. It
// captures the record handed to Create and the (id, owner) pair handed to
// Revoke, and mirrors the real Revoke owner-scoping semantics.
type rotateFakeRepo struct {
	keys      map[uint64]*model.APIKey
	created   *model.APIKey
	createErr error
	revokeErr error

	revokeCalled         bool
	revokedID            uint64
	revokedOwner         uint64
	usedFindByID         bool
	usedFindByIDForOwner bool
}

func (f *rotateFakeRepo) Create(_ context.Context, key *model.APIKey) error {
	if f.createErr != nil {
		return f.createErr
	}
	key.ID = 100
	key.CreatedAt = time.Now()
	f.created = key
	f.keys[key.ID] = key
	return nil
}

func (f *rotateFakeRepo) FindByIDForOwner(_ context.Context, id, ownerUserID uint64) (*model.APIKey, error) {
	f.usedFindByIDForOwner = true
	key, ok := f.keys[id]
	if !ok || key.OwnerUserID != ownerUserID {
		return nil, nil
	}
	return key, nil
}

func (f *rotateFakeRepo) FindByID(_ context.Context, id uint64) (*model.APIKey, error) {
	f.usedFindByID = true
	key, ok := f.keys[id]
	if !ok {
		return nil, nil
	}
	return key, nil
}

func (f *rotateFakeRepo) Revoke(_ context.Context, id, ownerUserID uint64) error {
	f.revokeCalled = true
	f.revokedID = id
	f.revokedOwner = ownerUserID
	if f.revokeErr != nil {
		return f.revokeErr
	}
	key, ok := f.keys[id]
	if !ok || key.OwnerUserID != ownerUserID {
		return repository.ErrAPIKeyNotFoundOrNotOwned
	}
	key.Status = model.APIKeyStatusRevoked
	return nil
}

func (f *rotateFakeRepo) FindByHash(context.Context, string) (*model.APIKey, error) {
	return nil, nil
}
func (f *rotateFakeRepo) TouchLastUsed(context.Context, uint64) error { return nil }
func (f *rotateFakeRepo) ListByOwner(context.Context, uint64) ([]*model.APIKey, error) {
	return nil, nil
}
func (f *rotateFakeRepo) ListAll(context.Context, repository.ListAPIKeysFilter) ([]*model.APIKey, error) {
	return nil, nil
}

// rotateTestKey builds an active key whose lifetime ends `lifetime` after
// `createdAt`; pass a negative lifetime with an old createdAt for expired
// keys, or nil out ExpiresAt afterwards for a permanent key.
func rotateTestKey(id, owner uint64, createdAt time.Time, lifetime time.Duration) *model.APIKey {
	expires := createdAt.Add(lifetime)
	return &model.APIKey{
		Base:            model.Base{ID: id, CreatedAt: createdAt},
		KeyPrefix:       "sk-agent-sec-bot",
		KeyHash:         "old-hash",
		AgentName:       "Sec Bot",
		OwnerUserID:     owner,
		PermissionScope: model.StringArray{"caasm:view_assets"},
		Status:          model.APIKeyStatusActive,
		ExpiresAt:       &expires,
	}
}

// newRotateTestEngine mounts only the rotate endpoint behind claim injection.
func newRotateTestEngine(repo repository.APIKeyRepository, cfg config.AgentConfig, claims *jwt.Claims) *gin.Engine {
	gin.SetMode(gin.TestMode)
	engine := gin.New()
	h := New(repo, cfg)
	engine.POST("/api/v1/agent/api-keys/:id/rotate", injectClaims(claims), h.Rotate)
	return engine
}

func doRotate(t *testing.T, engine *gin.Engine, id string) (int, *CreateAPIKeyResponse, string) {
	t.Helper()
	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/api-keys/"+id+"/rotate", nil)
	w := httptest.NewRecorder()
	engine.ServeHTTP(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("http status = %d, want 200 (envelope convention); body=%s", w.Code, w.Body.String())
	}
	var envelope struct {
		Code int                   `json:"code"`
		Data *CreateAPIKeyResponse `json:"data"`
	}
	if err := json.Unmarshal(w.Body.Bytes(), &envelope); err != nil {
		t.Fatalf("decode API response: %v; body=%s", err, w.Body.String())
	}
	return envelope.Code, envelope.Data, w.Body.String()
}

func TestRotateOwnKeySucceeds(t *testing.T) {
	createdAt := time.Now().Add(-24 * time.Hour)
	lifetime := 48 * time.Hour
	old := rotateTestKey(1, 7, createdAt, lifetime)
	repo := &rotateFakeRepo{keys: map[uint64]*model.APIKey{1: old}}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-"}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 7})

	t0 := time.Now()
	code, data, raw := doRotate(t, engine, "1")
	t1 := time.Now()

	if code != 0 {
		t.Fatalf("body code = %d, want 0 (success)", code)
	}
	if data == nil || data.APIKey == "" {
		t.Fatalf("data = %+v, want a rotated key with plaintext", data)
	}
	if !strings.HasPrefix(data.APIKey, "sk-agent-sec-bot-") {
		t.Fatalf("plaintext %q does not follow <prefix><slug>-<hex>", data.APIKey)
	}

	// 新记录沿用旧记录的 name/scope/owner。
	newKey := repo.created
	if newKey == nil {
		t.Fatal("repo.Create was not called")
	}
	if newKey.AgentName != old.AgentName || newKey.OwnerUserID != old.OwnerUserID {
		t.Fatalf("created key carries name=%q owner=%d, want name=%q owner=%d",
			newKey.AgentName, newKey.OwnerUserID, old.AgentName, old.OwnerUserID)
	}
	if !slices.Equal([]string(newKey.PermissionScope), []string(old.PermissionScope)) {
		t.Fatalf("created scope = %v, want %v", newKey.PermissionScope, old.PermissionScope)
	}
	if newKey.Status != model.APIKeyStatusActive {
		t.Fatalf("created status = %q, want active", newKey.Status)
	}

	// 只存哈希与前 16 位，且与响应明文对得上。
	if newKey.KeyHash != sha256Hex(data.APIKey) {
		t.Fatalf("stored hash does not match sha256 of the returned plaintext")
	}
	if newKey.KeyHash == data.APIKey {
		t.Fatal("plaintext was persisted into key_hash")
	}
	if newKey.KeyPrefix != data.APIKey[:16] {
		t.Fatalf("key_prefix = %q, want first 16 chars of plaintext", newKey.KeyPrefix)
	}

	// 续期策略：新 ExpiresAt = now + 原生命周期（48h），必然晚于旧 ExpiresAt。
	if newKey.ExpiresAt == nil {
		t.Fatal("created ExpiresAt = nil, want renewed expiry")
	}
	if !newKey.ExpiresAt.After(*old.ExpiresAt) {
		t.Fatalf("new ExpiresAt %v not after old %v", newKey.ExpiresAt, old.ExpiresAt)
	}
	if newKey.ExpiresAt.Before(t0.Add(lifetime)) || newKey.ExpiresAt.After(t1.Add(lifetime)) {
		t.Fatalf("new ExpiresAt %v outside now+48h window [%v, %v]",
			newKey.ExpiresAt, t0.Add(lifetime), t1.Add(lifetime))
	}

	// 旧记录在创建成功后被吊销，且按旧记录的真实 owner 吊销。
	if !repo.revokeCalled || repo.revokedID != 1 || repo.revokedOwner != 7 {
		t.Fatalf("revoke called=%v id=%d owner=%d, want id=1 owner=7",
			repo.revokeCalled, repo.revokedID, repo.revokedOwner)
	}
	if repo.keys[1].Status != model.APIKeyStatusRevoked {
		t.Fatalf("old key status = %q, want revoked", repo.keys[1].Status)
	}

	// 明文只在响应体的 api_key 字段出现一次，哈希不外泄。
	if strings.Count(raw, data.APIKey) != 1 {
		t.Fatalf("plaintext appears %d times in body, want exactly 1; body=%s",
			strings.Count(raw, data.APIKey), raw)
	}
	if strings.Contains(raw, "key_hash") || strings.Contains(raw, newKey.KeyHash) {
		t.Fatalf("key hash leaked into response body: %s", raw)
	}
}

func TestRotateOthersKeyReturns404(t *testing.T) {
	createdAt := time.Now().Add(-24 * time.Hour)
	repo := &rotateFakeRepo{keys: map[uint64]*model.APIKey{
		1: rotateTestKey(1, 8, createdAt, 48*time.Hour), // 属于别人
	}}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-"}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 7})

	code, _, _ := doRotate(t, engine, "1")
	if code != http.StatusNotFound {
		t.Fatalf("body code = %d, want 404", code)
	}
	if repo.created != nil || repo.revokeCalled {
		t.Fatal("no key should be created or revoked on 404")
	}
}

func TestRotateRevokedKeyReturns409(t *testing.T) {
	createdAt := time.Now().Add(-24 * time.Hour)
	old := rotateTestKey(1, 7, createdAt, 48*time.Hour)
	old.Status = model.APIKeyStatusRevoked
	repo := &rotateFakeRepo{keys: map[uint64]*model.APIKey{1: old}}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-"}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 7})

	code, _, _ := doRotate(t, engine, "1")
	if code != http.StatusConflict {
		t.Fatalf("body code = %d, want 409", code)
	}
	if repo.created != nil {
		t.Fatal("no new key should be created for an already revoked key")
	}
}

func TestRotateNonNumericIDReturns400(t *testing.T) {
	repo := &rotateFakeRepo{keys: map[uint64]*model.APIKey{}}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-"}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 7})

	code, _, _ := doRotate(t, engine, "abc")
	if code != http.StatusBadRequest {
		t.Fatalf("body code = %d, want 400", code)
	}
}

func TestRotatePermanentKeyRejectedWhenNotAllowed(t *testing.T) {
	createdAt := time.Now().Add(-24 * time.Hour)
	old := rotateTestKey(1, 7, createdAt, 48*time.Hour)
	old.ExpiresAt = nil // 永久 key
	repo := &rotateFakeRepo{keys: map[uint64]*model.APIKey{1: old}}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-", AllowPermanentKey: false}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 7})

	code, _, _ := doRotate(t, engine, "1")
	if code != http.StatusBadRequest {
		t.Fatalf("body code = %d, want 400", code)
	}
	if repo.created != nil || repo.revokeCalled {
		t.Fatal("permanent key rotation must be rejected before any write")
	}
}

func TestRotatePermanentKeyAllowedStaysPermanent(t *testing.T) {
	createdAt := time.Now().Add(-24 * time.Hour)
	old := rotateTestKey(1, 7, createdAt, 48*time.Hour)
	old.ExpiresAt = nil
	repo := &rotateFakeRepo{keys: map[uint64]*model.APIKey{1: old}}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-", AllowPermanentKey: true}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 7})

	code, data, _ := doRotate(t, engine, "1")
	if code != 0 {
		t.Fatalf("body code = %d, want 0 (success)", code)
	}
	if data == nil || data.ExpiresAt != nil || repo.created.ExpiresAt != nil {
		t.Fatalf("permanent key must rotate into a permanent key, data=%+v", data)
	}
	if repo.keys[1].Status != model.APIKeyStatusRevoked {
		t.Fatalf("old key status = %q, want revoked", repo.keys[1].Status)
	}
}

func TestRotateAsAdminRotatesOthersKey(t *testing.T) {
	createdAt := time.Now().Add(-24 * time.Hour)
	repo := &rotateFakeRepo{keys: map[uint64]*model.APIKey{
		1: rotateTestKey(1, 7, createdAt, 48*time.Hour), // admin(99) 轮换 user 7 的 key
	}}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-"}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 99, Roles: []string{"admin"}})

	code, data, _ := doRotate(t, engine, "1")
	if code != 0 {
		t.Fatalf("body code = %d, want 0 (success)", code)
	}
	if !repo.usedFindByID || repo.usedFindByIDForOwner {
		t.Fatal("admin path must use the unscoped FindByID lookup")
	}
	if data == nil || data.APIKey == "" || repo.created.OwnerUserID != 7 {
		t.Fatalf("rotated key must stay owned by 7, data=%+v", data)
	}
	if repo.revokedOwner != 7 {
		t.Fatalf("old key revoked under owner %d, want its real owner 7", repo.revokedOwner)
	}
}

func TestRotateRevokeFailureReturnsErrorButKeepsNewKey(t *testing.T) {
	createdAt := time.Now().Add(-24 * time.Hour)
	repo := &rotateFakeRepo{
		keys:      map[uint64]*model.APIKey{1: rotateTestKey(1, 7, createdAt, 48*time.Hour)},
		revokeErr: errors.New("db down"),
	}
	cfg := config.AgentConfig{Enabled: true, KeyPrefix: "sk-agent-"}
	engine := newRotateTestEngine(repo, cfg, &jwt.Claims{UserID: 7})

	code, data, _ := doRotate(t, engine, "1")
	if code != http.StatusInternalServerError {
		t.Fatalf("body code = %d, want 500", code)
	}
	if data != nil {
		t.Fatalf("no key material should be returned on failure, data=%+v", data)
	}
	// 取舍：新 key 已落库则不删除、不回滚，交由调用方重试。
	if repo.created == nil {
		t.Fatal("new key must be kept (no rollback) when revoking the old one fails")
	}
}
