package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/chromedp/chromedp"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run jackbox-player-count.go <ROOM_CODE>")
		os.Exit(1)
	}

	roomCode := strings.ToUpper(os.Args[1])
	count, err := getPlayerCount(roomCode)
	if err != nil {
		fmt.Printf("Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("%d\n", count)
}

func getPlayerCount(roomCode string) (int, error) {
	// Suppress cookie errors
	opts := append(chromedp.DefaultExecAllocatorOptions[:],
		chromedp.Flag("headless", true),
	)
	
	allocCtx, cancel := chromedp.NewExecAllocator(context.Background(), opts...)
	defer cancel()
	
	ctx, cancel := chromedp.NewContext(allocCtx, chromedp.WithLogf(func(s string, i ...interface{}) {
		msg := fmt.Sprintf(s, i...)
		if !strings.Contains(msg, "cookie") {
			log.Printf(msg)
		}
	}))
	defer cancel()

	ctx, cancel = context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	var playerCount int

	err := chromedp.Run(ctx,
		chromedp.Navigate("https://jackbox.tv/"),
		chromedp.WaitVisible(`input[placeholder*="ENTER 4-LETTER CODE"]`),
		
		// Inject WebSocket hook BEFORE joining
		chromedp.ActionFunc(func(ctx context.Context) error {
			return chromedp.Evaluate(`
				window.__jackboxPlayerCount = null;
				const OriginalWebSocket = window.WebSocket;
				window.WebSocket = function(...args) {
					const ws = new OriginalWebSocket(...args);
					const originalAddEventListener = ws.addEventListener;
					ws.addEventListener = function(type, listener, ...rest) {
						if (type === 'message') {
							const wrappedListener = function(event) {
								try {
									const data = JSON.parse(event.data);
									if (data.opcode === 'client/welcome' && data.result && data.result.here) {
										window.__jackboxPlayerCount = Object.keys(data.result.here).length;
									}
								} catch (e) {}
								return listener.call(this, event);
							};
							return originalAddEventListener.call(this, type, wrappedListener, ...rest);
						}
						return originalAddEventListener.call(this, type, listener, ...rest);
					};
					return ws;
				};
			`, nil).Do(ctx)
		}),
		
		// Now join
		chromedp.SendKeys(`input[placeholder*="ENTER 4-LETTER CODE"]`, roomCode+"\n\n"),
		
		// Poll for the value
		chromedp.ActionFunc(func(ctx context.Context) error {
			for i := 0; i < 60; i++ { // Try for 30 seconds
				time.Sleep(500 * time.Millisecond)
				var count int
				err := chromedp.Evaluate(`window.__jackboxPlayerCount || -1`, &count).Do(ctx)
				if err == nil && count > 0 {
					playerCount = count
					return nil
				}
			}
			return fmt.Errorf("timeout waiting for player count")
		}),
	)

	if err != nil {
		return 0, err
	}

	if playerCount < 0 {
		return 0, fmt.Errorf("could not find player count")
	}

	return playerCount, nil
}