package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"time"

	"github.com/playwright-community/playwright-go"
)

const (
	defaultRoomURL = "https://app.kosmi.io/room/@hyperspaceout"
	logFile        = "auth-monitor.log"
)

func main() {
	roomURL := flag.String("room", defaultRoomURL, "Kosmi room URL")
	loginMode := flag.Bool("login", false, "Monitor login flow (navigate to login page first)")
	testReconnect := flag.Bool("reconnect", false, "Test reconnection behavior (will pause network)")
	flag.Parse()

	log.Println("🔍 Starting Kosmi Auth & Reconnection Monitor")
	log.Printf("📡 Room URL: %s", *roomURL)
	log.Printf("🔐 Login mode: %v", *loginMode)
	log.Printf("🔄 Reconnect test: %v", *testReconnect)
	log.Printf("📝 Log file: %s", logFile)
	log.Println()

	// Create log file
	f, err := os.Create(logFile)
	if err != nil {
		log.Fatalf("Failed to create log file: %v", err)
	}
	defer f.Close()

	// Helper to log to both console and file
	logBoth := func(format string, args ...interface{}) {
		msg := fmt.Sprintf(format, args...)
		log.Println(msg)
		fmt.Fprintf(f, "%s %s\n", time.Now().Format("15:04:05.000"), msg)
	}

	// Set up interrupt handler
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	// Launch Playwright
	pw, err := playwright.Run()
	if err != nil {
		log.Fatalf("Failed to start Playwright: %v", err)
	}

	// Cleanup function
	cleanup := func() {
		logBoth("\n👋 Shutting down...")
		if pw != nil {
			pw.Stop()
		}
		os.Exit(0)
	}

	// Handle interrupt
	go func() {
		<-interrupt
		cleanup()
	}()

	// Launch browser (visible so we can interact for login)
	browser, err := pw.Chromium.Launch(playwright.BrowserTypeLaunchOptions{
		Headless: playwright.Bool(false),
	})
	if err != nil {
		log.Fatalf("Failed to launch browser: %v", err)
	}

	// Create context
	context, err := browser.NewContext(playwright.BrowserNewContextOptions{
		UserAgent: playwright.String("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"),
	})
	if err != nil {
		log.Fatalf("Failed to create context: %v", err)
	}

	// Monitor all network requests
	context.On("request", func(request playwright.Request) {
		url := request.URL()
		method := request.Method()
		
		// Log all Kosmi-related requests
		if containsKosmi(url) {
			logBoth("🌐 [HTTP REQUEST] %s %s", method, url)
			
			// Log POST data (likely contains login credentials or GraphQL mutations)
			if method == "POST" {
				postData, err := request.PostData()
				if err == nil && postData != "" {
					logBoth("   📤 POST Data: %s", postData)
				}
			}
			
			// Log headers for auth-related requests
			headers := request.Headers()
			if authHeader, ok := headers["authorization"]; ok {
				logBoth("   🔑 Authorization: %s", authHeader)
			}
			if cookie, ok := headers["cookie"]; ok {
				logBoth("   🍪 Cookie: %s", truncate(cookie, 100))
			}
		}
	})

	// Monitor all network responses
	context.On("response", func(response playwright.Response) {
		url := response.URL()
		status := response.Status()
		
		if containsKosmi(url) {
			logBoth("📨 [HTTP RESPONSE] %d %s", status, url)
			
			// Try to get response body for auth endpoints
			if status >= 200 && status < 300 {
				body, err := response.Body()
				if err == nil && len(body) > 0 && len(body) < 50000 {
					// Try to parse as JSON
					var jsonData interface{}
					if json.Unmarshal(body, &jsonData) == nil {
						prettyJSON, _ := json.MarshalIndent(jsonData, "   ", "  ")
						logBoth("   📦 Response Body: %s", string(prettyJSON))
					}
				}
			}
			
			// Log Set-Cookie headers
			headers := response.Headers()
			if setCookie, ok := headers["set-cookie"]; ok {
				logBoth("   🍪 Set-Cookie: %s", setCookie)
			}
		}
	})

	// Create page
	page, err := context.NewPage()
	if err != nil {
		log.Fatalf("Failed to create page: %v", err)
	}

	// Inject comprehensive monitoring script BEFORE navigation
	logBoth("📝 Injecting WebSocket, storage, and reconnection monitoring script...")
	monitorScript := getMonitoringScript()
	logBoth("   Script length: %d characters", len(monitorScript))
	if err := page.AddInitScript(playwright.Script{
		Content: playwright.String(monitorScript),
	}); err != nil {
		log.Fatalf("Failed to inject script: %v", err)
	}
	logBoth("   ✅ Script injected successfully")

	// Listen to console messages
	page.On("console", func(msg playwright.ConsoleMessage) {
		text := msg.Text()
		msgType := msg.Type()

		// Format the output nicely
		prefix := "💬"
		switch msgType {
		case "log":
			prefix = "📋"
		case "error":
			prefix = "❌"
		case "warning":
			prefix = "⚠️"
		case "info":
			prefix = "ℹ️"
		}

		logBoth("%s [BROWSER %s] %s", prefix, msgType, text)
	})

	// Navigate based on mode
	if *loginMode {
		logBoth("🔐 Login mode: Navigate to login page first")
		logBoth("   Please log in manually in the browser")
		logBoth("   We'll capture all auth traffic")
		
		// Navigate to main Kosmi page (will redirect to login if not authenticated)
		logBoth("🌐 Navigating to https://app.kosmi.io...")
		timeout := 30000.0 // 30 seconds
		_, err := page.Goto("https://app.kosmi.io", playwright.PageGotoOptions{
			WaitUntil: playwright.WaitUntilStateDomcontentloaded,
			Timeout:   &timeout,
		})
		if err != nil {
			logBoth("⚠️  Navigation error: %v", err)
			logBoth("   Continuing anyway - page might still be usable")
		} else {
			logBoth("✅ Page loaded successfully")
		}
		
		logBoth("")
		logBoth("📋 Instructions:")
		logBoth("   1. Log in with your credentials in the browser")
		logBoth("   2. Navigate to the room: %s", *roomURL)
		logBoth("   3. Browser console messages should appear here")
		logBoth("   4. Press Ctrl+C when done")
		logBoth("")
		logBoth("💡 Expected console messages:")
		logBoth("   - '[Monitor] Installing...'")
		logBoth("   - '[WS MONITOR] WebSocket created...'")
		logBoth("   - '[FETCH] Request to...'")
		logBoth("")
		
	} else {
		// Navigate directly to room (anonymous flow)
		logBoth("🌐 Navigating to %s...", *roomURL)
		timeout := 30000.0 // 30 seconds
		_, err := page.Goto(*roomURL, playwright.PageGotoOptions{
			WaitUntil: playwright.WaitUntilStateDomcontentloaded,
			Timeout:   &timeout,
		})
		if err != nil {
			logBoth("⚠️  Navigation error: %v", err)
			logBoth("   Continuing anyway - page might still be usable")
		} else {
			logBoth("✅ Page loaded successfully")
		}
	}

	// Wait for WebSocket to connect
	time.Sleep(5 * time.Second)

	// Capture storage data
	logBoth("\n📦 Capturing storage data...")
	captureStorage(page, logBoth)

	if *testReconnect {
		logBoth("\n🔄 Testing reconnection behavior...")
		logBoth("   This will simulate network disconnection")
		logBoth("   Watch how the browser handles reconnection")
		
		// Simulate network offline
		logBoth("   📡 Setting network to OFFLINE...")
		if err := context.SetOffline(true); err != nil {
			logBoth("   ❌ Failed to set offline: %v", err)
		} else {
			logBoth("   ✅ Network is now OFFLINE")
			logBoth("   ⏳ Waiting 10 seconds...")
			time.Sleep(10 * time.Second)
			
			// Restore network
			logBoth("   📡 Setting network to ONLINE...")
			if err := context.SetOffline(false); err != nil {
				logBoth("   ❌ Failed to set online: %v", err)
			} else {
				logBoth("   ✅ Network is now ONLINE")
				logBoth("   ⏳ Observing reconnection behavior...")
				time.Sleep(10 * time.Second)
			}
		}
	}

	// Capture final storage state
	logBoth("\n📦 Capturing final storage state...")
	captureStorage(page, logBoth)

	// Keep monitoring
	logBoth("\n⏳ Monitoring all traffic... Press Ctrl+C to stop")
	if *loginMode {
		logBoth("💡 TIP: Try logging in, navigating to rooms, and logging out")
		logBoth("      We'll capture all authentication flows")
	}

	// Wait forever (interrupt handler will exit)
	select {}
}

