package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/chromedp/cdproto/network"
	"github.com/chromedp/chromedp"
)

// PlayerInfo represents a player in the lobby
type PlayerInfo struct {
	ID   string `json:"id"`
	Role string `json:"role"`
	Name string `json:"name"`
}

// LobbyStatus contains all information about the current lobby
type LobbyStatus struct {
	RoomCode      string       `json:"roomCode"`
	AppTag        string       `json:"appTag"`
	PlayerCount   int          `json:"playerCount"`
	AudienceCount int          `json:"audienceCount"`
	MaxPlayers    int          `json:"maxPlayers"`
	GameState     string       `json:"gameState"`
	LobbyState    string       `json:"lobbyState"`
	Locked        bool         `json:"locked"`
	Full          bool         `json:"full"`
	Players       []PlayerInfo `json:"players"`
}

// WebSocketMessage represents a parsed WebSocket message
type WebSocketMessage struct {
	PC     int                    `json:"pc"`
	Opcode string                 `json:"opcode"`
	Result map[string]interface{} `json:"result"`
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run get-player-count.go <ROOM_CODE>")
		fmt.Println("Example: go run get-player-count.go JYET")
		fmt.Println("\nSet DEBUG=true for verbose output:")
		fmt.Println("DEBUG=true go run get-player-count.go JYET")
		os.Exit(1)
	}

	roomCode := strings.ToUpper(strings.TrimSpace(os.Args[1]))

	if len(roomCode) != 4 {
		fmt.Println("Error: Room code must be exactly 4 characters")
		os.Exit(1)
	}

	fmt.Printf("🎮 Jackbox Player Count Fetcher\n")
	fmt.Printf("Room Code: %s\n\n", roomCode)

	status, err := getPlayerCount(roomCode)
	if err != nil {
		log.Fatalf("Error: %v\n", err)
	}

	printStatus(status)
}

