// Command chat is a demo peer with two modes.
//
// Simple mode (default): joins a room subject and prints every received message
// live. It validates uniformity worker<->chat: a process publishes, a human UI
// subscribes, both speak the same protocol.
//
// Encrypted demo (--demo-encrypted): a self-contained script against a running
// membershipd that proves E2E encryption and forward secrecy. A creates an
// encrypted room, invites B, A publishes a message B can decrypt, then A kicks
// B and publishes again at the new epoch — B can no longer decrypt. Prints a
// PASS/FAIL summary of each assertion.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/enmanuel/unibus/pkg/client"
	"github.com/enmanuel/unibus/pkg/frame"
	"github.com/enmanuel/unibus/pkg/room"
)

func main() {
	var (
		natsURL  = flag.String("nats-url", "nats://127.0.0.1:4222", "NATS url")
		ctrlURL  = flag.String("ctrl-url", "http://127.0.0.1:8420", "membershipd control-plane url")
		roomSub  = flag.String("room", "proc.test.ticks", "room subject to subscribe to")
		idFile   = flag.String("id-file", "./local_files/chat.id", "identity file path")
		demoEnc  = flag.Bool("demo-encrypted", false, "run the encrypted forward-secrecy demo")
	)
	flag.Parse()

	log.SetFlags(log.LstdFlags | log.Lmsgprefix)
	log.SetPrefix("[chat] ")

	if *demoEnc {
		runEncryptedDemo(*natsURL, *ctrlURL)
		return
	}
	runSimple(*natsURL, *ctrlURL, *roomSub, *idFile)
}

// runSimple subscribes to a cleartext subject and prints messages live.
func runSimple(natsURL, ctrlURL, roomSub, idFile string) {
	id, err := client.LoadOrCreateIdentity(idFile)
	if err != nil {
		log.Fatalf("identity: %v", err)
	}
	c, err := client.New(natsURL, ctrlURL, id)
	if err != nil {
		log.Fatalf("connect: %v", err)
	}
	defer c.Close()
	log.Printf("endpoint: %s", c.Endpoint().ID)

	// A subscriber needs a room to resolve the subject + policy. For the
	// cleartext demo each peer owns a room mapped to the shared subject; NATS
	// fans out by subject so worker publishes reach this subscription.
	roomID, err := c.CreateRoom(roomSub, room.ModeNATS)
	if err != nil {
		log.Fatalf("create room: %v", err)
	}
	if err := c.Join(roomID); err != nil {
		log.Fatalf("join: %v", err)
	}
	sub, err := c.Subscribe(roomID, func(f frame.Frame, plaintext []byte) {
		fmt.Printf("[%s] %s: %s\n", f.Subject, shortID(f.Sender), string(plaintext))
	})
	if err != nil {
		log.Fatalf("subscribe: %v", err)
	}
	defer sub.Unsubscribe()
	log.Printf("subscribed to %q; waiting for messages (Ctrl-C to stop)", roomSub)

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
	<-stop
	log.Printf("bye")
}

func shortID(id string) string {
	if len(id) > 8 {
		return id[:8]
	}
	return id
}

// runEncryptedDemo proves E2E encryption + forward secrecy end-to-end.
func runEncryptedDemo(natsURL, ctrlURL string) {
	log.Printf("=== encrypted forward-secrecy demo ===")
	pass := true
	check := func(name string, ok bool) {
		status := "PASS"
		if !ok {
			status = "FAIL"
			pass = false
		}
		fmt.Printf("  [%s] %s\n", status, name)
	}

	// Two identities: A (owner) and B (invitee). In-memory only (demo).
	idA, err := newEphemeralIdentity()
	must(err, "generate A identity")
	idB, err := newEphemeralIdentity()
	must(err, "generate B identity")

	a, err := client.New(natsURL, ctrlURL, idA)
	must(err, "connect A")
	defer a.Close()
	b, err := client.New(natsURL, ctrlURL, idB)
	must(err, "connect B")
	defer b.Close()

	// A creates an encrypted room.
	roomID, err := a.CreateRoom("room.test", room.ModeMatrix)
	must(err, "A create room")
	fmt.Printf("  room.test -> %s (E2E, persisted, signed)\n", roomID)

	// A invites B (seals K to B's X25519 key).
	must(a.Invite(roomID, b.Endpoint()), "A invite B")

	// B joins (fetches + decrypts K).
	must(b.Join(roomID), "B join")

	// B subscribes; capture received plaintexts.
	recv := make(chan string, 4)
	subB, err := b.Subscribe(roomID, func(f frame.Frame, plaintext []byte) {
		recv <- string(plaintext)
	})
	must(err, "B subscribe")
	defer subB.Unsubscribe()
	time.Sleep(200 * time.Millisecond) // let the subscription settle

	// A publishes a message B can decrypt.
	const msg1 = "hola E2E"
	must(a.Publish(roomID, []byte(msg1)), "A publish msg1")
	got1, ok := waitMsg(recv, 2*time.Second)
	check("B decrypts pre-kick message", ok && got1 == msg1)

	// A kicks B (rotates K to a new epoch, re-sealed only for the remaining members).
	must(a.Kick(roomID, b.Endpoint().ID), "A kick B")
	time.Sleep(200 * time.Millisecond)

	// A publishes at the new epoch. B must NOT be able to decrypt: it was removed
	// from the member list, so /key returns no key for B at the new epoch.
	const msg2 = "secreto post-kick"
	must(a.Publish(roomID, []byte(msg2)), "A publish msg2 (post-kick)")
	got2, ok2 := waitMsg(recv, 1500*time.Millisecond)
	// Forward secrecy holds if B did NOT receive the post-kick plaintext.
	check("B cannot decrypt post-kick message (forward secrecy)", !(ok2 && got2 == msg2))
	if ok2 {
		fmt.Printf("  (unexpected: B received %q after kick)\n", got2)
	} else {
		fmt.Printf("  (B received nothing decryptable after kick — correct)\n")
	}

	fmt.Println()
	if pass {
		fmt.Println("RESULT: PASS — E2E encryption and forward secrecy verified")
		os.Exit(0)
	}
	fmt.Println("RESULT: FAIL — see assertions above")
	os.Exit(1)
}

func waitMsg(ch <-chan string, timeout time.Duration) (string, bool) {
	select {
	case m := <-ch:
		return m, true
	case <-time.After(timeout):
		return "", false
	}
}

func must(err error, what string) {
	if err != nil {
		log.Fatalf("%s: %v", what, err)
	}
}