package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
	"github.com/nats-io/nkeys"
)

type env struct {
	k             Keys
	op, sys, acc  nkeys.KeyPair
	sig           nkeys.KeyPair
	revoked       map[string]time.Time
	lastIat       int64
	accountDoneAt time.Time
}

func (e *env) url(ip string) string { return "nats://" + ip + ":4222" }

func say(format string, a ...any) {
	fmt.Printf("%s  %s\n", time.Now().Format("15:04:05.000"), fmt.Sprintf(format, a...))
}

func ms(d time.Duration) string { return fmt.Sprintf("%.1f ms", float64(d.Microseconds())/1000) }

type user struct {
	kp  nkeys.KeyPair
	opt nats.Option
}

func (e *env) newUser(name string) user {
	kp := must(nkeys.CreateUser())
	j := userJWT(e.acc, e.sig, kp, name)
	return user{kp, nats.UserJWT(func() (string, error) { return j, nil }, func(n []byte) ([]byte, error) { return kp.Sign(n) })}
}

func (e *env) sysConn(ip string) *nats.Conn {
	kp := must(nkeys.CreateUser())
	uc := newSysUser(e.sys, kp)
	nc, err := nats.Connect(e.url(ip), nats.UserJWT(func() (string, error) { return uc, nil }, func(n []byte) ([]byte, error) { return kp.Sign(n) }), nats.Timeout(5*time.Second))
	if err != nil {
		panic(fmt.Sprintf("system connection to %s: %v", ip, err))
	}
	return nc
}

// probe is one login attempt from this machine across the network.
func (e *env) probe(ip string, u user) error {
	nc, err := nats.Connect(e.url(ip), u.opt, nats.NoReconnect(), nats.Timeout(3*time.Second))
	if err == nil {
		nc.Close()
	}
	return err
}

type open struct {
	nc     *nats.Conn
	closed atomic.Int64
}

func (e *env) openConn(ip string, u user) *open {
	o := &open{}
	nc, err := nats.Connect(e.url(ip), u.opt, nats.NoReconnect(), nats.Timeout(3*time.Second),
		nats.ClosedHandler(func(*nats.Conn) { o.closed.Store(time.Now().UnixNano()) }))
	if err != nil {
		panic(fmt.Sprintf("open connection to %s: %v", ip, err))
	}
	nc.Flush()
	o.nc = nc
	return o
}

// push sends an account JWT to one server's resolver over the system account. The reply is not proof it applied.
func (e *env) push(sys *nats.Conn, j string) (string, time.Duration) {
	t0 := time.Now()
	m, err := sys.Request("$SYS.REQ.ACCOUNT."+pubOf(e.acc)+".CLAIMS.UPDATE", []byte(j), 5*time.Second)
	if err != nil {
		return "error: " + err.Error(), time.Since(t0)
	}
	return string(m.Data), time.Since(t0)
}

// nextAccountJWT builds the account JWT with the revocation list; its iat has to be a later second than the last one.
func (e *env) nextAccountJWT() string {
	for time.Now().Unix() <= e.lastIat {
		time.Sleep(50 * time.Millisecond)
	}
	e.lastIat = time.Now().Unix()
	return accountJWT(e.op, e.acc, e.sig, e.revoked)
}

func (e *env) refusedAfter(ip string, u user, t0 time.Time, limit time.Duration) (time.Duration, int) {
	n := 0
	for time.Since(t0) < limit {
		if err := e.probe(ip, u); err != nil {
			return time.Since(t0), n
		}
		n++
		time.Sleep(time.Millisecond)
	}
	return -1, n
}

func newSysUser(sys, ukp nkeys.KeyPair) string { return userClaimsFor(sys, ukp) }