func getPlayerCount(roomCode string) (*LobbyStatus, error) {
	// Create chrome context with less verbose logging
	opts := append(chromedp.DefaultExecAllocatorOptions[:],
		chromedp.Flag("headless", true),
		chromedp.Flag("disable-gpu", true),
		chromedp.Flag("no-sandbox", true),
		chromedp.Flag("disable-web-security", true),
	)

	allocCtx, cancel := chromedp.NewExecAllocator(context.Background(), opts...)
	defer cancel()

	// Create context without default logging to reduce cookie errors
	ctx, cancel := chromedp.NewContext(allocCtx, chromedp.WithLogf(func(s string, i ...interface{}) {
		// Only log non-cookie errors
		msg := fmt.Sprintf(s, i...)
		if !strings.Contains(msg, "cookiePart") && !strings.Contains(msg, "could not unmarshal") {
			log.Printf(msg)
		}
	}))
	defer cancel()

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

	var lobbyStatus *LobbyStatus
	welcomeMessageFound := false
	wsMessages := make([]string, 0)

	// Listen for WebSocket frames - this is the most reliable method
	debugMode := os.Getenv("DEBUG") == "true"
	
	chromedp.ListenTarget(ctx, func(ev interface{}) {
		if debugMode {
			fmt.Printf("[DEBUG] Event type: %T\n", ev)
		}
		
		switch ev := ev.(type) {
		case *network.EventWebSocketCreated:
			if debugMode {
				fmt.Printf("[DEBUG] WebSocket Created: %s\n", ev.URL)
			}
		case *network.EventWebSocketFrameReceived:
			// Capture all WebSocket frames
			wsMessages = append(wsMessages, ev.Response.PayloadData)
			
			if debugMode {
				fmt.Printf("[DEBUG] WS Frame received (%d bytes)\n", len(ev.Response.PayloadData))
				if len(ev.Response.PayloadData) < 200 {
					fmt.Printf("[DEBUG] Data: %s\n", ev.Response.PayloadData)
				} else {
					fmt.Printf("[DEBUG] Data (truncated): %s...\n", ev.Response.PayloadData[:200])
				}
			}
			
			// Try to parse immediately
			var wsMsg WebSocketMessage
			if err := json.Unmarshal([]byte(ev.Response.PayloadData), &wsMsg); err == nil {
				if debugMode {
					fmt.Printf("[DEBUG] Parsed opcode: %s\n", wsMsg.Opcode)
				}
				
				if wsMsg.Opcode == "client/welcome" && wsMsg.Result != nil {
					lobbyStatus = parseWelcomeMessage(&wsMsg)
					welcomeMessageFound = true
					fmt.Println("✓ Captured lobby data from WebSocket")
				}
			} else if debugMode {
				fmt.Printf("[DEBUG] Failed to parse JSON: %v\n", err)
			}
		case *network.EventWebSocketFrameSent:
			if debugMode {
				fmt.Printf("[DEBUG] WS Frame sent: %s\n", ev.Response.PayloadData)
			}
		}
	})

	fmt.Println("⏳ Navigating to jackbox.tv...")

	// Enable network tracking BEFORE navigation
	if err := chromedp.Run(ctx, network.Enable()); err != nil {
		return nil, fmt.Errorf("failed to enable network tracking: %w", err)
	}

	err := chromedp.Run(ctx,
		chromedp.Navigate("https://jackbox.tv/"),
		chromedp.WaitVisible(`input[placeholder*="ENTER 4-LETTER CODE"]`, chromedp.ByQuery),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to load jackbox.tv: %w", err)
	}

	fmt.Printf("✓ Loaded jackbox.tv\n")
	fmt.Printf("⏳ Joining room %s...\n", roomCode)

	// Type room code and press Enter to join
	err = chromedp.Run(ctx,
		chromedp.Focus(`input[placeholder*="ENTER 4-LETTER CODE"]`, chromedp.ByQuery),
		chromedp.SendKeys(`input[placeholder*="ENTER 4-LETTER CODE"]`, roomCode+"\n", chromedp.ByQuery),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to enter room code: %w", err)
	}

	if debugMode {
		fmt.Println("[DEBUG] Entered room code and pressed Enter")
	}

	// Wait for room code validation and page transition
	time.Sleep(2 * time.Second)

	fmt.Println("✓ Clicked Play button, waiting for WebSocket data...")

	// Check if we successfully joined (look for typical lobby UI elements)
	time.Sleep(2 * time.Second)
	
	var pageText string
	err = chromedp.Run(ctx,
		chromedp.Text("body", &pageText, chromedp.ByQuery),
	)
	if err == nil && debugMode {
		if strings.Contains(pageText, "Sit back") || strings.Contains(pageText, "waiting") {
			fmt.Println("[DEBUG] Successfully joined lobby (found lobby text)")
		} else {
			fmt.Printf("[DEBUG] Page text: %s\n", pageText[:min(300, len(pageText))])
		}
	}

	// Wait longer for WebSocket to connect and receive welcome message
	for i := 0; i < 15 && !welcomeMessageFound; i++ {
		time.Sleep(500 * time.Millisecond)
		if i%4 == 0 {
			fmt.Printf("⏳ Waiting for lobby data... (%ds)\n", i/2)
		}
	}

	// If we still didn't get it from WebSocket frames, try parsing all captured messages
	if !welcomeMessageFound && len(wsMessages) > 0 {
		fmt.Printf("⏳ Parsing %d captured WebSocket messages...\n", len(wsMessages))
		
		for _, msg := range wsMessages {
			var wsMsg WebSocketMessage
			if err := json.Unmarshal([]byte(msg), &wsMsg); err == nil {
				if wsMsg.Opcode == "client/welcome" && wsMsg.Result != nil {
					lobbyStatus = parseWelcomeMessage(&wsMsg)
					welcomeMessageFound = true
					fmt.Println("✓ Found lobby data in captured messages")
					break
				}
			}
		}
	}

	if lobbyStatus == nil {
		if len(wsMessages) == 0 {
			return nil, fmt.Errorf("no WebSocket messages captured - connection may have failed\nTry running with DEBUG=true for more details")
		}
		
		// Show what we captured
		fmt.Printf("\n⚠️  Captured %d WebSocket messages but couldn't find 'client/welcome'\n", len(wsMessages))
		fmt.Println("\nMessage types found:")
		opcodes := make(map[string]int)
		for _, msg := range wsMessages {
			var wsMsg WebSocketMessage
			if err := json.Unmarshal([]byte(msg), &wsMsg); err == nil {
				opcodes[wsMsg.Opcode]++
			}
		}
		for opcode, count := range opcodes {
			fmt.Printf("  - %s: %d\n", opcode, count)
		}
		
		return nil, fmt.Errorf("could not find player count data in WebSocket messages\nRoom may be invalid, closed, or not in lobby state")
	}

	lobbyStatus.RoomCode = roomCode
	
	// Fetch additional room info from REST API
	if err := enrichWithRestAPI(lobbyStatus); err != nil {
		fmt.Printf("Warning: Could not fetch additional room info: %v\n", err)
	}

	return lobbyStatus, nil
}

