package setup_test

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"log/slog"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/config"
	"github.com/dariuszjaniak/warp/server/internal/federation"
	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
	"github.com/dariuszjaniak/warp/server/internal/setup"
	"github.com/dariuszjaniak/warp/server/internal/userapi"
)

// The end-to-end federation tests: complete servers, each behind its own
// TLS listener, finding each other through a resolver that maps their
// names to those listeners. Everything between the client API on one
// server and the client API on the other is real — X-Matrix signing, key
// fetches, the routes, the gate, the outbound queue.

// fedNode is one server of an end-to-end federation.
type fedNode struct {
	name string
	srv  *setup.Server
	ts   *httptest.Server

	// handler is the server's handler, set once it exists. The listener
	// has to be up before the server is built — every server needs every
	// other's address — so it starts out answering 503.
	handler atomic.Pointer[http.Handler]
}

func (n *fedNode) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h := n.handler.Load()
	if h == nil {
		http.Error(w, "not ready", http.StatusServiceUnavailable)
		return
	}
	(*h).ServeHTTP(w, r)
}

// newFederation starts servers that can reach each other.
func newFederation(t *testing.T, names ...string) map[string]*fedNode {
	t.Helper()
	nodes := map[string]*fedNode{}
	for _, name := range names {
		node := &fedNode{name: name}
		node.ts = httptest.NewUnstartedServer(node)
		node.ts.StartTLS()
		t.Cleanup(node.ts.Close)
		nodes[name] = node
	}

	// Every test listener has the same certificate, for "example.com"; the
	// Host is that name with the listener's port, and the address is the
	// listener. That is the delegated-by-SRV shape: TLS is checked against
	// one name and the connection goes somewhere else.
	resolve := func(_ context.Context, name string) (federation.Resolved, error) {
		node, ok := nodes[name]
		if !ok {
			return federation.Resolved{}, &net.DNSError{Err: "no such host", Name: name, IsNotFound: true}
		}
		addr := node.ts.Listener.Addr().String()
		_, port, _ := net.SplitHostPort(addr)
		return federation.Resolved{Addresses: []string{addr}, Host: "example.com:" + port}, nil
	}
	anyTS := nodes[names[0]].ts
	clientTLS := anyTS.Client().Transport.(*http.Transport).TLSClientConfig.Clone()

	for _, name := range names {
		node := nodes[name]
		key, err := signing.Generate("")
		if err != nil {
			t.Fatal(err)
		}
		keyPath := filepath.Join(t.TempDir(), name+".key")
		if err := key.Save(keyPath); err != nil {
			t.Fatal(err)
		}
		cfg := config.Default()
		cfg.Global.ServerName = name
		cfg.Global.PrivateKeyPath = keyPath
		cfg.Database.ConnectionString = "file:" + filepath.ToSlash(filepath.Join(t.TempDir(), name+".db")) +
			"?_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)"
		cfg.Registration.Enabled = true
		cfg.ClientAPI.RateLimit.Enabled = false
		cfg.Federation.Enabled = true
		cfg.Media.BasePath = filepath.Join(t.TempDir(), name+"-media")
		if err := cfg.Validate(); err != nil {
			t.Fatalf("config: %v", err)
		}
		srv, err := setup.New(t.Context(), cfg, setup.Options{
			Logger:            slog.New(slog.DiscardHandler),
			FederationResolve: resolve,
			FederationTLS:     clientTLS,
		})
		if err != nil {
			t.Fatalf("setup.New(%s): %v", name, err)
		}
		t.Cleanup(func() { _ = srv.Close() })
		node.srv = srv
		h := srv.Handler()
		node.handler.Store(&h)
		t.Cleanup(srv.FederationModule().Stop)
		if err := srv.FederationModule().Start(t.Context()); err != nil {
			t.Fatalf("start federation on %s: %v", name, err)
		}
	}
	return nodes
}

