package federation

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

	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
	"github.com/dariuszjaniak/warp/server/internal/storage"
	_ "github.com/dariuszjaniak/warp/server/internal/storage/drivers"
)

// testStore opens a migrated SQLite database for the key cache.
func testStore(t *testing.T) storage.FederationStore {
	t.Helper()
	db, err := storage.Open(t.Context(), storage.Config{
		Driver:           storage.DriverSQLite,
		ConnectionString: "file:" + strings.ReplaceAll(t.TempDir(), `\`, "/") + "/warp.db",
	}, nil)
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })
	if err := db.Migrate(t.Context()); err != nil {
		t.Fatalf("migrate: %v", err)
	}
	return db.Federation()
}

// remote is a fake homeserver: a TLS server answering for a server name,
// with its own key.
type remote struct {
	name    string
	key     *signing.KeyPair
	srv     *httptest.Server
	mux     *http.ServeMux
	keyHits atomic.Int32

	// validFor is the validity its key response claims.
	validFor time.Duration
}

func newRemote(t *testing.T, name string) *remote {
	t.Helper()
	key, err := signing.Generate("")
	if err != nil {
		t.Fatal(err)
	}
	r := &remote{name: name, key: key, mux: http.NewServeMux(), validFor: time.Hour}
	r.mux.HandleFunc("GET /_matrix/key/v2/server", func(w http.ResponseWriter, _ *http.Request) {
		r.keyHits.Add(1)
		doc, _ := json.Marshal(serverKeyResponse{
			ServerName:    r.name,
			ValidUntilTS:  time.Now().Add(r.validFor).UnixMilli(),
			VerifyKeys:    map[string]verifyKey{r.key.KeyID(): {Key: r.key.PublicBase64()}},
			OldVerifyKeys: map[string]oldVerifyKey{},
		})
		signed, _ := signing.SignJSON(r.name, r.key, doc)
		_, _ = w.Write(signed)
	})
	r.srv = httptest.NewTLSServer(r.mux)
	t.Cleanup(r.srv.Close)
	return r
}

// clientFor builds a client that resolves every remote in the list to its
// test server. The Host is "example.com", the name the test certificate is
// issued for, and the address is the listener: the dial-target path is
// exercised on every request.
func clientFor(t *testing.T, serverName string, key *signing.KeyPair, remotes ...*remote) *Client {
	t.Helper()
	byName := map[string]*remote{}
	for _, r := range remotes {
		byName[r.name] = r
	}
	var tlsConfig = remotes[0].srv.Client().Transport.(*http.Transport).TLSClientConfig.Clone()
	return NewClient(ClientOptions{
		ServerName: serverName,
		Key:        key,
		TLSConfig:  tlsConfig,
		Enabled:    func() bool { return true },
		Resolve: func(_ context.Context, name string) (Resolved, error) {
			r, ok := byName[name]
			if !ok {
				return Resolved{}, errors.New("no such test server")
			}
			_, port, _ := net.SplitHostPort(r.srv.Listener.Addr().String())
			return Resolved{Addresses: []string{r.srv.Listener.Addr().String()}, Host: "example.com:" + port}, nil
		},
	})
}

func newKeyRing(t *testing.T, name string, remotes ...*remote) (*KeyRing, *signing.KeyPair) {
	t.Helper()
	key, err := signing.Generate("")
	if err != nil {
		t.Fatal(err)
	}
	return NewKeyRing(KeyRingOptions{
		ServerName: name,
		Key:        key,
		Store:      testStore(t),
		Client:     clientFor(t, name, key, remotes...),
	}), key
}

func TestKeyRingFetchesAndCaches(t *testing.T) {
	r := newRemote(t, "remote.example")
	ring, _ := newKeyRing(t, "warp.example", r)

	pub, err := ring.PublicKey(t.Context(), "remote.example", r.key.KeyID(), time.Now().UnixMilli())
	if err != nil {
		t.Fatalf("PublicKey: %v", err)
	}
	if !pub.Equal(r.key.Public) {
		t.Error("fetched the wrong key")
	}
	for range 3 {
		if _, err := ring.PublicKey(t.Context(), "remote.example", r.key.KeyID(), time.Now().UnixMilli()); err != nil {
			t.Fatal(err)
		}
	}
	if n := r.keyHits.Load(); n != 1 {
		t.Errorf("%d key fetches, want 1", n)
	}

	// A key needed for a time past its validity is fetched again, in case
	// the server has extended it.
	if _, err := ring.PublicKey(t.Context(), "remote.example", r.key.KeyID(),
		time.Now().Add(2*time.Hour).UnixMilli()); err == nil {
		t.Error("a key was accepted for a time past its validity")
	}
	if n := r.keyHits.Load(); n != 2 {
		t.Errorf("%d key fetches after asking past validity, want 2", n)
	}

	// Validity is capped at a week, whatever the server claims.
	r.validFor = 365 * 24 * time.Hour
	ring2, _ := newKeyRing(t, "warp.example", r)
	if _, err := ring2.PublicKey(t.Context(), "remote.example", r.key.KeyID(),
		time.Now().Add(30*24*time.Hour).UnixMilli()); err == nil {
		t.Error("a key was believed for longer than a week")
	}
}

func TestKeyRingRefusesAForgedResponse(t *testing.T) {
	r := newRemote(t, "remote.example")
	impostor, _ := signing.Generate("")
	r.mux = http.NewServeMux()
	r.mux.HandleFunc("GET /_matrix/key/v2/server", func(w http.ResponseWriter, _ *http.Request) {
		// Lists the real key but is signed with another one.
		doc, _ := json.Marshal(serverKeyResponse{
			ServerName:   r.name,
			ValidUntilTS: time.Now().Add(time.Hour).UnixMilli(),
			VerifyKeys:   map[string]verifyKey{r.key.KeyID(): {Key: impostor.PublicBase64()}},
		})
		signed, _ := signing.SignJSON(r.name, r.key, doc)
		_, _ = w.Write(signed)
	})
	r.srv.Config.Handler = r.mux

	ring, _ := newKeyRing(t, "warp.example", r)
	if _, err := ring.PublicKey(t.Context(), "remote.example", r.key.KeyID(), 0); !errors.Is(err, ErrUnknownKey) {
		t.Errorf("PublicKey with a forged response = %v, want ErrUnknownKey", err)
	}
}

func TestOwnKeyResponseVerifies(t *testing.T) {
	ring, key := newKeyRing(t, "warp.example", newRemote(t, "unused.example"))
	doc, err := ring.OwnKeyResponse()
	if err != nil {
		t.Fatal(err)
	}
	if err := signing.VerifyJSON("warp.example", key.KeyID(), key.Public, doc); err != nil {
		t.Errorf("our own key response does not verify: %v", err)
	}
	var parsed serverKeyResponse
	if err := json.Unmarshal(doc, &parsed); err != nil {
		t.Fatal(err)
	}
	if parsed.VerifyKeys[key.KeyID()].Key != key.PublicBase64() {
		t.Errorf("verify_keys = %+v", parsed.VerifyKeys)
	}
}

func TestSignedRequestIsAuthenticatedEndToEnd(t *testing.T) {
	// The remote here is the *sender*: a key server for remote.example,
	// and a client signing as remote.example to warp.example.
	sender := newRemote(t, "remote.example")
	receiver := newRemote(t, "warp.example")
	ring := NewKeyRing(KeyRingOptions{
		ServerName: "warp.example",
		Key:        receiver.key,
		Store:      testStore(t),
		Client:     clientFor(t, "warp.example", receiver.key, sender),
	})

	var gotOrigin atomic.Value
	receiver.mux.HandleFunc("PUT /_matrix/federation/v1/send/{txn}", func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		origin, err := ring.AuthenticateRequest(r.Context(), r, body)
		if err != nil {
			http.Error(w, `{"errcode":"M_UNAUTHORIZED","error":"`+err.Error()+`"}`, http.StatusUnauthorized)
			return
		}
		gotOrigin.Store(origin)
		_, _ = w.Write([]byte(`{"pdus":{}}`))
	})

	client := clientFor(t, "remote.example", sender.key, receiver)
	if _, err := client.SendTransaction(t.Context(), "warp.example", "txn/1 with spaces", Transaction{
		Origin: "remote.example", OriginServerTS: 1, PDUs: []json.RawMessage{},
	}); err != nil {
		t.Fatalf("SendTransaction: %v", err)
	}
	if got, _ := gotOrigin.Load().(string); got != "remote.example" {
		t.Errorf("authenticated origin = %q", got)
	}

	// A request signed by a key the origin does not publish is refused.
	stranger, _ := signing.Generate("")
	forged := clientFor(t, "remote.example", stranger, receiver)
	_, err := forged.SendTransaction(t.Context(), "warp.example", "2", Transaction{Origin: "remote.example"})
	var httpErr *HTTPError
	if !errors.As(err, &httpErr) || httpErr.Status != http.StatusUnauthorized {
		t.Errorf("forged request = %v, want 401", err)
	}
}

func TestClientRefusesWhileDisabledOrDenied(t *testing.T) {
	r := newRemote(t, "remote.example")
	key, _ := signing.Generate("")
	enabled := false
	client := clientFor(t, "warp.example", key, r)
	client.enabled = func() bool { return enabled }

	if _, err := client.Version(t.Context(), "remote.example"); !errors.Is(err, ErrDisabled) {
		t.Errorf("Version while disabled = %v, want ErrDisabled", err)
	}
	enabled = true
	client.policy = NewPolicy(nil, []string{"remote.example"})
	if _, err := client.Version(t.Context(), "remote.example"); !errors.Is(err, ErrNotAllowed) {
		t.Errorf("Version to a denied server = %v, want ErrNotAllowed", err)
	}
	if _, err := client.Version(t.Context(), "warp.example"); !errors.Is(err, ErrSelf) {
		t.Errorf("Version to ourselves = %v, want ErrSelf", err)
	}
	if r.keyHits.Load() != 0 {
		t.Error("a refused request reached the network")
	}
}