package service

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/klarkxy/nekonest-cloud/relay/internal/authsnapshot"
	"github.com/klarkxy/nekonest-cloud/relay/internal/controlplane"
	"github.com/klarkxy/nekonest-cloud/relay/internal/forwarder"
	"github.com/klarkxy/nekonest-cloud/relay/internal/registry"
	"github.com/klarkxy/nekonest/relaycore/protocol"
)

type fakeControlPlane struct {
	snapshot controlplane.SnapshotResponse
	register func(string, string, json.RawMessage) (protocol.DeviceRegistrationResponse, error)
	route    *controlplane.RouteResolution
	consume  func(context.Context, any) (controlplane.ConsumedHandoff, error)
	complete func(context.Context, string, string, string, string) error
}

func localRoute() controlplane.RouteResolution {
	return controlplane.RouteResolution{
		RelayNodeID: "node_test", PlacementGeneration: 1,
		HomeRegion: "test", Local: true,
	}
}

func (f *fakeControlPlane) ResolveDeviceRoute(context.Context, string, string) (controlplane.RouteResolution, error) {
	if f.route != nil {
		return *f.route, nil
	}
	return localRoute(), nil
}
func (f *fakeControlPlane) ResolvePhoneRoute(context.Context, string, string) (controlplane.RouteResolution, error) {
	if f.route != nil {
		return *f.route, nil
	}
	return localRoute(), nil
}
func (f *fakeControlPlane) ResolveTenantRoute(context.Context, string, int64) (controlplane.RouteResolution, error) {
	if f.route != nil {
		return *f.route, nil
	}
	return localRoute(), nil
}
func (f *fakeControlPlane) ResolveHandoffRoute(context.Context, string, string) (controlplane.RouteResolution, error) {
	if f.route != nil {
		return *f.route, nil
	}
	return localRoute(), nil
}

func (f *fakeControlPlane) AuthorizeDevice(_ context.Context, _ string, _ string) (controlplane.SnapshotResponse, error) {
	return f.snapshot, nil
}
func (f *fakeControlPlane) AuthorizationSnapshot(_ context.Context, _ string, _ int64) (controlplane.SnapshotResponse, error) {
	return f.snapshot, nil
}
func (f *fakeControlPlane) AuthorizationDelta(_ context.Context, _ string, _ int64) (controlplane.DeltaResponse, error) {
	return controlplane.DeltaResponse{}, nil
}
func (f *fakeControlPlane) AuthorizePhone(context.Context, string, string) (controlplane.PhoneAuthorization, error) {
	return controlplane.PhoneAuthorization{}, nil
}
func (f *fakeControlPlane) ConsumePhoneHandoff(ctx context.Context, request any) (controlplane.ConsumedHandoff, error) {
	if f.consume != nil {
		return f.consume(ctx, request)
	}
	return controlplane.ConsumedHandoff{}, nil
}
func (f *fakeControlPlane) CompletePhoneHandoff(ctx context.Context, handoffID, phoneID, phoneTokenHash, routeHandleHash string) error {
	if f.complete != nil {
		return f.complete(ctx, handoffID, phoneID, phoneTokenHash, routeHandleHash)
	}
	return nil
}
func (f *fakeControlPlane) RevokePhone(context.Context, string, string, string) error { return nil }
func (f *fakeControlPlane) RegisterDevice(_ context.Context, bootstrap, source string, request json.RawMessage) (protocol.DeviceRegistrationResponse, error) {
	return f.register(bootstrap, source, request)
}

