package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/nats-io/jwt/v2"
	"github.com/nats-io/nkeys"
)

// Keys is what the driver on machine A needs to act as the operator, the identity seat and a system user.
type Keys struct {
	Op, Sys, Acc, Sig string // seeds
	A, B, C           string // addresses
}

func must[T any](v T, err error) T {
	if err != nil {
		panic(err)
	}
	return v
}

func pubOf(kp nkeys.KeyPair) string { return must(kp.PublicKey()) }

func fromSeed(s string) nkeys.KeyPair { return must(nkeys.FromSeed([]byte(s))) }

func seedOf(kp nkeys.KeyPair) string { return string(must(kp.Seed())) }

// accountJWT is the MESH account: JetStream limits, one signing key, a revocation list.
func accountJWT(op, acc, sig nkeys.KeyPair, revoked map[string]time.Time) string {
	ac := jwt.NewAccountClaims(pubOf(acc))
	ac.Name = "MESH"
	ac.Limits.JetStreamLimits = jwt.JetStreamLimits{MemoryStorage: 64 << 20, DiskStorage: 1 << 30, Streams: 20, Consumer: 50}
	ac.SigningKeys.Add(pubOf(sig))
	if len(revoked) > 0 {
		ac.Revocations = jwt.RevocationList{}
		for k, t := range revoked {
			ac.Revocations.Revoke(k, t)
		}
	}
	return must(ac.Encode(op))
}

// userJWT: a user of the MESH account signed by the account's signing key, allowed everything.
func userJWT(acc, sig, ukp nkeys.KeyPair, name string) string {
	uc := jwt.NewUserClaims(pubOf(ukp))
	uc.Name = name
	uc.IssuerAccount = pubOf(acc)
	uc.Pub.Allow.Add(">")
	uc.Sub.Allow.Add(">")
	return must(uc.Encode(sig))
}

func creds(token string, ukp nkeys.KeyPair) []byte {
	return must(jwt.FormatUserConfig(token, []byte(seedOf(ukp))))
}

func gen(args []string) {
	fs := flag.NewFlagSet("gen", flag.ExitOnError)
	out := fs.String("out", "", "output directory")
	a := fs.String("a", "", "address of A")
	b := fs.String("b", "", "address of B")
	c := fs.String("c", "", "address of C")
	fs.Parse(args)
	os.MkdirAll(*out, 0o755)
	op := must(nkeys.CreateOperator())
	sys := must(nkeys.CreateAccount())
	acc := must(nkeys.CreateAccount())
	sig := must(nkeys.CreateAccount())
	oc := jwt.NewOperatorClaims(pubOf(op))
	oc.Name = "op"
	oc.SystemAccount = pubOf(sys)
	opJWT := must(oc.Encode(op))
	sc := jwt.NewAccountClaims(pubOf(sys))
	sc.Name = "SYS"
	sysJWT := must(sc.Encode(op))
	accJ := accountJWT(op, acc, sig, nil)

	keys := Keys{seedOf(op), seedOf(sys), seedOf(acc), seedOf(sig), *a, *b, *c}
	kj, _ := json.MarshalIndent(keys, "", " ")
	w := func(name string, data []byte) { must(0, os.WriteFile(filepath.Join(*out, name), data, 0o600)) }
	w("keys.json", kj)
	w("op.jwt", []byte(opJWT))

	// the legacy server joins as a leaf with one user in each of the two accounts
	lk := must(nkeys.CreateUser())
	w("leaf-mesh.creds", creds(userJWT(acc, sig, lk, "legacy-leaf"), lk))
	sk := must(nkeys.CreateUser())
	suc := jwt.NewUserClaims(pubOf(sk))
	suc.Name = "legacy-leaf-sys"
	w("leaf-sys.creds", creds(must(suc.Encode(sys)), sk))

	common := func(name, domain string) string {
		return fmt.Sprintf(`listen: 0.0.0.0:4222
server_name: %s
http_port: 8222
ping_interval: "2s"
ping_max: 2
log_file: "/var/log/q8-%s.log"
jetstream: {store_dir: "/var/lib/q8/%s/js", domain: %s}
`, name, name, name, domain)
	}
	hub := func(name, peer string) string {
		return common(name, "hub") + fmt.Sprintf(`operator: "/root/q8/op.jwt"
system_account: %s
resolver: {type: full, dir: "/var/lib/q8/%s/accts", allow_delete: false}
resolver_preload: {%s: %q, %s: %q}
cluster { name: HUB, listen: 0.0.0.0:6222, routes: ["nats-route://%s:6222"] }
leafnodes { listen: 0.0.0.0:7422 }
`, pubOf(sys), name, pubOf(sys), sysJWT, pubOf(acc), accJ, peer)
	}
	w("B.conf", []byte(hub("B", *c)))
	w("C.conf", []byte(hub("C", *b)))
	w("A.conf", []byte(common("A", "legacy")+fmt.Sprintf(`accounts {
  MESH { jetstream: enabled
    users: [ {user: w, password: pw} ] }
  SYS { users: [ {user: sys, password: sys} ] }
}
system_account: SYS
leafnodes { remotes [
  {url: "nats-leaf://%s:7422", account: MESH, credentials: "/root/q8/leaf-mesh.creds"},
  {url: "nats-leaf://%s:7422", account: SYS, credentials: "/root/q8/leaf-sys.creds"}
] }
`, *b, *b)))
}

func userClaimsFor(sys, ukp nkeys.KeyPair) string {
	uc := jwt.NewUserClaims(pubOf(ukp))
	uc.Name = "sys"
	return must(uc.Encode(sys))
}