package main

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"io"
	"math/big"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"git.sechmachine.io.vn/sechmachine/VerseVDI-Data-Plane/gateway"
	protocol "git.sechmachine.io.vn/sechmachine/VerseVDI-Protocol/gen/go/protocol"
)

func TestHeartbeatReportsMeasuredEgressInsteadOfConfiguredCapacity(t *testing.T) {
	heartbeats := make(chan protocol.GatewayHeartbeat, 1)
	control := httptest.NewTLSServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/api/v1/gateway/heartbeat" {
			heartbeat, err := protocol.DecodeGatewayHeartbeat(mustReadBody(t, request))
			if err != nil {
				t.Errorf("decode heartbeat: %v", err)
			} else {
				heartbeats <- heartbeat
			}
		}
		response.WriteHeader(http.StatusNoContent)
	}))
	defer control.Close()

	server, err := gateway.NewServer(gateway.ServerConfig{
		ListenAddress: "127.0.0.1:0", TLSConfig: heartbeatTestTLS(t), GatewayID: "gateway-1",
		Admission: gateway.AdmissionFunc(func(context.Context, protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error) {
			return protocol.SessionAuthority{}, gateway.ErrAdmissionRejected
		}),
		Provider: gateway.NewFakeApollo(gateway.FakeApolloConfig{}),
	})
	if err != nil {
		t.Fatal(err)
	}
	defer server.Close()
	registration := protocol.GatewayRegistration{GatewayID: "gateway-1", BandwidthCapacityKbps: 100000}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go heartbeatLoop(ctx, gateway.NewControlPlaneClient(control.URL, control.Client()), server, registration)

	select {
	case heartbeat := <-heartbeats:
		cancel()
		if heartbeat.EgressKbps != 0 {
			t.Fatalf("idle measured egress = %d Kbps, want 0; configured capacity is not traffic", heartbeat.EgressKbps)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("heartbeat was not published")
	}
}

func TestHeartbeatSamplerUsesByteDeltaAndMonotonicElapsed(t *testing.T) {
	var sampler heartbeatSampler
	start := time.Now()
	if egress, _ := sampler.sample(start, gateway.MetricsSnapshot{MediaBytes: 1000}); egress != 0 {
		t.Fatalf("first sample egress = %d, want baseline 0", egress)
	}
	egress, telemetry := sampler.sample(start.Add(2*time.Second), gateway.MetricsSnapshot{
		AdmittedSessions: 2, AdmissionRejects: 3, Reconnects: 4, DrainTransitions: 5,
		MediaDrops: 6, MediaPackets: 7, MediaBytes: 17000, QueueDelayNanos: 9000,
		ProcessingDelayNanos: 10000, ProcessingSamples: 11, PacingDelayNanos: 12000,
		ProviderErrors: 13, InputRejected: 14, ControlRTTNanos: 15000,
		ControlJitterNanos: 16000, ControlLossPPM: 17, PendingReliable: 18, ProviderState: 2,
	})
	if egress != 64 || telemetry.MediaBytes != 17000 || telemetry.MediaPackets != 7 ||
		telemetry.QueueDelayMicros != 9 || telemetry.ProviderState != gateway.ProviderStateReady {
		t.Fatalf("sample = egress:%d telemetry:%#v", egress, telemetry)
	}
}

func mustReadBody(t *testing.T, request *http.Request) []byte {
	t.Helper()
	defer request.Body.Close()
	data, err := io.ReadAll(request.Body)
	if err != nil {
		t.Fatal(err)
	}
	return data
}

func heartbeatTestTLS(t *testing.T) *tls.Config {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	template := &x509.Certificate{
		SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "gateway.test"},
		NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour),
		IsCA: true, BasicConstraintsValid: true, KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	certificate := tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}
	pool := x509.NewCertPool()
	pool.AddCert(template)
	return &tls.Config{
		MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{certificate},
		ClientAuth: tls.RequireAndVerifyClientCert, ClientCAs: pool,
	}
}