package token

import (
	"crypto/ed25519"
	"encoding/base64"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"go.yaml.in/yaml/v3"

	"github.com/tidbcloud/AgentTeam/desktop/internal/crypto"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
)

func TestNewManager_CreatesKeyPair(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")

	mgr, err := NewKeyManager(path)
	if err != nil {
		t.Fatalf("NewKeyManager() error = %v", err)
	}

	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("failed to read config file: %v", err)
	}

	var cfg profile.Config
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		t.Fatalf("failed to unmarshal config: %v", err)
	}

	if len(cfg.Security.PublicKey) != ed25519.PublicKeySize {
		t.Errorf("public key size = %d, want %d", len(cfg.Security.PublicKey), ed25519.PublicKeySize)
	}
	if len(cfg.Security.PrivateKey) != ed25519.PrivateKeySize {
		t.Errorf("private key size = %d, want %d", len(cfg.Security.PrivateKey), ed25519.PrivateKeySize)
	}

	if mgr == nil {
		t.Fatal("NewKeyManager() returned nil manager")
	}
}

func TestNewManagerInitializesCryptoKeys(t *testing.T) {
	crypto.ResetKeys()
	t.Cleanup(crypto.ResetKeys)

	_, err := NewKeyManager(filepath.Join(t.TempDir(), "config.yaml"))
	require.NoError(t, err)
	require.True(t, crypto.KeysSet())

	encrypted, err := crypto.Encrypt("local-secret")
	require.NoError(t, err)
	require.NotEmpty(t, encrypted)
}

func TestNewManager_LoadsExistingKeyPair(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")

	mgr1, err := NewKeyManager(path)
	if err != nil {
		t.Fatalf("first NewKeyManager() error = %v", err)
	}

	mgr2, err := NewKeyManager(path)
	if err != nil {
		t.Fatalf("second NewKeyManager() error = %v", err)
	}

	key, err := mgr1.Generate()
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	if err := mgr2.Validate(key); err != nil {
		t.Errorf("Validate() error = %v", err)
	}
}

func TestFingerprintTokenAnalysisPayloadIsStableKeyedAndBodyFree(t *testing.T) {
	firstPath := filepath.Join(t.TempDir(), "first.yaml")
	first, err := NewKeyManager(firstPath)
	require.NoError(t, err)
	firstAgain, err := NewKeyManager(firstPath)
	require.NoError(t, err)
	second, err := NewKeyManager(filepath.Join(t.TempDir(), "second.yaml"))
	require.NoError(t, err)

	payload := []byte("private task body")
	fingerprint := first.FingerprintTokenAnalysisPayload(payload)

	require.Len(t, fingerprint, 32)
	require.Equal(t, fingerprint, firstAgain.FingerprintTokenAnalysisPayload(payload))
	require.NotEqual(t, fingerprint, first.FingerprintTokenAnalysisPayload([]byte("different task body")))
	require.NotEqual(t, fingerprint, second.FingerprintTokenAnalysisPayload(payload))
	require.NotContains(t, fingerprint, "private")
}

func TestGenerate_ReusesPersistedKeyPairWhenAnotherProcessAlreadyCreatedIt(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")

	winner, err := NewKeyManager(path)
	require.NoError(t, err)

	// Simulate a stale process that already decided generation was needed
	// before another process finished writing the shared keypair to disk.
	stale := &KeyManager{}
	require.NoError(t, stale.generate(path))

	key, err := winner.GenerateTempKey()
	require.NoError(t, err)
	require.NoError(t, stale.ValidateTempKey(key))
}

func TestGenerate_ValidKey(t *testing.T) {
	mgr := newTestManager(t)

	key, err := mgr.Generate()
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	if key == "" {
		t.Fatal("Generate() returned empty key")
	}

	if len(key) < len(Prefix) || key[:len(Prefix)] != Prefix {
		t.Errorf("key prefix = %q, want %q", key[:min(len(Prefix), len(key))], Prefix)
	}

	parts := splitParts(key)
	if len(parts) != 3 {
		t.Fatalf("key has %d parts, want 3: %v", len(parts), parts)
	}

	if parts[0] != Prefix {
		t.Errorf("key prefix = %q, want %q", parts[0], Prefix)
	}

	if _, err := base64.RawURLEncoding.DecodeString(parts[1]); err != nil {
		t.Errorf("payload is not valid base64url: %v", err)
	}

	if _, err := base64.RawURLEncoding.DecodeString(parts[2]); err != nil {
		t.Errorf("signature is not valid base64url: %v", err)
	}
}

