package sesame

import (
	"bytes"
	"context"
	"crypto/ecdh"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"sync/atomic"
	"testing"
	"time"
)

func waitRegistration(t *testing.T, p *RegistrationConnection, phase string) {
	t.Helper()
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		if p.State().Phase == phase {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("registration phase = %+v, want %s", p.State(), phase)
}
func linkClosed(f *fakeLink) bool { f.mu.Lock(); defer f.mu.Unlock(); return f.closed }

func TestPreparedRegistrationWaitsAndHandsOffOriginalSession(t *testing.T) {
	f := newLink()
	device, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	var peer *sessionCipher
	var registrations, logins, dials atomic.Int32
	var expectedSecret string
	f.write = func(kind byte, payload []byte) error {
		if kind == 1 {
			if payload[0] == 2 {
				logins.Add(1)
				return errors.New("registration must not login again")
			}
			if len(payload) != 69 || payload[0] != 1 {
				return errors.New("bad registration")
			}
			registrations.Add(1)
			pub, err := ecdh.P256().NewPublicKey(append([]byte{4}, payload[1:65]...))
			if err != nil {
				return err
			}
			shared, err := device.ECDH(pub)
			if err != nil {
				return err
			}
			expectedSecret = hex.EncodeToString(shared[:16])
			peer, _, err = newCipher(shared[:16], unhex("11223344"))
			if err != nil {
				return err
			}
			reply := make([]byte, 77)
			copy(reply, []byte{0, 0, 0, 128, 45, 0, 18})
			copy(reply[7:], []byte{90, 0, 0, 0, 0, 0})
			copy(reply[13:], device.PublicKey().Bytes()[1:])
			f.emit(1, append([]byte{7, 1, 0}, reply...))
			return nil
		}
		plain, err := peer.decrypt(payload)
		if err != nil {
			return err
		}
		if !bytes.Equal(plain, []byte{83, 0, 14}) {
			return errors.New("wrong live-session operation")
		}
		ack, err := peer.encrypt([]byte{7, 83, 0})
		if err != nil {
			return err
		}
		f.emit(2, ack)
		status, err := peer.encrypt([]byte{8, 81, 0, 0, 0, 128, 46, 0, 16})
		if err != nil {
			return err
		}
		f.emit(2, status)
		return nil
	}
	p := PrepareRegistration(context.Background(), testTarget, func(context.Context, Target) (Link, Target, error) { dials.Add(1); return f, testTarget, nil })
	defer p.Close()
	waitRegistration(t, p, "ready")
	if registrations.Load() != 0 || linkClosed(f) {
		t.Fatal("acted before Register")
	}
	saved := ""
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	c, err := p.Register(ctx, func(key Credential) error { saved = key.Secret; return nil })
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	if saved == "" || saved != expectedSecret {
		t.Fatal("credential not persisted")
	}
	p.Close()
	if linkClosed(f) {
		t.Fatal("handoff closed the live connection")
	}
	if state := c.State(); !state.Online || state.Position == nil || *state.Position != 45 || state.LockPosition == nil || *state.LockPosition != 90 {
		t.Fatalf("registration state lost: %+v", state)
	}
	life, stop := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { c.Run(life); close(done) }()
	defer func() { stop(); <-done }()
	if err = c.Action(ctx, "unlock"); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(time.Second)
	for c.State().Position != nil && *c.State().Position != 46 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if c.State().Position == nil || *c.State().Position != 46 {
		t.Fatal("live notifications not handed off")
	}
	if dials.Load() != 1 || registrations.Load() != 1 || logins.Load() != 0 {
		t.Fatal("unexpected reconnect, replay or login")
	}
}

func TestPreparedRegistrationCannotRegisterBeforeInitialAndClosesOnCancel(t *testing.T) {
	f := newLink()
	<-f.packets
	var writes atomic.Int32
	f.write = func(byte, []byte) error { writes.Add(1); return nil }
	p := PrepareRegistration(context.Background(), testTarget, func(context.Context, Target) (Link, Target, error) { return f, testTarget, nil })
	waitRegistration(t, p, "waiting-token")
	if _, err := p.Register(context.Background(), func(Credential) error { return nil }); err == nil {
		t.Fatal("registered before initial token")
	}
	p.Close()
	if !linkClosed(f) || writes.Load() != 0 {
		t.Fatal("cancel did not clean up idle link")
	}
}

func TestPreparedRegistrationDisconnectAndTimeoutNeverReplay(t *testing.T) {
	for _, lost := range []bool{true, false} {
		t.Run(map[bool]string{true: "disconnect", false: "exchange-timeout"}[lost], func(t *testing.T) {
			f := newLink()
			var writes atomic.Int32
			f.write = func(byte, []byte) error { writes.Add(1); return nil }
			p := PrepareRegistration(context.Background(), testTarget, func(context.Context, Target) (Link, Target, error) { return f, testTarget, nil })
			defer p.Close()
			waitRegistration(t, p, "ready")
			if lost {
				f.done <- errors.New("lost link")
				waitRegistration(t, p, "disconnected")
			}
			ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
			defer cancel()
			_, err := p.Register(ctx, func(Credential) error { t.Fatal("persisted without response"); return nil })
			if err == nil {
				t.Fatal("accepted incomplete registration")
			}
			<-p.done
			want := int32(1)
			if lost {
				want = 0
			}
			if writes.Load() != want || !linkClosed(f) {
				t.Fatal("replayed command or leaked link")
			}
			if _, err = p.Register(context.Background(), func(Credential) error { return nil }); err == nil {
				t.Fatal("reused failed registration")
			}
			if writes.Load() != want {
				t.Fatal("replayed failed registration")
			}
		})
	}
}

func TestPreparedRegistrationCancelsInFlightSetup(t *testing.T) {
	started := make(chan struct{})
	p := PrepareRegistration(context.Background(), testTarget, func(ctx context.Context, _ Target) (Link, Target, error) {
		close(started)
		<-ctx.Done()
		return nil, Target{}, ctx.Err()
	})
	<-started
	p.Close()
	if p.State().Connected {
		t.Fatal("connected after setup cancellation")
	}
}