// eventually polls until check passes or the deadline is reached.
func eventually(t *testing.T, what string, check func() bool) {
	t.Helper()
	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		if check() {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

// messageBodies reads the message bodies of a room's recent history.
func messageBodies(t *testing.T, srv *setup.Server, roomID, token string) []string {
	t.Helper()
	rec := getWithToken(t, srv, "/_matrix/client/v3/rooms/"+url.PathEscape(roomID)+"/messages?dir=b&limit=50", token)
	if rec.Code != http.StatusOK {
		return nil
	}
	var resp struct {
		Chunk []struct {
			Type    string `json:"type"`
			Content struct {
				Body string `json:"body"`
			} `json:"content"`
		} `json:"chunk"`
	}
	_ = json.Unmarshal(rec.Body.Bytes(), &resp)
	var out []string
	for _, e := range resp.Chunk {
		if e.Type == "m.room.message" {
			out = append(out, e.Content.Body)
		}
	}
	return out
}

func contains(list []string, s string) bool {
	for _, v := range list {
		if v == s {
			return true
		}
	}
	return false
}

func sendMessage(t *testing.T, srv *setup.Server, roomID, token, txn, body string) {
	t.Helper()
	rec := put(t, srv, "/_matrix/client/v3/rooms/"+url.PathEscape(roomID)+"/send/m.room.message/"+txn,
		map[string]any{"msgtype": "m.text", "body": body}, token)
	if rec.Code != http.StatusOK {
		t.Fatalf("send %q: status %d: %s", body, rec.Code, rec.Body.String())
	}
}

func TestFederationEndToEnd(t *testing.T) {
	if testing.Short() {
		t.Skip("starts two servers")
	}
	nodes := newFederation(t, "hs-a.example", "hs-b.example")
	a, b := nodes["hs-a.example"].srv, nodes["hs-b.example"].srv

	aliceToken := registerThrough(t, a, "alice")
	bobToken := registerThrough(t, b, "bob")

	rec := post(t, a, "/_matrix/client/v3/createRoom", map[string]any{
		"preset": "public_chat", "name": "Across", "room_alias_name": "across",
	}, aliceToken)
	if rec.Code != http.StatusOK {
		t.Fatalf("createRoom: %d %s", rec.Code, rec.Body.String())
	}
	roomID, _ := decode[map[string]any](t, rec)["room_id"].(string)

	// Bob joins by alias: the alias is resolved by asking A, then the join
	// goes through A.
	rec = post(t, b, "/_matrix/client/v3/join/"+url.PathEscape("#across:hs-a.example"), map[string]any{}, bobToken)
	if rec.Code != http.StatusOK {
		t.Fatalf("remote join: %d %s", rec.Code, rec.Body.String())
	}
	if got, _ := decode[map[string]any](t, rec)["room_id"].(string); got != roomID {
		t.Fatalf("joined %q, want %q", got, roomID)
	}

	// Alice hears of the join without any queue involved: it went
	// through her server.
	rec = getWithToken(t, a, "/_matrix/client/v3/rooms/"+url.PathEscape(roomID)+"/joined_members", aliceToken)
	if !strings.Contains(rec.Body.String(), "@bob:hs-b.example") {
		t.Fatalf("A's joined members lack bob: %s", rec.Body.String())
	}

	// Messages in both directions go through the queues.
	sendMessage(t, a, roomID, aliceToken, "a1", "hello from A")
	eventually(t, "A's message on B", func() bool {
		return contains(messageBodies(t, b, roomID, bobToken), "hello from A")
	})
	sendMessage(t, b, roomID, bobToken, "b1", "hello from B")
	eventually(t, "B's message on A", func() bool {
		return contains(messageBodies(t, a, roomID, aliceToken), "hello from B")
	})

	// Remote profile lookup.
	rec = getWithToken(t, b, "/_matrix/client/v3/profile/@alice:hs-a.example", bobToken)
	if rec.Code != http.StatusOK {
		t.Errorf("remote profile: %d %s", rec.Code, rec.Body.String())
	}

	// Switching A's federation off: B's next message cannot be delivered,
	// and B records the failure. Switching it back on and clearing the
	// backoff delivers it.
	if err := a.Settings().SetFederationEnabled(t.Context(), false); err != nil {
		t.Fatal(err)
	}
	a.FederationModule().Stop()
	sendMessage(t, b, roomID, bobToken, "b2", "while A is away")
	eventually(t, "B to record A as failing", func() bool {
		ds, err := b.FederationModule().Queue().Destinations(t.Context())
		if err != nil {
			return false
		}
		for _, d := range ds {
			if d.ServerName == "hs-a.example" && d.Failures > 0 && d.QueuedPDUs > 0 {
				return true
			}
		}
		return false
	})
	if contains(messageBodies(t, a, roomID, aliceToken), "while A is away") {
		t.Fatal("a message reached A while its federation was off")
	}

	// The admin API shows the backlog and the backoff.
	root, err := b.Users().CreateAccount(t.Context(), userapi.NewAccountRequest{
		Localpart: "root", Password: "sUp3rs3kr1t", Admin: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	session, err := b.Users().CreateSession(t.Context(), userapi.SessionRequest{Localpart: root.Localpart})
	if err != nil {
		t.Fatal(err)
	}
	rec = getWithToken(t, b, "/_warp/admin/v1/federation/destinations", session.AccessToken)
	if rec.Code != http.StatusOK {
		t.Fatalf("GET destinations: %d %s", rec.Code, rec.Body.String())
	}
	var listing struct {
		Running      bool `json:"running"`
		Destinations []struct {
			ServerName string `json:"server_name"`
			Failures   int    `json:"failures"`
			QueuedPDUs int    `json:"queued_pdus"`
		} `json:"destinations"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &listing); err != nil {
		t.Fatal(err)
	}
	if !listing.Running || len(listing.Destinations) == 0 ||
		listing.Destinations[0].ServerName != "hs-a.example" || listing.Destinations[0].Failures == 0 {
		t.Errorf("destinations = %s", rec.Body.String())
	}

	if err := a.Settings().SetFederationEnabled(t.Context(), true); err != nil {
		t.Fatal(err)
	}
	if err := a.FederationModule().Start(t.Context()); err != nil {
		t.Fatal(err)
	}
	if err := b.FederationModule().Queue().ResetBackoff(t.Context(), "hs-a.example"); err != nil {
		t.Fatal(err)
	}
	eventually(t, "the held message on A", func() bool {
		return contains(messageBodies(t, a, roomID, aliceToken), "while A is away")
	})
}

func TestFederationEndpointsAreGated(t *testing.T) {
	srv := newSigningServer(t, nil)
	for _, path := range []string{
		"/_matrix/key/v2/server",
		"/_matrix/federation/v1/version",
	} {
		if rec := get(t, srv, http.MethodGet, path); rec.Code != http.StatusNotFound {
			t.Errorf("%s with federation off = %d, want 404", path, rec.Code)
		}
	}

	if err := srv.Settings().SetFederationEnabled(t.Context(), true); err != nil {
		t.Fatal(err)
	}
	rec := get(t, srv, http.MethodGet, "/_matrix/key/v2/server")
	if rec.Code != http.StatusOK {
		t.Fatalf("key server with federation on = %d", rec.Code)
	}
	if name, _ := decode[map[string]any](t, rec)["server_name"].(string); name != "warp.example" {
		t.Errorf("key response server_name = %q", name)
	}

	// An unsigned request to an authenticated endpoint is refused.
	rec = get(t, srv, http.MethodGet, "/_matrix/federation/v1/event/$abc")
	if rec.Code != http.StatusUnauthorized {
		t.Errorf("unsigned /event = %d, want 401", rec.Code)
	}
}

// writeSelfSignedCert writes a certificate for 127.0.0.1 and its key, and
// returns their paths and a pool that trusts the certificate.
func writeSelfSignedCert(t *testing.T) (string, string, *x509.CertPool) {
	t.Helper()
	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	tmpl := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject:      pkix.Name{CommonName: "warp.example"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(time.Hour),
		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")},
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &priv.PublicKey, priv)
	if err != nil {
		t.Fatal(err)
	}
	keyDER, err := x509.MarshalECPrivateKey(priv)
	if err != nil {
		t.Fatal(err)
	}
	dir := t.TempDir()
	certPath, keyPath := filepath.Join(dir, "cert.pem"), filepath.Join(dir, "key.pem")
	if err := os.WriteFile(certPath, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(keyPath, pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}), 0o600); err != nil {
		t.Fatal(err)
	}
	cert, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatal(err)
	}
	pool := x509.NewCertPool()
	pool.AddCert(cert)
	return certPath, keyPath, pool
}

func TestFederationListenerServesTLS(t *testing.T) {
	certPath, keyPath, pool := writeSelfSignedCert(t)
	srv := newSigningServer(t, func(c *config.Config) {
		c.Listen.Address = "127.0.0.1:0"
		c.Listen.FederationAddress = "127.0.0.1:0"
		c.Listen.TLSCertFile = certPath
		c.Listen.TLSKeyFile = keyPath
		c.Federation.Enabled = true
		c.Database.ConnectionString = "file:" + filepath.ToSlash(filepath.Join(t.TempDir(), "warp.db"))
	})

	ctx, cancel := context.WithCancel(t.Context())
	done := make(chan error, 1)
	go func() { done <- srv.Run(ctx) }()
	defer func() {
		cancel()
		<-done
	}()

	var addr string
	for range 200 {
		if addr = srv.FederationAddr(); addr != "" {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if addr == "" {
		t.Fatal("the federation listener never started")
	}
	if !srv.FederationModule().Running() {
		t.Error("Run did not start the federation module with federation on")
	}

	client := &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12}}}
	resp, err := client.Get("https://" + addr + "/_matrix/key/v2/server")
	if err != nil {
		t.Fatalf("GET over TLS: %v", err)
	}
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Errorf("key server over TLS = %d", resp.StatusCode)
	}

	// The runtime switch stops the module while Run is serving.
	if err := srv.Settings().SetFederationEnabled(t.Context(), false); err != nil {
		t.Fatal(err)
	}
	if srv.FederationModule().Running() {
		t.Error("switching federation off did not stop the module")
	}
	if err := srv.Settings().SetFederationEnabled(t.Context(), true); err != nil {
		t.Fatal(err)
	}
	if !srv.FederationModule().Running() {
		t.Error("switching federation back on did not start the module")
	}
}

func TestFederationListenerNeedsACertificate(t *testing.T) {
	cfg := config.Default()
	cfg.Global.ServerName = "warp.example"
	cfg.Listen.FederationAddress = ":8448"
	if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "tls_cert_file") {
		t.Errorf("Validate = %v, want a complaint about the certificate", err)
	}
}