func parseWelcomeMessage(msg *WebSocketMessage) *LobbyStatus {
	status := &LobbyStatus{
		Players: []PlayerInfo{},
	}

	// Parse "here" object for players
	if here, ok := msg.Result["here"].(map[string]interface{}); ok {
		status.PlayerCount = len(here)

		for id, playerData := range here {
			if pd, ok := playerData.(map[string]interface{}); ok {
				player := PlayerInfo{ID: id}

				if roles, ok := pd["roles"].(map[string]interface{}); ok {
					if _, hasHost := roles["host"]; hasHost {
						player.Role = "host"
						player.Name = "Host"
					} else if playerRole, ok := roles["player"].(map[string]interface{}); ok {
						player.Role = "player"
						if name, ok := playerRole["name"].(string); ok {
							player.Name = name
						} else {
							player.Name = "Unknown"
						}
					}
				}

				status.Players = append(status.Players, player)
			}
		}
	}

	// Parse entities for additional info
	if entities, ok := msg.Result["entities"].(map[string]interface{}); ok {
		// Audience count
		if audience, ok := entities["audience"].([]interface{}); ok && len(audience) > 1 {
			if audienceData, ok := audience[1].(map[string]interface{}); ok {
				if count, ok := audienceData["count"].(float64); ok {
					status.AudienceCount = int(count)
				}
			}
		}

		// Room state
		if bcRoom, ok := entities["bc:room"].([]interface{}); ok && len(bcRoom) > 1 {
			if roomData, ok := bcRoom[1].(map[string]interface{}); ok {
				if val, ok := roomData["val"].(map[string]interface{}); ok {
					if gameState, ok := val["state"].(string); ok {
						status.GameState = gameState
					}
					if lobbyState, ok := val["lobbyState"].(string); ok {
						status.LobbyState = lobbyState
					}
				}
			}
		}
	}

	return status
}

func enrichWithRestAPI(status *LobbyStatus) error {
	// Fetch additional room info from REST API
	url := fmt.Sprintf("https://ecast.jackboxgames.com/api/v2/rooms/%s", status.RoomCode)
	
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	var result struct {
		OK   bool `json:"ok"`
		Body struct {
			AppTag     string `json:"appTag"`
			MaxPlayers int    `json:"maxPlayers"`
			Locked     bool   `json:"locked"`
			Full       bool   `json:"full"`
		} `json:"body"`
	}

	if err := json.Unmarshal(body, &result); err != nil {
		return err
	}

	if result.OK {
		status.AppTag = result.Body.AppTag
		status.MaxPlayers = result.Body.MaxPlayers
		status.Locked = result.Body.Locked
		status.Full = result.Body.Full
	}

	return nil
}

func printStatus(status *LobbyStatus) {
	fmt.Println()
	fmt.Println("═══════════════════════════════════════════")
	fmt.Println("  Jackbox Room Status")
	fmt.Println("═══════════════════════════════════════════")
	fmt.Println()
	fmt.Printf("Room Code:      %s\n", status.RoomCode)
	fmt.Printf("Game:           %s\n", status.AppTag)
	fmt.Printf("Game State:     %s\n", status.GameState)
	fmt.Printf("Lobby State:    %s\n", status.LobbyState)
	fmt.Printf("Locked:         %t\n", status.Locked)
	fmt.Printf("Full:           %t\n", status.Full)
	fmt.Println()
	fmt.Printf("Players:        %d / %d\n", status.PlayerCount, status.MaxPlayers)
	fmt.Printf("Audience:       %d\n", status.AudienceCount)
	fmt.Println()

	if len(status.Players) > 0 {
		fmt.Println("Current Players:")
		for i, player := range status.Players {
			fmt.Printf("  %d. %s (%s)\n", i+1, player.Name, player.Role)
		}
		fmt.Println()
	}

	fmt.Println("═══════════════════════════════════════════")
	fmt.Println()
}