func drive(args []string) {
	fs := flag.NewFlagSet("drive", flag.ExitOnError)
	keysFile := fs.String("keys", "/root/q8/keys.json", "keys")
	only := fs.String("only", "2,4,3", "which parts, in order")
	trials := fs.Int("trials", 3, "trials per part")
	fs.Parse(args)
	raw, err := os.ReadFile(*keysFile)
	if err != nil {
		panic(err)
	}
	e := &env{revoked: map[string]time.Time{}}
	json.Unmarshal(raw, &e.k)
	e.op, e.sys, e.acc, e.sig = fromSeed(e.k.Op), fromSeed(e.k.Sys), fromSeed(e.k.Acc), fromSeed(e.k.Sig)
	e.lastIat = time.Now().Unix()
	// the account JWT the servers were preloaded with has an iat from gen time; ours must be later
	time.Sleep(1100 * time.Millisecond)
	if d := waitRoutes(e, 1, 90*time.Second); d < 0 {
		say("the route between B and C did not come up")
	} else {
		say("cluster formed: route B-C up after %s", ms(d))
	}
	for _, p := range strings.Split(*only, ",") {
		switch p {
		case "2":
			part2(e, *trials)
		case "3":
			part3(e, *trials)
		case "4":
			part4(e)
		}
	}
	say("DRIVE-DONE")
}

func routes(ip string) int {
	r, err := http.Get("http://" + ip + ":8222/routez")
	if err != nil {
		return -1
	}
	defer r.Body.Close()
	var v struct {
		NumRoutes int `json:"num_routes"`
	}
	json.NewDecoder(r.Body).Decode(&v)
	return v.NumRoutes
}

func waitRoutes(e *env, want int, limit time.Duration) time.Duration {
	t0 := time.Now()
	for time.Since(t0) < limit {
		if routes(e.k.B) == want && routes(e.k.C) == want {
			return time.Since(t0)
		}
		time.Sleep(100 * time.Millisecond)
	}
	return -1
}

func agent(e *env, what string) (string, time.Time) {
	r, err := http.Get("http://" + e.k.C + ":9099/" + what)
	t := time.Now()
	if err != nil {
		return "error: " + err.Error(), t
	}
	defer r.Body.Close()
	b, _ := io.ReadAll(r.Body)
	return string(b), t
}

// ---- part 2: revocation across the network ----

func part2(e *env, n int) {
	say("=== PART 2: revocation over the network (B and C, pushed to B from A)")
	sysB := e.sysConn(e.k.B)
	defer sysB.Close()
	for run := 1; run <= n; run++ {
		u := e.newUser(fmt.Sprintf("p2-%d", run))
		ctl := e.newUser(fmt.Sprintf("p2-ctl-%d", run))
		ob, oc := e.openConn(e.k.B, u), e.openConn(e.k.C, u)
		cb, cc := e.openConn(e.k.B, ctl), e.openConn(e.k.C, ctl)
		e.revoked[pubOf(u.kp)] = time.Now()
		j := e.nextAccountJWT()
		t0 := time.Now()
		reply, rtt := e.push(sysB, j)
		type res struct {
			ip      string
			refused time.Duration
			acc     int
		}
		var wg sync.WaitGroup
		out := make([]res, 2)
		for i, ip := range []string{e.k.B, e.k.C} {
			wg.Add(1)
			go func() {
				defer wg.Done()
				d, a := e.refusedAfter(ip, u, t0, 30*time.Second)
				out[i] = res{ip, d, a}
			}()
		}
		wg.Wait()
		time.Sleep(500 * time.Millisecond)
		cl := func(o *open) string {
			if v := o.closed.Load(); v != 0 {
				return ms(time.Duration(v - t0.UnixNano()))
			}
			return "NOT closed"
		}
		alive := func(o *open) bool { return o.nc.IsConnected() }
		say("RESULT part2 run %d: push reply %q rtt %s; new login refused after: B %s (%d accepted before), C %s (%d accepted before); open connection closed after: B %s, C %s; control user still connected: B %v, C %v",
			run, strings.TrimSpace(reply)[:min(len(strings.TrimSpace(reply)), 60)], ms(rtt), ms(out[0].refused), out[0].acc, ms(out[1].refused), out[1].acc, cl(ob), cl(oc), alive(cb), alive(cc))
		for _, o := range []*open{ob, oc, cb, cc} {
			o.nc.Close()
		}
	}
}

