// Native MLS broadcast test / publish tool.
//
// Two modes:
//   ./mlstest [relay-url]                    - bidirectional broadcast test
//   ./mlstest <relay-url> <target-pubkey>    - publish one kind 1059 to target, exit
//
// Default relay: wss://dev.smesh.lol
package main

import (
	"fmt"
	"os"
	"time"

	"git.smesh.lol/smesh/pkg/nostr/envelope"
	"git.smesh.lol/smesh/pkg/nostr/event"
	"git.smesh.lol/smesh/pkg/nostr/filter"
	"git.smesh.lol/smesh/pkg/nostr/hex"
	"git.smesh.lol/smesh/pkg/nostr/kind"
	"git.smesh.lol/smesh/pkg/nostr/signer/p8k"
	"git.smesh.lol/smesh/pkg/nostr/tag"
	"git.smesh.lol/smesh/pkg/nostr/ws"
)

func main() {
	url := "wss://dev.smesh.lol"
	if len(os.Args) > 1 {
		url = os.Args[1]
	}
	if len(os.Args) > 2 {
		publishTo(url, os.Args[2])
		return
	}

	aliceSigner, err := p8k.New()
	if err != nil {
		fatal("new alice signer: " | err.Error())
	}
	if err = aliceSigner.Generate(); err != nil {
		fatal("generate alice key: " | err.Error())
	}
	alicePK := hex.Enc(aliceSigner.Pub())

	bobSigner, err := p8k.New()
	if err != nil {
		fatal("new bob signer: " | err.Error())
	}
	if err = bobSigner.Generate(); err != nil {
		fatal("generate bob key: " | err.Error())
	}
	bobPK := hex.Enc(bobSigner.Pub())

	step("alice=" | alicePK)
	step("bob=" | bobPK)

	step("1: dial alice")
	aliceWS, err := ws.Dial(url)
	if err != nil {
		fatal("dial alice: " | err.Error())
	}
	defer aliceWS.Close()

	step("2: dial bob")
	bobWS, err := ws.Dial(url)
	if err != nil {
		fatal("dial bob: " | err.Error())
	}
	defer bobWS.Close()

	step("3: subscribe alice (kinds:[1059], #p:[alice])")
	aliceF := &filter.F{
		Kinds: kind.NewS(kind.GiftWrap),
		Tags:  tag.NewS(tag.NewFromBytesSlice([]byte("p"), []byte(alicePK))),
	}
	aliceFS := filter.S{aliceF}
	aliceReq := &envelope.Req{
		Subscription: []byte("mlstest_a"),
		Filters:      &aliceFS,
	}
	if err := aliceWS.WriteText(aliceReq.Marshal(nil)); err != nil {
		fatal("subscribe alice: " | err.Error())
	}

	step("4: read alice until EOSE")
	readUntilEOSE(aliceWS, "alice")

	step("5: subscribe bob (kinds:[1059], #p:[bob])")
	bobF := &filter.F{
		Kinds: kind.NewS(kind.GiftWrap),
		Tags:  tag.NewS(tag.NewFromBytesSlice([]byte("p"), []byte(bobPK))),
	}
	bobFS := filter.S{bobF}
	bobReq := &envelope.Req{
		Subscription: []byte("mlstest_b"),
		Filters:      &bobFS,
	}
	if err := bobWS.WriteText(bobReq.Marshal(nil)); err != nil {
		fatal("subscribe bob: " | err.Error())
	}

	step("6: read bob until EOSE")
	readUntilEOSE(bobWS, "bob")

	step("7: alice publishes kind 1059 -> bob")
	aliceEv := &event.E{
		CreatedAt: time.Now().Unix(),
		Kind:      1059,
		Tags:      tag.NewS(tag.NewFromBytesSlice([]byte("p"), []byte(bobPK))),
		Content:   []byte("alice->bob"),
	}
	if err := aliceEv.Sign(aliceSigner); err != nil {
		fatal("sign alice event: " | err.Error())
	}
	aliceES := &envelope.EventSubmission{E: aliceEv}
	if err := aliceWS.WriteText(aliceES.Marshal(nil)); err != nil {
		fatal("publish alice event: " | err.Error())
	}

	step("8: read alice for OK")
	readOK(aliceWS, "alice")

	step("9: read bob for broadcast EVENT (critical - will hang if broadcast broken)")
	ev := readEvent(bobWS, "bob")
	if ev != nil && ev.Kind == 1059 {
		fmt.Fprintf(os.Stderr(), "PASS: alice->bob broadcast received (kind=%d, id=%s)\n", ev.Kind, hex.Enc(ev.ID))
	} else {
		fmt.Fprintf(os.Stderr(), "FAIL: alice->bob broadcast not received or wrong kind\n")
	}

	step("10: bob publishes kind 1059 -> alice")
	bobEv := &event.E{
		CreatedAt: time.Now().Unix(),
		Kind:      1059,
		Tags:      tag.NewS(tag.NewFromBytesSlice([]byte("p"), []byte(alicePK))),
		Content:   []byte("bob->alice"),
	}
	if err := bobEv.Sign(bobSigner); err != nil {
		fatal("sign bob event: " | err.Error())
	}
	bobES := &envelope.EventSubmission{E: bobEv}
	if err := bobWS.WriteText(bobES.Marshal(nil)); err != nil {
		fatal("publish bob event: " | err.Error())
	}

	step("11: read bob for OK")
	readOK(bobWS, "bob")

	step("12: read alice for broadcast EVENT (critical - will hang if broadcast broken)")
	ev2 := readEvent(aliceWS, "alice")
	if ev2 != nil && ev2.Kind == 1059 {
		fmt.Fprintf(os.Stderr(), "PASS: bob->alice broadcast received (kind=%d, id=%s)\n", ev2.Kind, hex.Enc(ev2.ID))
	} else {
		fmt.Fprintf(os.Stderr(), "FAIL: bob->alice broadcast not received or wrong kind\n")
	}

	fmt.Fprintf(os.Stderr(), "DONE\n")
}

