//go:build e2e

// Package api contains e2e tests for REST API endpoints and token auth.
package api

import (
	"fmt"
	"testing"

	"github.com/pingcap/failpoint"

	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
)

func setupTokenTest(t *testing.T) *token.KeyManager {
	t.Helper()
	cfgDir := t.TempDir()
	failpoint.Enable("agentteam-config-dir", fmt.Sprintf(`return("%s")`, cfgDir))
	t.Cleanup(func() { failpoint.Disable("agentteam-config-dir") })
	mgr, err := token.NewKeyManager("")
	if err != nil {
		t.Fatalf("NewKeyManager: %v", err)
	}
	return mgr
}

func TestToken_GenerateAndValidate(t *testing.T) {
	mgr := setupTokenTest(t)

	key, err := mgr.Generate()
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if len(key) < 10 {
		t.Fatalf("key too short: %s", key)
	}
	prefix := "sk-dev."
	if key[:len(prefix)] != prefix {
		t.Errorf("key=%s, want prefix %s", key, prefix)
	}
	if err := mgr.Validate(key); err != nil {
		t.Fatalf("Validate: %v", err)
	}
}

func TestToken_ValidateInvalidFormat(t *testing.T) {
	mgr := setupTokenTest(t)

	cases := []struct {
		name string
		key  string
	}{
		{"empty", ""},
		{"too short", "sk-dev"},
		{"no dots", "sk-dev-abc"},
		{"one dot", "sk-dev.abc"},
		{"invalid base64 payload", "sk-dev.!!!.abc"},
		{"invalid base64 signature", "sk-dev.YWJj.!!!"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if err := mgr.Validate(tc.key); err == nil {
				t.Fatalf("expected error for key=%q", tc.key)
			}
		})
	}
}

func TestToken_ValidateTamperedPayload(t *testing.T) {
	mgr := setupTokenTest(t)
	key, err := mgr.Generate()
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if err := mgr.Validate(key + "x"); err == nil {
		t.Fatal("expected error for tampered key")
	}
}

func TestToken_ValidateDifferentKeyManager(t *testing.T) {
	mgr1 := setupTokenTest(t)
	mgr2 := setupTokenTest(t)

	key, err := mgr1.Generate()
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	if err := mgr2.Validate(key); err == nil {
		t.Fatal("expected error: key from different KeyManager should fail")
	}
}

func TestToken_AgentTokenGenerateAndValidate(t *testing.T) {
	mgr := setupTokenTest(t)

	at, err := mgr.GenerateAgentTokenForServer("my-agent", "")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer: %v", err)
	}
	if len(at) < 10 {
		t.Fatalf("agent token too short: %s", at)
	}
	prefix := "at-dev."
	if at[:len(prefix)] != prefix {
		t.Errorf("token=%s, want prefix %s", at, prefix)
	}
	if _, err := mgr.ValidateAgentTokenInfo("my-agent", at); err != nil {
		t.Fatalf("ValidateAgentTokenInfo: %v", err)
	}
}

func TestToken_AgentTokenMismatch(t *testing.T) {
	mgr := setupTokenTest(t)

	at, err := mgr.GenerateAgentTokenForServer("agent-a", "")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer: %v", err)
	}
	if _, err := mgr.ValidateAgentTokenInfo("agent-b", at); err == nil {
		t.Fatal("expected error for agent mismatch")
	}
}

func TestToken_AgentTokenInvalidFormat(t *testing.T) {
	mgr := setupTokenTest(t)

	cases := []struct {
		name  string
		token string
	}{
		{"empty", ""},
		{"too short", "at-dev"},
		{"no dots", "at-dev-abc"},
		{"one dot", "at-dev.abc"},
		{"invalid base64 payload", "at-dev.!!!.abc"},
		{"invalid base64 sig", "at-dev.YQ.!!!"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := mgr.ValidateAgentTokenInfo("any-agent", tc.token); err == nil {
				t.Fatalf("expected error for token=%q", tc.token)
			}
		})
	}
}

func TestToken_AgentTokenEmptyID(t *testing.T) {
	mgr := setupTokenTest(t)
	if _, err := mgr.GenerateAgentTokenForServer("", ""); err == nil {
		t.Fatal("expected error for empty agent ID")
	}
}

func TestToken_AgentTokenWrongKey(t *testing.T) {
	mgr1 := setupTokenTest(t)
	mgr2 := setupTokenTest(t)

	at, err := mgr1.GenerateAgentTokenForServer("agent-1", "")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer: %v", err)
	}
	if _, err := mgr2.ValidateAgentTokenInfo("agent-1", at); err == nil {
		t.Fatal("expected error: token signed by different key")
	}
}

func TestToken_MultipleAgents(t *testing.T) {
	mgr := setupTokenTest(t)

	at1, _ := mgr.GenerateAgentTokenForServer("agent-1", "")
	at2, _ := mgr.GenerateAgentTokenForServer("agent-2", "")

	if _, err := mgr.ValidateAgentTokenInfo("agent-1", at1); err != nil {
		t.Fatalf("agent-1 validate own token: %v", err)
	}
	if _, err := mgr.ValidateAgentTokenInfo("agent-2", at2); err != nil {
		t.Fatalf("agent-2 validate own token: %v", err)
	}
	if _, err := mgr.ValidateAgentTokenInfo("agent-1", at2); err == nil {
		t.Fatal("agent-1 should not validate agent-2 token")
	}
	if _, err := mgr.ValidateAgentTokenInfo("agent-2", at1); err == nil {
		t.Fatal("agent-2 should not validate agent-1 token")
	}
}

func TestToken_KeyPairPersistence(t *testing.T) {
	cfgDir := t.TempDir()
	failpoint.Enable("agentteam-config-dir", fmt.Sprintf(`return("%s")`, cfgDir))
	mgr1, err := token.NewKeyManager("")
	if err != nil {
		t.Fatalf("first NewKeyManager: %v", err)
	}
	key, err := mgr1.Generate()
	if err != nil {
		t.Fatalf("Generate: %v", err)
	}
	failpoint.Disable("agentteam-config-dir")
	t.Cleanup(func() { failpoint.Disable("agentteam-config-dir") })

	failpoint.Enable("agentteam-config-dir", fmt.Sprintf(`return("%s")`, cfgDir))
	mgr2, err := token.NewKeyManager("")
	if err != nil {
		t.Fatalf("second NewKeyManager: %v", err)
	}
	failpoint.Disable("agentteam-config-dir")

	if err := mgr2.Validate(key); err != nil {
		t.Fatalf("key from mgr1 should be valid with mgr2: %v", err)
	}
}

func TestToken_RoundtripAPIKeys(t *testing.T) {
	mgr := setupTokenTest(t)
	for i := 0; i < 10; i++ {
		key, err := mgr.Generate()
		if err != nil {
			t.Fatalf("Generate %d: %v", i, err)
		}
		if err := mgr.Validate(key); err != nil {
			t.Fatalf("Validate %d: %v", i, err)
		}
	}
}

func TestToken_RoundtripAgentTokens(t *testing.T) {
	mgr := setupTokenTest(t)
	for i := 0; i < 10; i++ {
		agentID := fmt.Sprintf("agent-%d", i)
		at, err := mgr.GenerateAgentTokenForServer(agentID, "")
		if err != nil {
			t.Fatalf("GenerateAgentTokenForServer %d: %v", i, err)
		}
		if _, err := mgr.ValidateAgentTokenInfo(agentID, at); err != nil {
			t.Fatalf("ValidateAgentTokenInfo %d: %v", i, err)
		}
	}
}