// captureStorage captures localStorage, sessionStorage, and cookies
func captureStorage(page playwright.Page, logBoth func(string, ...interface{})) {
	// Capture localStorage
	localStorageResult, err := page.Evaluate(`
		(function() {
			const data = {};
			for (let i = 0; i < localStorage.length; i++) {
				const key = localStorage.key(i);
				data[key] = localStorage.getItem(key);
			}
			return data;
		})();
	`)
	if err == nil {
		if localStorage, ok := localStorageResult.(map[string]interface{}); ok {
			logBoth("   📦 localStorage:")
			for key, value := range localStorage {
				logBoth("      %s: %s", key, truncate(fmt.Sprintf("%v", value), 100))
			}
		}
	}

	// Capture sessionStorage
	sessionStorageResult, err := page.Evaluate(`
		(function() {
			const data = {};
			for (let i = 0; i < sessionStorage.length; i++) {
				const key = sessionStorage.key(i);
				data[key] = sessionStorage.getItem(key);
			}
			return data;
		})();
	`)
	if err == nil {
		if sessionStorage, ok := sessionStorageResult.(map[string]interface{}); ok {
			logBoth("   📦 sessionStorage:")
			for key, value := range sessionStorage {
				logBoth("      %s: %s", key, truncate(fmt.Sprintf("%v", value), 100))
			}
		}
	}

	// Capture cookies
	cookiesResult, err := page.Evaluate(`
		(function() {
			return document.cookie.split(';').map(c => {
				const parts = c.trim().split('=');
				return {
					name: parts[0],
					value: parts.slice(1).join('=')
				};
			}).filter(c => c.name && c.value);
		})();
	`)
	if err == nil {
		if cookiesArray, ok := cookiesResult.([]interface{}); ok {
			logBoth("   🍪 Cookies:")
			for _, cookieItem := range cookiesArray {
				if cookie, ok := cookieItem.(map[string]interface{}); ok {
					name := cookie["name"]
					value := cookie["value"]
					logBoth("      %s: %s", name, truncate(fmt.Sprintf("%v", value), 100))
				}
			}
		}
	}

	// Capture WebSocket status
	wsStatusResult, err := page.Evaluate(`
		(function() {
			return window.__KOSMI_WS_STATUS__ || { error: "No WebSocket found" };
		})();
	`)
	if err == nil {
		if wsStatus, ok := wsStatusResult.(map[string]interface{}); ok {
			logBoth("   🔌 WebSocket Status:")
			for key, value := range wsStatus {
				logBoth("      %s: %v", key, value)
			}
		}
	}
}