// ---- part 3: C cut off while a user is revoked ----

type seen struct {
	at      time.Time
	subject string
	data    []byte
}

func part3(e *env, n int) {
	say("=== PART 3: C cut off from B while a user is revoked; modes a (periodic sync only) and b (re-push on C's return)")
	for _, mode := range []string{"a", "b"} {
		for run := 1; run <= n; run++ {
			trial3(e, mode, run)
		}
	}
}

func trial3(e *env, mode string, run int) {
	say("--- trial %s%d", mode, run)
	if d := waitRoutes(e, 1, 90*time.Second); d < 0 {
		say("RESULT part3 %s%d: the route B-C did not come up (routes B=%d C=%d); trial skipped", mode, run, routes(e.k.B), routes(e.k.C))
		return
	}
	time.Sleep(3 * time.Second)
	sysB, sysC := e.sysConn(e.k.B), e.sysConn(e.k.C)
	defer sysB.Close()
	defer sysC.Close()
	u, ctl := e.newUser(fmt.Sprintf("p3-%s%d", mode, run)), e.newUser(fmt.Sprintf("p3-ctl-%s%d", mode, run))
	oc, cc := e.openConn(e.k.C, u), e.openConn(e.k.C, ctl)
	defer oc.nc.Close()
	defer cc.nc.Close()
	var mu sync.Mutex
	var events []seen
	sysB.Subscribe("$SYS.>", func(m *nats.Msg) {
		mu.Lock()
		events = append(events, seen{time.Now(), m.Subject, m.Data})
		mu.Unlock()
	})
	sysB.Flush()

	r, tCut := agent(e, "block")
	say("cut: agent said %q", r)
	dr := waitRoutes(e, 0, 60*time.Second)
	say("routes gone on B and C after %s (ping interval 2s, 2 missed pings)", ms(dr))
	time.Sleep(1100 * time.Millisecond)
	e.revoked[pubOf(u.kp)] = time.Now()
	j := e.nextAccountJWT()
	tPush := time.Now()
	reply, rtt := e.push(sysB, j)
	dB, _ := e.refusedAfter(e.k.B, u, tPush, 10*time.Second)
	time.Sleep(4 * time.Second)
	accC := e.probe(e.k.C, u)
	say("partitioned: push to B reply %q rtt %s; B refused the user after %s; C still accepts the user: %v (cut lasted %s so far)", strings.TrimSpace(reply)[:min(len(strings.TrimSpace(reply)), 40)], ms(rtt), ms(dB), accC == nil, time.Since(tCut).Round(time.Second))

	mu.Lock()
	events = nil
	mu.Unlock()
	r, tHeal := agent(e, "unblock")
	say("healed: agent said %q", r)

	var trigger atomic.Int64 // offset ns of the event used as trigger
	var triggerSubj atomic.Value
	var pushedRTT atomic.Int64
	var pushedReply atomic.Value
	if mode == "b" {
		go func() {
			deadline := time.Now().Add(150 * time.Second)
			for time.Now().Before(deadline) && trigger.Load() == 0 {
				mu.Lock()
				var hit *seen
				for i := range events {
					ev := events[i]
					if strings.HasPrefix(ev.subject, "$SYS.SERVER.") && strings.HasSuffix(ev.subject, ".STATSZ") && ev.at.After(tHeal) {
						var v struct {
							Server struct{ Name string } `json:"server"`
						}
						if json.Unmarshal(ev.data, &v) == nil && v.Server.Name == "C" {
							hit = &ev
							break
						}
					}
				}
				mu.Unlock()
				if hit != nil {
					trigger.Store(int64(hit.at.Sub(tHeal)))
					triggerSubj.Store(hit.subject)
					rep, rt := e.push(sysC, j)
					pushedRTT.Store(int64(rt))
					pushedReply.Store(rep)
					return
				}
				time.Sleep(time.Millisecond)
			}
		}()
	}
	refused, accepted := e.refusedAfter(e.k.C, u, tHeal, 150*time.Second)
	time.Sleep(700 * time.Millisecond)
	closedAt := "NOT closed"
	if v := oc.closed.Load(); v != 0 {
		closedAt = ms(time.Duration(v - tHeal.UnixNano()))
	}
	line := fmt.Sprintf("RESULT part3 %s%d: after the link came back C accepted the revoked user for %s (%d logins accepted), open connection closed after %s, control user still connected %v",
		mode, run, ms(refused), accepted, closedAt, cc.nc.IsConnected())
	if mode == "b" {
		if t := trigger.Load(); t != 0 {
			subj, _ := triggerSubj.Load().(string)
			rep, _ := pushedReply.Load().(string)
			line += fmt.Sprintf("; trigger: first STATSZ from server C arrived on B's system account %s after the link came back (%s), push to C reply %q rtt %s",
				ms(time.Duration(t)), strings.ReplaceAll(subj, "$SYS.SERVER.", "$SYS.SERVER.<id>."), strings.TrimSpace(rep)[:min(len(strings.TrimSpace(rep)), 30)], ms(time.Duration(pushedRTT.Load())))
		} else {
			line += "; trigger NEVER fired"
		}
	}
	say("%s", line)
	if run == 1 {
		mu.Lock()
		first := map[string]time.Duration{}
		count := map[string]int{}
		for _, ev := range events {
			k := normalise(ev.subject)
			if _, ok := first[k]; !ok {
				first[k] = ev.at.Sub(tHeal)
			}
			count[k]++
		}
		mu.Unlock()
		var ks []string
		for k := range first {
			ks = append(ks, k)
		}
		sort.Slice(ks, func(i, j int) bool { return first[ks[i]] < first[ks[j]] })
		say("system-account subjects seen on B after the link came back (trial %s%d), first-seen offset, count:", mode, run)
		for _, k := range ks {
			say("    %9s  x%-3d %s", ms(first[k]), count[k], k)
		}
	}
}