func TestValidate_ValidKey(t *testing.T) {
	mgr := newTestManager(t)

	key, err := mgr.Generate()
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	if err := mgr.Validate(key); err != nil {
		t.Errorf("Validate() error = %v", err)
	}
}

func TestValidate_CrossManager(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")

	mgr1, err := NewKeyManager(path)
	if err != nil {
		t.Fatalf("NewKeyManager() error = %v", err)
	}

	key, err := mgr1.Generate()
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	mgr2, err := NewKeyManager(path)
	if err != nil {
		t.Fatalf("second NewKeyManager() error = %v", err)
	}

	if err := mgr2.Validate(key); err != nil {
		t.Errorf("cross-manager Validate() error = %v", err)
	}
}

func TestValidate_DifferentKeyPair(t *testing.T) {
	mgr1 := newTestManager(t)
	mgr2 := newTestManager(t)

	key, err := mgr1.Generate()
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	if err := mgr2.Validate(key); err != ErrInvalidSignature {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidSignature)
	}
}

func TestValidate_EmptyKey(t *testing.T) {
	mgr := newTestManager(t)

	if err := mgr.Validate(""); err != ErrInvalidFormat {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidFormat)
	}
}

func TestValidate_MissingPrefix(t *testing.T) {
	mgr := newTestManager(t)

	if err := mgr.Validate("abc.123.456"); err != ErrInvalidFormat {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidFormat)
	}
}

func TestValidate_NoDots(t *testing.T) {
	mgr := newTestManager(t)

	if err := mgr.Validate(Prefix + "nodots"); err != ErrInvalidFormat {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidFormat)
	}
}

func TestValidate_OnlyPrefix(t *testing.T) {
	mgr := newTestManager(t)

	if err := mgr.Validate(Prefix); err != ErrInvalidFormat {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidFormat)
	}
}

func TestValidate_InvalidBase64(t *testing.T) {
	mgr := newTestManager(t)

	if err := mgr.Validate(Prefix + ".!!invalid!!.abc"); !errors.Is(err, ErrInvalidFormat) {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidFormat)
	}
}

func TestValidate_TamperedPayload(t *testing.T) {
	mgr := newTestManager(t)

	key, err := mgr.Generate()
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	parts := splitParts(key)

	tampered := parts[0] + "." + base64.RawURLEncoding.EncodeToString([]byte("9999999999999999")) + "." + parts[2]

	if err := mgr.Validate(tampered); err != ErrInvalidSignature {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidSignature)
	}
}

func TestValidate_TamperedSignature(t *testing.T) {
	mgr := newTestManager(t)

	key, err := mgr.Generate()
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	parts := splitParts(key)

	tampered := parts[0] + "." + parts[1] + "." + base64.RawURLEncoding.EncodeToString([]byte("tampered"))

	if err := mgr.Validate(tampered); err != ErrInvalidSignature {
		t.Errorf("Validate() error = %v, want %v", err, ErrInvalidSignature)
	}
}

func TestValidate_ExpiredKey(t *testing.T) {
	mgr := newTestManager(t)

	expiredTs := time.Now().Add(-8 * 24 * time.Hour).UnixNano()
	payload := base64.RawURLEncoding.EncodeToString([]byte(intStr(expiredTs)))

	mgr.mu.RLock()
	priKey := mgr.priKey
	mgr.mu.RUnlock()

	sig := ed25519.Sign(priKey, []byte(intStr(expiredTs)))
	sigB64 := base64.RawURLEncoding.EncodeToString(sig)

	key := Prefix + "." + payload + "." + sigB64

	if err := mgr.Validate(key); err != ErrExpired {
		t.Errorf("Validate() error = %v, want %v", err, ErrExpired)
	}
}