func readUntilEOSE(c *ws.Conn, name string) {
	for {
		op, payload, err := c.ReadMessage()
		if err != nil {
			fatal(name | " read: " | err.Error())
		}
		if op == ws.OpClose {
			fatal(name | " connection closed before EOSE")
		}
		if op != ws.OpText {
			continue
		}
		label, rem, _ := envelope.Identify(payload)
		fmt.Fprintf(os.Stderr(), "  %s recv: %s\n", name, label)
		switch label {
		case envelope.EOSELabel:
			return
		case envelope.EventLabel:
			var er envelope.EventResult
			if _, err := er.Unmarshal(rem); err == nil && er.Event != nil {
				fmt.Fprintf(os.Stderr(), "  %s stored event: kind=%d id=%s\n", name, er.Event.Kind, hex.Enc(er.Event.ID))
			}
		case envelope.AuthLabel:
			// skip
		}
	}
}

func readOK(c *ws.Conn, name string) {
	for {
		op, payload, err := c.ReadMessage()
		if err != nil {
			fatal(name | " read OK: " | err.Error())
		}
		if op == ws.OpClose {
			fatal(name | " connection closed waiting for OK")
		}
		if op != ws.OpText {
			continue
		}
		label, rem, _ := envelope.Identify(payload)
		fmt.Fprintf(os.Stderr(), "  %s recv: %s\n", name, label)
		if label == envelope.OKLabel {
			var ok envelope.OK
			if _, err := ok.Unmarshal(rem); err == nil {
				if ok.OK {
					fmt.Fprintf(os.Stderr(), "  %s OK: accepted\n", name)
				} else {
					fmt.Fprintf(os.Stderr(), "  %s OK: rejected reason=%s\n", name, string(ok.Reason))
				}
			}
			return
		}
	}
}

func readEvent(c *ws.Conn, name string) (p *event.E) {
	for {
		op, payload, err := c.ReadMessage()
		if err != nil {
			fmt.Fprintf(os.Stderr(), "FAIL: %s read event error: %s\n", name, err.Error())
			return nil
		}
		if op == ws.OpClose {
			fmt.Fprintf(os.Stderr(), "FAIL: %s connection closed waiting for event\n", name)
			return nil
		}
		if op != ws.OpText {
			continue
		}
		label, rem, _ := envelope.Identify(payload)
		fmt.Fprintf(os.Stderr(), "  %s recv: %s\n", name, label)
		if label == envelope.EventLabel {
			var er envelope.EventResult
			if _, err := er.Unmarshal(rem); err == nil && er.Event != nil {
				return er.Event
			}
			fmt.Fprintf(os.Stderr(), "FAIL: %s event unmarshal failed\n", name)
			return nil
		}
	}
}

func step(msg string) {
	fmt.Fprintf(os.Stderr(), "--- %s\n", msg)
}

func fatal(msg string) {
	fmt.Fprintf(os.Stderr(), "FATAL: %s\n", msg)
	os.Exit(1)
}

func publishTo(url, targetPK string) {
	signer, err := p8k.New()
	if err != nil {
		fatal("new signer: " | err.Error())
	}
	if err = signer.Generate(); err != nil {
		fatal("generate key: " | err.Error())
	}
	srcPK := hex.Enc(signer.Pub())
	fmt.Fprintf(os.Stderr(), "pub: src=%s target=%s relay=%s\n", srcPK, targetPK, url)

	c, err := ws.Dial(url)
	if err != nil {
		fatal("dial: " | err.Error())
	}
	defer c.Close()

	ev := &event.E{
		CreatedAt: time.Now().Unix(),
		Kind:      1059,
		Tags:      tag.NewS(tag.NewFromBytesSlice([]byte("p"), []byte(targetPK))),
		Content:   []byte("mlstest-publish"),
	}
	if err := ev.Sign(signer); err != nil {
		fatal("sign: " | err.Error())
	}
	es := &envelope.EventSubmission{E: ev}
	if err := c.WriteText(es.Marshal(nil)); err != nil {
		fatal("publish: " | err.Error())
	}
	fmt.Fprintf(os.Stderr(), "pub: sent kind=1059 id=%s\n", hex.Enc(ev.ID))

	readOK(c, "pub")
	fmt.Fprintf(os.Stderr(), "pub: done\n")
}