package securityevents

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	"github.com/google/uuid"
)

type memoryConnectionRepository struct {
	mutex      sync.Mutex
	connection *store.SecurityEventConnection
	mode       string
	verifiedAt *time.Time
	evaluation store.SecurityEventEvaluation
}

func (*memoryConnectionRepository) ApplySessionRevoked(context.Context, store.ApplySessionRevokedInput) (store.ApplySecurityEventResult, error) {
	return store.ApplySecurityEventResult{}, nil
}
func (*memoryConnectionRepository) ApplySecurityEventStreamUpdated(context.Context, string, string, string, string, string, string, time.Time) (bool, error) {
	return true, nil
}
func (*memoryConnectionRepository) ConfirmSecurityEventVerification(context.Context, string, string, string, string, []byte, time.Time) (bool, error) {
	return true, nil
}
func (*memoryConnectionRepository) RecordSecurityEventReceipt(context.Context, string, string, string, string, string) (bool, error) {
	return true, nil
}
func (r *memoryConnectionRepository) EnsureSecurityEventStreamState(context.Context, string, string, string) error {
	r.mode = "bootstrap"
	return nil
}
func (*memoryConnectionRepository) BeginSecurityEventVerification(context.Context, string, string, []byte, time.Time) error {
	return nil
}
func (*memoryConnectionRepository) RecordSecurityEventHeartbeatFailure(context.Context, string, string, string) error {
	return nil
}
func (*memoryConnectionRepository) AdvanceSecurityEventStreamState(context.Context, string, string, time.Time, time.Duration) error {
	return nil
}
func (r *memoryConnectionRepository) EvaluateOIDCSecurityEvent(context.Context, string, string, string, string, string, time.Time, time.Time, time.Duration) (store.SecurityEventEvaluation, error) {
	if r.evaluation.Mode == "" {
		return store.SecurityEventEvaluation{Mode: "bootstrap", RequireIntrospection: true}, nil
	}
	return r.evaluation, nil
}
func (r *memoryConnectionRepository) CreateSecurityEventConnection(_ context.Context, input store.CreateSecurityEventConnectionInput) (store.SecurityEventConnection, error) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	if r.connection != nil {
		return store.SecurityEventConnection{}, store.ErrSecurityEventConnectionConflict
	}
	now := time.Now().UTC()
	value := store.SecurityEventConnection{ConnectionID: input.ConnectionID, TransmitterIssuer: input.TransmitterIssuer,
		EndpointURL: input.EndpointURL, CredentialRef: input.CredentialRef, ManagementClientID: input.ManagementClientID,
		ManagementCredentialRef: &input.ManagementCredentialRef, LifecycleStatus: "connecting", Version: 1,
		IdempotencyKey: input.IdempotencyKey, CreatedAt: now, UpdatedAt: now, HealthMode: "disabled"}
	r.connection = &value
	return value, nil
}
func (r *memoryConnectionRepository) SetSecurityEventManagementCredential(_ context.Context, id, clientID, reference string) (store.SecurityEventConnection, error) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	if r.connection == nil || r.connection.ConnectionID != id {
		return store.SecurityEventConnection{}, store.ErrSecurityEventConnectionNotFound
	}
	r.connection.ManagementClientID, r.connection.ManagementCredentialRef = clientID, &reference
	return *r.connection, nil
}
func (r *memoryConnectionRepository) SecurityEventConnection(context.Context) (store.SecurityEventConnection, error) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	if r.connection == nil {
		return store.SecurityEventConnection{}, store.ErrSecurityEventConnectionNotFound
	}
	value := *r.connection
	value.HealthMode, value.LastVerificationAt = r.mode, r.verifiedAt
	return value, nil
}
func (r *memoryConnectionRepository) BindSecurityEventConnection(_ context.Context, id, audience, streamID, lifecycle string, expected int64) (store.SecurityEventConnection, error) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	if r.connection == nil || r.connection.ConnectionID != id || r.connection.Version != expected {
		return store.SecurityEventConnection{}, store.ErrSecurityEventConnectionConflict
	}
	r.connection.Audience, r.connection.StreamID = &audience, &streamID
	r.connection.LifecycleStatus, r.connection.Version = lifecycle, r.connection.Version+1
	return *r.connection, nil
}
func (r *memoryConnectionRepository) UpdateSecurityEventConnectionLifecycle(_ context.Context, id, lifecycle string, category *string) (store.SecurityEventConnection, error) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	if r.connection == nil || r.connection.ConnectionID != id {
		return store.SecurityEventConnection{}, store.ErrSecurityEventConnectionNotFound
	}
	r.connection.LifecycleStatus, r.connection.LastErrorCategory = lifecycle, category
	r.connection.Version++
	r.connection.UpdatedAt = time.Now().UTC()
	return *r.connection, nil
}
func (r *memoryConnectionRepository) SetSecurityEventNextCredential(context.Context, string, string) (store.SecurityEventConnection, error) {
	return store.SecurityEventConnection{}, errors.New("not used")
}