func TestValidate_JustBeforeExpiry(t *testing.T) {
	mgr := newTestManager(t)

	ts := time.Now().Add(-(6*24*time.Hour + 23*time.Hour)).UnixNano()
	payload := base64.RawURLEncoding.EncodeToString([]byte(intStr(ts)))

	mgr.mu.RLock()
	priKey := mgr.priKey
	mgr.mu.RUnlock()

	sig := ed25519.Sign(priKey, []byte(intStr(ts)))
	sigB64 := base64.RawURLEncoding.EncodeToString(sig)

	key := Prefix + "." + payload + "." + sigB64

	if err := mgr.Validate(key); err != nil {
		t.Errorf("Validate() error = %v, want nil (key should not be expired yet)", err)
	}
}

func TestParseKey(t *testing.T) {
	tests := []struct {
		name      string
		key       string
		wantErr   error
		wantParts int
	}{
		{"empty", "", ErrInvalidFormat, 0},
		{"too short", "sk", ErrInvalidFormat, 0},
		{"no dots", "sk-devabc", ErrInvalidFormat, 0},
		{"one dot", "sk-dev.abc", ErrInvalidFormat, 0},
		{"invalid payload b64", "sk-dev.!!.abc", ErrInvalidFormat, 0},
		{"invalid sig b64", "sk-dev.MQ.!!", ErrInvalidFormat, 0},
		{"valid", "sk-dev.MQ.yMGB9yRdFwKxZ7a_4NHNqGJrG2JVAjE", nil, 2},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			payload, sig, err := parseKey(tt.key)
			if tt.wantErr != nil && !errors.Is(err, tt.wantErr) {
				t.Errorf("parseKey() error = %v, want %v", err, tt.wantErr)
				return
			}
			if err == nil && tt.wantParts == 0 {
				t.Error("expected error but got nil")
			}
			if err == nil {
				if len(payload) == 0 {
					t.Error("payload is empty")
				}
				if len(sig) == 0 {
					t.Error("sig is empty")
				}
			}
		})
	}
}

func TestParseTimestamp(t *testing.T) {
	tests := []struct {
		payload string
		want    int64
		wantErr bool
	}{
		{"0", 0, false},
		{"1740000000000000000", 1740000000000000000, false},
		{"abc", 0, true},
		{"", 0, true},
	}

	for _, tt := range tests {
		ts, err := parseTimestamp([]byte(tt.payload))
		if tt.wantErr && err == nil {
			t.Errorf("parseTimestamp(%q) expected error, got nil", tt.payload)
		}
		if !tt.wantErr && err != nil {
			t.Errorf("parseTimestamp(%q) unexpected error: %v", tt.payload, err)
		}
		if !tt.wantErr && ts != tt.want {
			t.Errorf("parseTimestamp(%q) = %d, want %d", tt.payload, ts, tt.want)
		}
	}
}

func TestGenerate_UninitializedManager(t *testing.T) {
	mgr := &KeyManager{}

	_, err := mgr.Generate()
	if err == nil {
		t.Error("Generate() expected error for uninitialized manager")
	}
}

func TestValidate_UninitializedManager(t *testing.T) {
	mgr := &KeyManager{}

	err := mgr.Validate("some.key")
	if err == nil {
		t.Error("Validate() expected error for uninitialized manager")
	}
}

func newTestManager(t *testing.T) *KeyManager {
	t.Helper()

	dir := t.TempDir()
	path := filepath.Join(dir, "config.yaml")

	mgr, err := NewKeyManager(path)
	if err != nil {
		t.Fatalf("newTestManager: %v", err)
	}
	return mgr
}

func splitParts(key string) []string {
	var parts []string
	start := 0
	for i, c := range key {
		if c == '.' {
			parts = append(parts, key[start:i])
			start = i + 1
		}
	}
	parts = append(parts, key[start:])
	return parts
}

func intStr(n int64) string {
	return fmt.Sprintf("%d", n)
}