func signedFixture(t *testing.T, token string) (controlplane.SnapshotResponse, authsnapshot.Keyring) {
	t.Helper()
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now().UTC()
	keyString := base64.RawURLEncoding.EncodeToString(make([]byte, 32))
	payload := authsnapshot.Payload{
		SnapshotVersion:       1,
		TenantID:              "tenant_0123456789abcdef0123456789abcdef",
		TenantStatus:          "active",
		HomeRegion:            "test",
		RelayNodeID:           "node_test",
		PlacementGeneration:   1,
		AuthorizationRevision: 1,
		Devices: []authsnapshot.Device{{
			DeviceID:            "host_0123456789abcdef0123456789abcdef",
			CredentialHash:      controlplane.SHA256Hex(token),
			IdentityFingerprint: strings.Repeat("a", 64),
			Ed25519Public:       keyString,
			X25519Public:        keyString,
			Name:                "Test host",
			OS:                  "windows",
		}},
		IssuedAt:  now.Add(-time.Second).Format(time.RFC3339Nano),
		ExpiresAt: now.Add(4 * time.Minute).Format(time.RFC3339Nano),
	}
	message, err := authsnapshot.SigningBytes(payload)
	if err != nil {
		t.Fatal(err)
	}
	signed := authsnapshot.Signed{
		Algorithm: "Ed25519",
		KID:       "test-key",
		Payload:   payload,
		Signature: base64.RawURLEncoding.EncodeToString(ed25519.Sign(privateKey, message)),
	}
	jwk, _ := json.Marshal(map[string]string{
		"kty": "OKP", "crv": "Ed25519", "kid": "test-key",
		"x": base64.RawURLEncoding.EncodeToString(publicKey),
	})
	return controlplane.SnapshotResponse{Snapshot: signed, PublicKeyJWK: jwk}, authsnapshot.Keyring{
		"test-key": {PublicKey: publicKey, NotBefore: now.Add(-time.Hour), RetainUntil: now.Add(time.Hour)},
	}
}

func newTestService(t *testing.T, control *fakeControlPlane, keyring authsnapshot.Keyring) *Service {
	return newTestServiceWithNode(t, control, keyring, "node_test", nil)
}

func newTestServiceWithNode(
	t *testing.T, control *fakeControlPlane, keyring authsnapshot.Keyring,
	nodeID string, relayForwarder *forwarder.Forwarder,
) *Service {
	t.Helper()
	registryValue, err := registry.New(registry.Config{
		DataRoot:        t.TempDir(),
		NodeID:          nodeID,
		AllowedOrigins:  []string{"https://pwa.example"},
		RouteSecret:     []byte("0123456789abcdef0123456789abcdef"),
		Keyring:         keyring,
		ControlPlane:    control,
		RefreshInterval: time.Hour,
		DeltaInterval:   time.Hour,
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = registryValue.Close() })
	service, err := New(Config{
		ControlPlane:      control,
		Registry:          registryValue,
		AllowedPWAOrigins: []string{"https://pwa.example"},
		SourceHashSecret:  []byte("abcdef0123456789abcdef0123456789"),
		NodeID:            nodeID,
		ForwardSecret:     []byte("fedcba9876543210fedcba9876543210"),
		HandoffSecret:     []byte("handoff0123456789handoff0123456789"),
		Forwarder:         relayForwarder,
	})
	if err != nil {
		t.Fatal(err)
	}
	return service
}

func TestRegistrationHashesTrustedSourceAndReturnsStableContract(t *testing.T) {
	snapshot, keys := signedFixture(t, "device-token")
	control := &fakeControlPlane{snapshot: snapshot}
	control.register = func(bootstrap, source string, body json.RawMessage) (protocol.DeviceRegistrationResponse, error) {
		if bootstrap != "pair_secret" || len(source) != 64 {
			t.Fatalf("bootstrap=%q source=%q", bootstrap, source)
		}
		if strings.Contains(string(body), "pair_secret") {
			t.Fatal("relay injected bootstrap into the client body before the trusted control call")
		}
		return protocol.DeviceRegistrationResponse{
			DeviceID:        "host_0123456789abcdef0123456789abcdef",
			Token:           "device-token",
			Name:            "Host",
			TransportMode:   protocol.TransportSealed,
			ConnectionState: protocol.ConnectionReady,
		}, nil
	}
	service := newTestService(t, control, keys)
	request := httptest.NewRequest(http.MethodPost, "https://connect.example/api/devices/register", strings.NewReader(`{"os":"windows"}`))
	request.RemoteAddr = "203.0.113.5:12345"
	request.Header.Set("X-Neko-Bootstrap", "pair_secret")
	response := httptest.NewRecorder()
	service.Handler().ServeHTTP(response, request)
	if response.Code != http.StatusOK || !strings.Contains(response.Body.String(), `"connection_state":"ready"`) {
		t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
	}
}