// getMonitoringScript returns the comprehensive monitoring script
func getMonitoringScript() string {
	return `
		(function() {
			if (window.__KOSMI_MONITOR_INSTALLED__) return;
			
			console.log('[Monitor] Installing comprehensive monitoring hooks...');
			
			// Store original WebSocket constructor
			const OriginalWebSocket = window.WebSocket;
			let wsInstance = null;
			let messageCount = 0;
			let reconnectAttempts = 0;
			
			// Hook WebSocket constructor
			window.WebSocket = function(url, protocols) {
				console.log('🔌 [WS MONITOR] WebSocket created:', url, 'protocols:', protocols);
				const socket = new OriginalWebSocket(url, protocols);
				
				if (url.includes('engine.kosmi.io') || url.includes('gql-ws')) {
					wsInstance = socket;
					
					// Track connection lifecycle
					socket.addEventListener('open', (event) => {
						console.log('✅ [WS MONITOR] WebSocket OPENED');
						reconnectAttempts = 0;
						updateWSStatus('OPEN', url);
					});
					
					socket.addEventListener('close', (event) => {
						console.log('🔴 [WS MONITOR] WebSocket CLOSED:', event.code, event.reason, 'wasClean:', event.wasClean);
						reconnectAttempts++;
						updateWSStatus('CLOSED', url, { code: event.code, reason: event.reason, reconnectAttempts });
					});
					
					socket.addEventListener('error', (event) => {
						console.error('❌ [WS MONITOR] WebSocket ERROR:', event);
						updateWSStatus('ERROR', url);
					});
					
					// Intercept outgoing messages
					const originalSend = socket.send;
					socket.send = function(data) {
						messageCount++;
						console.log('📤 [WS MONITOR] SEND #' + messageCount + ':', data);
						try {
							const parsed = JSON.parse(data);
							console.log('   Type:', parsed.type, 'ID:', parsed.id);
							if (parsed.payload) {
								console.log('   Payload:', JSON.stringify(parsed.payload, null, 2));
								
								// Check for auth-related messages
								if (parsed.type === 'connection_init' && parsed.payload.token) {
									console.log('   🔑 [AUTH] Connection init with token:', parsed.payload.token.substring(0, 50) + '...');
								}
							}
						} catch (e) {
							// Not JSON
						}
						return originalSend.call(this, data);
					};
					
					// Intercept incoming messages
					socket.addEventListener('message', (event) => {
						messageCount++;
						console.log('📥 [WS MONITOR] RECEIVE #' + messageCount + ':', event.data);
						try {
							const parsed = JSON.parse(event.data);
							console.log('   Type:', parsed.type, 'ID:', parsed.id);
							if (parsed.payload) {
								console.log('   Payload:', JSON.stringify(parsed.payload, null, 2));
							}
						} catch (e) {
							// Not JSON
						}
					});
				}
				
				return socket;
			};
			
			// Preserve WebSocket properties
			window.WebSocket.prototype = OriginalWebSocket.prototype;
			window.WebSocket.CONNECTING = OriginalWebSocket.CONNECTING;
			window.WebSocket.OPEN = OriginalWebSocket.OPEN;
			window.WebSocket.CLOSING = OriginalWebSocket.CLOSING;
			window.WebSocket.CLOSED = OriginalWebSocket.CLOSED;
			
			// Monitor localStorage changes
			const originalSetItem = Storage.prototype.setItem;
			Storage.prototype.setItem = function(key, value) {
				console.log('💾 [STORAGE] localStorage.setItem:', key, '=', value.substring(0, 100));
				if (key.toLowerCase().includes('token') || key.toLowerCase().includes('auth')) {
					console.log('   🔑 [AUTH] Auth-related storage detected!');
				}
				return originalSetItem.call(this, key, value);
			};
			
			// Monitor sessionStorage changes
			const originalSessionSetItem = sessionStorage.setItem;
			sessionStorage.setItem = function(key, value) {
				console.log('💾 [STORAGE] sessionStorage.setItem:', key, '=', value.substring(0, 100));
				if (key.toLowerCase().includes('token') || key.toLowerCase().includes('auth')) {
					console.log('   🔑 [AUTH] Auth-related storage detected!');
				}
				return originalSessionSetItem.call(this, key, value);
			};
			
			// Monitor fetch requests (for GraphQL mutations)
			const originalFetch = window.fetch;
			window.fetch = function(url, options) {
				if (typeof url === 'string' && url.includes('kosmi.io')) {
					console.log('🌐 [FETCH] Request to:', url);
					if (options && options.body) {
						console.log('   📤 Body:', options.body);
						try {
							const parsed = JSON.parse(options.body);
							if (parsed.query) {
								console.log('   📝 GraphQL Query:', parsed.query.substring(0, 200));
								if (parsed.query.includes('login') || parsed.query.includes('auth')) {
									console.log('   🔑 [AUTH] Auth-related GraphQL detected!');
								}
							}
						} catch (e) {
							// Not JSON
						}
					}
				}
				return originalFetch.apply(this, arguments).then(response => {
					if (typeof url === 'string' && url.includes('kosmi.io')) {
						console.log('📨 [FETCH] Response from:', url, 'status:', response.status);
					}
					return response;
				});
			};
			
			// Helper to update WebSocket status
			function updateWSStatus(state, url, extra = {}) {
				window.__KOSMI_WS_STATUS__ = {
					state,
					url,
					timestamp: new Date().toISOString(),
					messageCount,
					...extra
				};
			}
			
			// Monitor network connectivity
			window.addEventListener('online', () => {
				console.log('🌐 [NETWORK] Browser is ONLINE');
			});
			
			window.addEventListener('offline', () => {
				console.log('🌐 [NETWORK] Browser is OFFLINE');
			});
			
			window.__KOSMI_MONITOR_INSTALLED__ = true;
			console.log('[Monitor] ✅ All monitoring hooks installed');
		})();
	`
}

// Helper functions
func containsKosmi(url string) bool {
	return contains(url, "kosmi.io") || contains(url, "engine.kosmi.io") || contains(url, "app.kosmi.io")
}

func contains(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}

func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}