func normalise(s string) string {
	parts := strings.Split(s, ".")
	for i, p := range parts {
		if len(p) >= 40 && strings.ToUpper(p) == p {
			parts[i] = "<id>"
		}
	}
	return strings.Join(parts, ".")
}

// ---- part 4: moving a stream and a KV bucket from A to B over the leaf link ----

func part4(e *env) {
	say("=== PART 4: stream S (10000 msgs) and KV bucket cfg moved from A (legacy) to B/C (hub) over the leaf link, writer running")
	ctx := context.Background()
	lc, err := nats.Connect(e.url(e.k.A), nats.UserInfo("w", "pw"))
	if err != nil {
		panic(err)
	}
	defer lc.Close()
	ljs := must(jetstream.New(lc))
	must(ljs.CreateStream(ctx, jetstream.StreamConfig{Name: "S", Subjects: []string{"mesh.s.>"}}))
	t0 := time.Now()
	for i := 0; i < 10000; i++ {
		if _, err := ljs.Publish(ctx, fmt.Sprintf("mesh.s.%d", i%50), []byte(fmt.Sprintf("seed-%d", i)), jetstream.WithMsgID(fmt.Sprintf("seed-%d", i))); err != nil {
			panic(err)
		}
	}
	say("seeded S with 10000 messages on A in %s", time.Since(t0).Round(time.Millisecond))
	kv := must(ljs.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "cfg", History: 5}))
	for i := 0; i < 1000; i++ {
		must(kv.Put(ctx, fmt.Sprintf("k%d", i), []byte(fmt.Sprintf("v%d", i))))
	}
	for i := 0; i < 200; i++ {
		for r := 0; r < 3; r++ {
			must(kv.Put(ctx, fmt.Sprintf("k%d", i), []byte(fmt.Sprintf("v%d-r%d", i, r))))
		}
	}
	kv.Delete(ctx, "k999")
	say("seeded KV cfg on A")

	for t := time.Now(); time.Since(t) < 60*time.Second; time.Sleep(500 * time.Millisecond) {
		r, err := http.Get("http://" + e.k.B + ":8222/leafz")
		if err != nil {
			continue
		}
		var v struct {
			Leafs int `json:"leafnodes"`
		}
		json.NewDecoder(r.Body).Decode(&v)
		r.Body.Close()
		if v.Leafs >= 2 {
			say("leaf links from A to B up (%d connections: MESH and SYS)", v.Leafs)
			break
		}
	}
	u := e.newUser("p4")
	hc, err := nats.Connect(e.url(e.k.B), u.opt)
	if err != nil {
		panic(err)
	}
	defer hc.Close()
	hjs := must(jetstream.New(hc))
	ready := time.Now()
	for time.Since(ready) < 60*time.Second {
		if _, err := hjs.AccountInfo(ctx); err == nil {
			break
		}
		time.Sleep(500 * time.Millisecond)
	}
	say("hub JetStream answered after %s", time.Since(ready).Round(time.Millisecond))

	var stop atomic.Bool
	var acked, kvacked, failed atomic.Int64
	done := make(chan bool)
	go func() {
		for i := 0; !stop.Load(); i++ {
			if _, err := ljs.Publish(ctx, fmt.Sprintf("mesh.s.%d", i%50), []byte(fmt.Sprintf("live-%d", i)), jetstream.WithMsgID(fmt.Sprintf("live-%d", i))); err == nil {
				acked.Add(1)
			} else {
				failed.Add(1)
			}
			if i%10 == 0 {
				if _, err := kv.Put(ctx, fmt.Sprintf("live%d", i%20), []byte(fmt.Sprintf("x%d", i))); err == nil {
					kvacked.Add(1)
				}
			}
			time.Sleep(2 * time.Millisecond)
		}
		close(done)
	}()
	time.Sleep(300 * time.Millisecond)

	tm := time.Now()
	for _, n := range []string{"S", "KV_cfg"} {
		l := must(ljs.Stream(ctx, n))
		cfg := l.CachedInfo().Config
		cfg.Subjects = nil
		cfg.Sources = []*jetstream.StreamSource{{Name: n, External: &jetstream.ExternalStream{APIPrefix: "$JS.legacy.API"}}}
		if _, err := hjs.CreateStream(ctx, cfg); err != nil {
			panic(fmt.Sprintf("create source stream %s: %v", n, err))
		}
	}
	msgs := func(js jetstream.JetStream, n string) uint64 {
		s, err := js.Stream(ctx, n)
		if err != nil {
			return 0
		}
		i, err := s.Info(ctx)
		if err != nil {
			return 0
		}
		return i.State.Msgs
	}
	waitEq := func(n string, limit time.Duration) time.Duration {
		t := time.Now()
		for time.Since(t) < limit {
			if msgs(hjs, n) == msgs(ljs, n) && msgs(hjs, n) > 0 {
				return time.Since(tm)
			}
			time.Sleep(20 * time.Millisecond)
		}
		return -1
	}
	say("catch-up under writes, counts equal after: S %s, KV_cfg %s (the writer keeps going)", ms(waitEq("S", 60*time.Second)), ms(waitEq("KV_cfg", 60*time.Second)))
	time.Sleep(3 * time.Second)
	stop.Store(true)
	<-done
	say("writer stopped: S acked %d (failed %d), KV puts acked %d", acked.Load(), failed.Load(), kvacked.Load())
	tq := time.Now()
	for _, n := range []string{"S", "KV_cfg"} {
		for time.Since(tq) < 30*time.Second && msgs(hjs, n) != msgs(ljs, n) {
			time.Sleep(10 * time.Millisecond)
		}
	}
	say("quiesced: equal after %s; counts S legacy %d hub %d (expect %d); KV_cfg legacy %d hub %d", time.Since(tq).Round(time.Millisecond),
		msgs(ljs, "S"), msgs(hjs, "S"), 10000+acked.Load(), msgs(ljs, "KV_cfg"), msgs(hjs, "KV_cfg"))

	// content: every message once on both sides
	readAll := func(js jetstream.JetStream) (map[string]int, int) {
		s := must(js.Stream(ctx, "S"))
		info := must(s.Info(ctx))
		set := map[string]int{}
		total := 0
		for seq := info.State.FirstSeq; seq <= info.State.LastSeq; seq++ {
			m, err := s.GetMsg(ctx, seq)
			if err != nil {
				continue
			}
			set[string(m.Data)]++
			total++
		}
		return set, total
	}
	ls, lt := readAll(ljs)
	hs, ht := readAll(hjs)
	dup, missing, extra := 0, 0, 0
	for k, c := range hs {
		if c > 1 {
			dup += c - 1
		}
		if _, ok := ls[k]; !ok {
			extra++
		}
	}
	for k := range ls {
		if _, ok := hs[k]; !ok {
			missing++
		}
	}
	say("RESULT part4 stream S: messages legacy %d (distinct %d) hub %d (distinct %d); on hub: duplicates %d, missing vs legacy %d, not on legacy %d; expected %d = 10000 seed + %d acked live",
		lt, len(ls), ht, len(hs), dup, missing, extra, 10000+acked.Load(), acked.Load())

	hkv := must(hjs.KeyValue(ctx, "cfg"))
	lk := must(kv.Keys(ctx))
	hk := must(hkv.Keys(ctx))
	diff := 0
	for _, k := range lk {
		a, err1 := kv.Get(ctx, k)
		b, err2 := hkv.Get(ctx, k)
		if err1 != nil || err2 != nil || string(a.Value()) != string(b.Value()) {
			diff++
		}
	}
	say("RESULT part4 KV cfg: keys legacy %d hub %d; keys whose latest value differs or is missing on hub: %d", len(lk), len(hk), diff)

	// cut-over: the hub stream stops sourcing and takes over the subjects
	tc := time.Now()
	subj := map[string][]string{"S": {"mesh.s.>"}, "KV_cfg": {"$KV.cfg.>"}}
	for n, ss := range subj {
		h := must(hjs.Stream(ctx, n))
		cfg := h.CachedInfo().Config
		cfg.Sources = nil
		cfg.Subjects = ss
		if _, err := hjs.UpdateStream(ctx, cfg); err != nil {
			say("cut-over %s FAILED: %v", n, err)
		}
	}
	say("cut-over update (sources removed, subjects added, 2 streams): %s", time.Since(tc).Round(time.Millisecond))
	before := msgs(hjs, "S")
	if _, err := hjs.Publish(ctx, "mesh.s.1", []byte("after-cutover")); err != nil {
		say("RESULT part4 publish on hub after cut-over FAILED: %v", err)
	} else {
		say("RESULT part4 after cut-over a publish on the hub lands in the hub stream: %d -> %d messages (legacy stays %d)", before, msgs(hjs, "S"), msgs(ljs, "S"))
	}
	if _, err := hkv.Put(ctx, "after", []byte("1")); err != nil {
		say("RESULT part4 KV put on hub after cut-over FAILED: %v", err)
	} else {
		g, err := hkv.Get(ctx, "after")
		say("RESULT part4 KV put on hub after cut-over: readable %v", err == nil && string(g.Value()) == "1")
	}
}