func TestSourceAddressOnlyTrustsForwardedIPFromConfiguredProxy(t *testing.T) {
	request := httptest.NewRequest(http.MethodPost, "https://connect.example/api/devices/register", nil)
	request.RemoteAddr = "192.0.2.10:443"
	request.Header.Set("CF-Connecting-IP", "203.0.113.9")
	if got := sourceAddress(request, nil); got != "192.0.2.10" {
		t.Fatalf("untrusted source = %q", got)
	}
	prefix := netip.MustParsePrefix("192.0.2.0/24")
	if got := sourceAddress(request, []netip.Prefix{prefix}); got != "203.0.113.9" {
		t.Fatalf("trusted source = %q", got)
	}
}

func TestPhoneHandoffRejectsTrailingJSON(t *testing.T) {
	snapshot, keys := signedFixture(t, "device-token")
	control := &fakeControlPlane{snapshot: snapshot, register: func(string, string, json.RawMessage) (protocol.DeviceRegistrationResponse, error) {
		return protocol.DeviceRegistrationResponse{}, nil
	}}
	service := newTestService(t, control, keys)
	request := httptest.NewRequest(http.MethodPost, "https://connect.example/api/pwa/handoff/exchange", strings.NewReader(
		`{"ticket":"ticket","pwa_origin":"https://pwa.example"} {}`,
	))
	request.Header.Set("Origin", "https://pwa.example")
	response := httptest.NewRecorder()
	service.Handler().ServeHTTP(response, request)
	if response.Code != http.StatusBadRequest {
		t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
	}
}

func TestPhoneHandoffRetriesAmbiguousCompletionWithStableCredentials(t *testing.T) {
	snapshot, keys := signedFixture(t, "device-token")
	key := base64.RawURLEncoding.EncodeToString(make([]byte, 32))
	control := &fakeControlPlane{snapshot: snapshot, register: func(string, string, json.RawMessage) (protocol.DeviceRegistrationResponse, error) {
		return protocol.DeviceRegistrationResponse{}, nil
	}}
	control.consume = func(context.Context, any) (controlplane.ConsumedHandoff, error) {
		return controlplane.ConsumedHandoff{
			HandoffID:           "handoff_0123456789abcdef0123456789abcdef",
			TenantID:            snapshot.Snapshot.Payload.TenantID,
			Name:                "Cloud PWA",
			PhoneEd25519Public:  key,
			PhoneX25519Public:   key,
			IdentityFingerprint: strings.Repeat("f", 64),
			PlacementGeneration: snapshot.Snapshot.Payload.PlacementGeneration,
		}, nil
	}
	type completion struct{ handoff, phone, tokenHash, routeHash string }
	var completions []completion
	control.complete = func(_ context.Context, handoff, phone, tokenHash, routeHash string) error {
		completions = append(completions, completion{handoff, phone, tokenHash, routeHash})
		if len(completions) == 1 {
			return errors.New("control-plane response lost after commit")
		}
		return nil
	}
	service := newTestService(t, control, keys)
	body := `{"ticket":"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef","pwa_origin":"https://pwa.example","name":"Cloud PWA","phone_ed25519_public":"` + key + `","phone_x25519_public":"` + key + `","identity_fingerprint":"` + strings.Repeat("f", 64) + `"}`
	exchange := func() map[string]string {
		request := httptest.NewRequest(http.MethodPost, "https://connect.example/api/pwa/handoff/exchange", strings.NewReader(body))
		request.Header.Set("Origin", "https://pwa.example")
		response := httptest.NewRecorder()
		service.Handler().ServeHTTP(response, request)
		if response.Code != http.StatusOK {
			t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
		}
		var result map[string]string
		if err := json.Unmarshal(response.Body.Bytes(), &result); err != nil {
			t.Fatal(err)
		}
		return result
	}
	first := exchange()
	second := exchange()
	if len(completions) != 3 || completions[0] != completions[1] || completions[1] != completions[2] {
		t.Fatalf("completion credentials changed across retry: %#v", completions)
	}
	if first["phone_id"] != second["phone_id"] || first["phone_token"] != second["phone_token"] || first["route_handle"] != second["route_handle"] {
		t.Fatalf("handoff response changed across retry: first=%#v second=%#v", first, second)
	}
	if _, cached := service.registry.ByPhoneRoute(first["route_handle"]); cached {
		t.Fatal("pending handoff route became usable before first possession proof")
	}
}