func TestConnectionManagerBootstrapOnlyBecomesEnabledWhenPushIsHealthy(t *testing.T) {
	for _, test := range []struct {
		mode, wantLifecycle string
	}{
		{mode: "push_healthy", wantLifecycle: "enabled"},
		{mode: "introspection_fallback", wantLifecycle: "degraded"},
	} {
		t.Run(test.mode, func(t *testing.T) {
			audience := "urn:easyai:ssf:receiver:" + uuid.NewString()
			repository := &memoryConnectionRepository{mode: test.mode, connection: &store.SecurityEventConnection{
				ConnectionID: uuid.NewString(), TransmitterIssuer: "https://auth.example/ssf", Audience: &audience,
				LifecycleStatus: "bootstrap", Version: 1, UpdatedAt: time.Now().UTC(),
			}}
			manager := &ConnectionManager{ctx: context.Background(), repository: repository,
				config: ConnectionManagerConfig{BootstrapDuration: time.Millisecond}}
			manager.finishBootstrap(repository.connection.ConnectionID)
			connection, err := repository.SecurityEventConnection(context.Background())
			if err != nil || connection.LifecycleStatus != test.wantLifecycle {
				t.Fatalf("connection=%#v err=%v", connection, err)
			}
		})
	}
}

func TestConnectionManagerResumesBootstrapAfterRestart(t *testing.T) {
	repository := &memoryConnectionRepository{mode: "push_healthy", connection: &store.SecurityEventConnection{
		ConnectionID: uuid.NewString(), TransmitterIssuer: "https://auth.example/ssf",
		LifecycleStatus: "bootstrap", Version: 1, UpdatedAt: time.Now().Add(-time.Second),
	}}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if _, err := NewConnectionManager(ctx, repository, &memorySecretStore{}, ConnectionManagerConfig{
		OIDCEnabled: true, OIDCIssuer: "https://auth.example/issuer", OIDCTenantID: uuid.NewString(),
		BootstrapDuration: 10 * time.Millisecond,
	}, &Metrics{}); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		connection, err := repository.SecurityEventConnection(ctx)
		if err == nil && connection.LifecycleStatus == "enabled" {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("bootstrap lifecycle did not resume after process restart")
}

func TestRemoteNotFoundIsRecognizedForIdempotentDisconnect(t *testing.T) {
	if !isRemoteHTTPStatus(remoteHTTPError{status: http.StatusNotFound}, http.StatusNotFound) {
		t.Fatal("remote 404 was not recognized as an already-deleted Stream")
	}
}

func TestConnectionLifecycleCanResumeOnlyIncompleteOrFailedConnections(t *testing.T) {
	for _, test := range []struct {
		status string
		want   bool
	}{
		{status: "connecting", want: true},
		{status: "verifying", want: true},
		{status: "degraded", want: true},
		{status: "error", want: true},
		{status: "bootstrap"},
		{status: "enabled"},
		{status: "rotating"},
		{status: "disconnect_pending"},
		{status: "retiring"},
	} {
		t.Run(test.status, func(t *testing.T) {
			if got := resumableConnectionLifecycle(test.status); got != test.want {
				t.Fatalf("resumableConnectionLifecycle(%q)=%t want %t", test.status, got, test.want)
			}
		})
	}
}

func TestExistingConnectionCredentialRepairDoesNotRequirePublicBaseURL(t *testing.T) {
	manager := &ConnectionManager{config: ConnectionManagerConfig{
		AppEnv: "test", OIDCEnabled: true, OIDCTenantID: uuid.NewString(),
	}}
	secret := []byte("machine-secret-for-existing-connection")
	defer clear(secret)
	if err := manager.validatePrerequisites("https://auth.example/ssf", "gateway-machine", secret, false); err != nil {
		t.Fatalf("existing connection credential repair unexpectedly required a public base URL: %v", err)
	}
	if err := manager.validatePrerequisites("https://auth.example/ssf", "gateway-machine", secret, true); err == nil {
		t.Fatal("a new connection must still require AI_GATEWAY_PUBLIC_BASE_URL")
	}
}

func TestTransmitterSSRFProtectionBlocksReservedNetworks(t *testing.T) {
	for _, address := range []string{
		"127.0.0.1", "10.0.0.1", "169.254.169.254", "100.64.0.1", "192.0.2.1", "198.51.100.1", "203.0.113.1", "2001:db8::1",
	} {
		if !blockedAddress(net.ParseIP(address)) {
			t.Fatalf("reserved address %s was allowed", address)
		}
	}
	if blockedAddress(net.ParseIP("8.8.8.8")) {
		t.Fatal("public address was blocked")
	}
}

func TestDialResolvedAddressesFallsBackFromIPv6ToIPv4(t *testing.T) {
	listener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer listener.Close()
	port := listener.Addr().(*net.TCPAddr).Port
	accepted := make(chan struct{})
	go func() {
		connection, acceptErr := listener.Accept()
		if acceptErr == nil {
			_ = connection.Close()
			close(accepted)
		}
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	connection, err := dialResolvedAddresses(ctx, "tcp", strconv.Itoa(port), []net.IPAddr{
		{IP: net.ParseIP("::1")},
		{IP: net.ParseIP("127.0.0.1")},
	}, &net.Dialer{Timeout: 250 * time.Millisecond})
	if err != nil {
		t.Fatalf("IPv4 fallback was not attempted: %v", err)
	}
	_ = connection.Close()
	select {
	case <-accepted:
	case <-ctx.Done():
		t.Fatal("IPv4 listener did not receive the fallback connection")
	}
}

func TestRetiringConnectionStillAppliesRevocationWatermark(t *testing.T) {
	audience := "urn:easyai:ssf:receiver:" + uuid.NewString()
	repository := &memoryConnectionRepository{
		evaluation: store.SecurityEventEvaluation{Mode: "introspection_fallback", Revoked: true, RequireIntrospection: true},
		connection: &store.SecurityEventConnection{
			ConnectionID: uuid.NewString(), TransmitterIssuer: "https://auth.example/ssf", Audience: &audience,
			LifecycleStatus: "retiring", Version: 1,
		},
	}
	manager := &ConnectionManager{ctx: context.Background(), repository: repository, config: ConnectionManagerConfig{StaleAfter: 3 * time.Minute}}
	evaluation, err := manager.Evaluate(context.Background(), auth.OIDCSecurityEventIdentity{IssuedAt: time.Now().Add(-time.Minute)})
	if err != nil || !evaluation.Enabled || !evaluation.Revoked || !evaluation.RequireIntrospection {
		t.Fatalf("retiring evaluation=%#v err=%v", evaluation, err)
	}
}
func (r *memoryConnectionRepository) PromoteSecurityEventCredential(context.Context, string, string) (store.SecurityEventConnection, error) {
	return store.SecurityEventConnection{}, errors.New("not used")
}
func (r *memoryConnectionRepository) DeleteSecurityEventConnection(context.Context, string) error {
	r.connection = nil
	return nil
}
func (r *memoryConnectionRepository) SecurityEventMetrics(context.Context, string, string) (string, *time.Time, error) {
	return r.mode, r.verifiedAt, nil
}

type memorySecretStore struct {
	values map[string][]byte
}

func (s *memorySecretStore) Put(_ context.Context, reference string, value []byte) error {
	if s.values == nil {
		s.values = map[string][]byte{}
	}
	s.values[reference] = append([]byte(nil), value...)
	return nil
}
func (s *memorySecretStore) Get(_ context.Context, reference string) ([]byte, error) {
	value, ok := s.values[reference]
	if !ok {
		return nil, ErrSecretNotFound
	}
	return append([]byte(nil), value...), nil
}
func (s *memorySecretStore) Delete(_ context.Context, reference string) error {
	delete(s.values, reference)
	return nil
}

func TestConnectionManagerCreatesPausedStreamWithBackendGeneratedBearer(t *testing.T) {
	repository := &memoryConnectionRepository{}
	secrets := &memorySecretStore{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var transmitter *httptest.Server
	verificationRequested := make(chan struct{}, 1)
	var managementScopeRequested atomic.Bool
	transmitter = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/.well-known/ssf-configuration/ssf":
			_ = json.NewEncoder(w).Encode(map[string]any{
				"spec_version": "1_0", "issuer": transmitter.URL + "/ssf", "jwks_uri": transmitter.URL + "/ssf/jwks.json",
				"configuration_endpoint": transmitter.URL + "/ssf/v1/stream", "status_endpoint": transmitter.URL + "/ssf/v1/status",
				"verification_endpoint": transmitter.URL + "/ssf/v1/verify", "delivery_methods_supported": []string{pushDeliveryMethod},
			})
		case "/oidc/token":
			clientID, secret, ok := r.BasicAuth()
			if !ok || clientID != "gateway-machine" || secret != "machine-secret-for-test" {
				t.Fatal("RFC 7662 machine client was not reused")
			}
			_ = r.ParseForm()
			if strings.Contains(r.Form.Get("scope"), "ssf.stream.manage") {
				managementScopeRequested.Store(true)
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"access_token": "machine-token", "token_type": "Bearer"})
		case "/ssf/v1/stream":
			if r.Method == http.MethodGet {
				_, _ = w.Write([]byte(`[]`))
				return
			}
			var request struct {
				Delivery struct {
					AuthorizationHeader string `json:"authorization_header"`
					EndpointURL         string `json:"endpoint_url"`
				} `json:"delivery"`
				Description string `json:"description"`
			}
			_ = json.NewDecoder(r.Body).Decode(&request)
			if len(strings.TrimPrefix(request.Delivery.AuthorizationHeader, "Bearer ")) < 43 {
				t.Fatal("generated Push Bearer has less than 256 bits")
			}
			_ = json.NewEncoder(w).Encode(map[string]any{
				"stream_id": uuid.NewString(), "iss": transmitter.URL + "/ssf",
				"aud": "urn:easyai:ssf:receiver:" + uuid.NewString(), "events_requested": []string{sessionRevokedEvent},
				"delivery":    map[string]any{"method": pushDeliveryMethod, "endpoint_url": request.Delivery.EndpointURL},
				"description": request.Description,
			})
		case "/ssf/v1/verify":
			w.WriteHeader(http.StatusNoContent)
			select {
			case verificationRequested <- struct{}{}:
			default:
			}
		default:
			http.NotFound(w, r)
		}
	}))
	defer transmitter.Close()
	manager, err := NewConnectionManager(ctx, repository, secrets, ConnectionManagerConfig{
		AppEnv: "test", OIDCEnabled: true, OIDCIssuer: transmitter.URL + "/oidc", OIDCTenantID: uuid.NewString(),
		ManagementClientID: "gateway-machine", ManagementClientSecret: "machine-secret-for-test",
		PublicBaseURL: transmitter.URL, HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour, HTTPClient: transmitter.Client(),
	}, &Metrics{})
	if err != nil {
		t.Fatal(err)
	}
	evaluation, err := manager.Evaluate(ctx, auth.OIDCSecurityEventIdentity{})
	if err != nil || evaluation.Enabled {
		t.Fatalf("unconfigured manager evaluation=%#v err=%v", evaluation, err)
	}
	view, err := manager.Connect(ctx, transmitter.URL+"/ssf", "", nil, "connect-once")
	if err != nil {
		t.Fatal(err)
	}
	if view.LifecycleStatus != "verifying" || view.StreamID == nil || view.Audience == nil {
		t.Fatalf("connection view=%#v", view)
	}
	encoded, _ := json.Marshal(view)
	if strings.Contains(string(encoded), "credential") || strings.Contains(string(encoded), "machine-secret-for-test") || strings.Contains(string(encoded), "ssf-push-") {
		t.Fatalf("secret metadata escaped in response: %s", encoded)
	}
	if len(secrets.values) != 2 {
		t.Fatalf("secret count=%d", len(secrets.values))
	}
	restarted := &ConnectionManager{repository: repository, secrets: secrets, config: ConnectionManagerConfig{}}
	restartedClientID, restartedSecret, err := restarted.IntrospectionCredential(ctx)
	if err != nil || restartedClientID != "gateway-machine" || string(restartedSecret) != "machine-secret-for-test" {
		t.Fatalf("persisted management credential was not recoverable after restart: client=%q secret_present=%v err=%v", restartedClientID, len(restartedSecret) > 0, err)
	}
	clear(restartedSecret)
	if !managementScopeRequested.Load() {
		t.Fatal("management scopes were not requested automatically")
	}
	select {
	case <-verificationRequested:
	case <-time.After(2 * time.Second):
		t.Fatal("verification was not started without a restart")
	}
}