func TestSecondLiveDaemonWithSameIdentityIsRejected(t *testing.T) {
	const token = "device-token"
	snapshot, keys := signedFixture(t, token)
	control := &fakeControlPlane{snapshot: snapshot, register: func(string, string, json.RawMessage) (protocol.DeviceRegistrationResponse, error) {
		return protocol.DeviceRegistrationResponse{}, nil
	}}
	service := newTestService(t, control, keys)
	server := httptest.NewServer(service.Handler())
	defer server.Close()

	dial := func() *websocket.Conn {
		connection, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http")+"/ws/daemon", nil)
		if err != nil {
			t.Fatal(err)
		}
		if err := connection.WriteJSON(protocol.NekoMessage{
			ProtocolVersion: "1.3",
			TransportMode:   protocol.TransportSealed,
			Type:            protocol.MsgRegisterDevice,
			Timestamp:       time.Now().Unix(),
			Payload: map[string]any{
				"device_id":      "host_0123456789abcdef0123456789abcdef",
				"token":          token,
				"daemon_version": "0.2.6",
			},
		}); err != nil {
			t.Fatal(err)
		}
		return connection
	}
	first := dial()
	defer first.Close()
	var authenticated protocol.NekoMessage
	if err := first.ReadJSON(&authenticated); err != nil || authenticated.Type != protocol.MsgAuthResponse {
		t.Fatalf("first authentication err=%v message=%+v", err, authenticated)
	}
	second := dial()
	defer second.Close()
	var rejected protocol.NekoMessage
	if err := second.ReadJSON(&rejected); err != nil {
		t.Fatal(err)
	}
	if rejected.Type != protocol.MsgError || rejected.Payload["error_code"] != string(protocol.ServiceErrorDeviceAlreadyConnected) {
		t.Fatalf("rejected = %+v", rejected)
	}
}

func TestRemoteIngressTunnelsDaemonWithoutRedirectingClient(t *testing.T) {
	const token = "device-token"
	snapshot, keys := signedFixture(t, token)
	targetControl := &fakeControlPlane{snapshot: snapshot, register: func(string, string, json.RawMessage) (protocol.DeviceRegistrationResponse, error) {
		return protocol.DeviceRegistrationResponse{}, nil
	}}
	targetService := newTestServiceWithNode(t, targetControl, keys, "node_test", nil)
	targetServer := httptest.NewServer(targetService.Handler())
	defer targetServer.Close()

	relayForwarder, err := forwarder.New(forwarder.Config{
		NodeID: "node_source", Secret: []byte("fedcba9876543210fedcba9876543210"),
		Endpoints:  map[string]string{"target-ref": targetServer.URL},
		HTTPClient: targetServer.Client(), WebSocketDialer: websocket.DefaultDialer,
	})
	if err != nil {
		t.Fatal(err)
	}
	remoteRoute := controlplane.RouteResolution{
		RelayNodeID: "node_test", PlacementGeneration: 1, HomeRegion: "test",
		Local: false, EndpointRef: "target-ref",
	}
	sourceControl := &fakeControlPlane{snapshot: snapshot, route: &remoteRoute, register: targetControl.register}
	sourceService := newTestServiceWithNode(t, sourceControl, keys, "node_source", relayForwarder)
	sourceServer := httptest.NewServer(sourceService.Handler())
	defer sourceServer.Close()

	connection, response, err := websocket.DefaultDialer.Dial(
		"ws"+strings.TrimPrefix(sourceServer.URL, "http")+"/ws/daemon", nil,
	)
	if response != nil && response.Body != nil {
		_ = response.Body.Close()
	}
	if err != nil {
		t.Fatal(err)
	}
	defer connection.Close()
	if err := connection.WriteJSON(protocol.NekoMessage{
		ProtocolVersion: "1.3", TransportMode: protocol.TransportSealed,
		Type: protocol.MsgRegisterDevice, Timestamp: time.Now().Unix(),
		Payload: map[string]any{
			"device_id": "host_0123456789abcdef0123456789abcdef",
			"token":     token, "daemon_version": "0.2.6",
		},
	}); err != nil {
		t.Fatal(err)
	}
	var authenticated protocol.NekoMessage
	if err := connection.ReadJSON(&authenticated); err != nil || authenticated.Type != protocol.MsgAuthResponse {
		t.Fatalf("forwarded authentication err=%v message=%+v", err, authenticated)
	}
}