package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"time"

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

func main() {
	roomURL := flag.String("room", "https://app.kosmi.io/room/@hyperspaceout", "Kosmi room URL")
	output := flag.String("output", "auth-data.json", "Output file for captured data")
	flag.Parse()

	fmt.Printf("Capturing authentication data from: %s\n", *roomURL)
	fmt.Printf("Output will be saved to: %s\n\n", *output)

	// Storage for captured data
	authData := &AuthData{
		RoomURL:       *roomURL,
		CaptureTime:   time.Now(),
		Cookies:       []Cookie{},
		RequestHeaders: map[string]interface{}{},
		ResponseHeaders: map[string]interface{}{},
		WebSocketFrames: []WebSocketFrame{},
		NetworkRequests: []NetworkRequest{},
	}

	// Create Chrome context with network logging
	opts := append(chromedp.DefaultExecAllocatorOptions[:],
		chromedp.Flag("headless", true),
		chromedp.Flag("disable-gpu", false),
		chromedp.Flag("no-sandbox", true),
		chromedp.Flag("disable-dev-shm-usage", true),
		chromedp.Flag("disable-blink-features", "AutomationControlled"),
		chromedp.Flag("disable-infobars", true),
		chromedp.Flag("window-size", "1920,1080"),
		chromedp.UserAgent("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"),
	)

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

	ctx, cancel := chromedp.NewContext(allocCtx)
	defer cancel()

	// Enable network tracking
	chromedp.ListenTarget(ctx, func(ev interface{}) {
		switch ev := ev.(type) {
		case *network.EventRequestWillBeSent:
			if containsKosmiDomain(ev.Request.URL) {
				fmt.Printf("→ REQUEST: %s %s\n", ev.Request.Method, ev.Request.URL)
				authData.NetworkRequests = append(authData.NetworkRequests, NetworkRequest{
					URL:     ev.Request.URL,
					Method:  ev.Request.Method,
					Headers: ev.Request.Headers,
					Time:    time.Now(),
				})
			}

		case *network.EventResponseReceived:
			if containsKosmiDomain(ev.Response.URL) {
				fmt.Printf("← RESPONSE: %d %s\n", ev.Response.Status, ev.Response.URL)
				if ev.Response.Status >= 200 && ev.Response.Status < 300 {
					authData.ResponseHeaders[ev.Response.URL] = ev.Response.Headers
				}
			}

		case *network.EventWebSocketCreated:
			fmt.Printf("🔌 WebSocket Created: %s\n", ev.URL)
			authData.WebSocketURL = ev.URL
			authData.WebSocketRequestID = ev.RequestID.String()

		case *network.EventWebSocketFrameSent:
			data := string(ev.Response.PayloadData)
			fmt.Printf("📤 WS SEND: %s\n", truncate(data, 100))
			authData.WebSocketFrames = append(authData.WebSocketFrames, WebSocketFrame{
				Direction: "sent",
				Data:      data,
				Time:      time.Now(),
			})

		case *network.EventWebSocketFrameReceived:
			data := string(ev.Response.PayloadData)
			fmt.Printf("📥 WS RECV: %s\n", truncate(data, 100))
			authData.WebSocketFrames = append(authData.WebSocketFrames, WebSocketFrame{
				Direction: "received",
				Data:      data,
				Time:      time.Now(),
			})
		}
	})

	// Inject WebSocket hook script
	hookScript := getWebSocketHookScript()

	// Run the capture
	err := chromedp.Run(ctx,
		network.Enable(),
		chromedp.ActionFunc(func(ctx context.Context) error {
			_, err := page.AddScriptToEvaluateOnNewDocument(hookScript).Do(ctx)
			return err
		}),
		chromedp.Navigate(*roomURL),
		chromedp.WaitReady("body"),
	)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error navigating: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("\nWaiting for page to fully load and WebSocket to connect...")
	time.Sleep(5 * time.Second)

	// Capture all cookies
	fmt.Println("\n📋 Capturing cookies...")
	var cookiesData []map[string]interface{}
	script := `
		(function() {
			return document.cookie.split(';').map(c => {
				const parts = c.trim().split('=');
				return {
					name: parts[0],
					value: parts.slice(1).join('='),
					domain: window.location.hostname,
					path: '/'
				};
			}).filter(c => c.name && c.value);
		})();
	`
	if err := chromedp.Run(ctx, chromedp.Evaluate(script, &cookiesData)); err != nil {
		fmt.Fprintf(os.Stderr, "Error capturing cookies: %v\n", err)
	} else {
		for _, c := range cookiesData {
			if name, ok := c["name"].(string); ok {
				if value, ok := c["value"].(string); ok {
					authData.Cookies = append(authData.Cookies, Cookie{
						Name:   name,
						Value:  value,
						Domain: c["domain"].(string),
						Path:   c["path"].(string),
					})
					fmt.Printf("  🍪 %s=%s\n", name, truncate(value, 40))
				}
			}
		}
	}

	// Get CDP cookies (includes HTTPOnly)
	cookies, err := network.GetCookies().Do(ctx)
	if err == nil {
		fmt.Println("\n📋 CDP Cookies (including HTTPOnly):")
		for _, c := range cookies {
			if containsKosmiDomain(c.Domain) {
				authData.HTTPOnlyCookies = append(authData.HTTPOnlyCookies, Cookie{
					Name:     c.Name,
					Value:    c.Value,
					Domain:   c.Domain,
					Path:     c.Path,
					Secure:   c.Secure,
					HTTPOnly: c.HTTPOnly,
					SameSite: string(c.SameSite),
				})
				fmt.Printf("  🔒 %s=%s (HTTPOnly=%v, Secure=%v)\n", 
					c.Name, truncate(c.Value, 40), c.HTTPOnly, c.Secure)
			}
		}
	}

	// Check WebSocket status
	var wsStatus map[string]interface{}
	checkScript := `
		(function() {
			return {
				hookInstalled: !!window.__KOSMI_WS_HOOK_INSTALLED__,
				wsFound: !!window.__KOSMI_WS__,
				wsConnected: window.__KOSMI_WS__ ? window.__KOSMI_WS__.readyState === WebSocket.OPEN : false,
				wsURL: window.__KOSMI_WS__ ? window.__KOSMI_WS__.url : null,
				messageQueueSize: window.__KOSMI_MESSAGE_QUEUE__ ? window.__KOSMI_MESSAGE_QUEUE__.length : 0
			};
		})();
	`
	if err := chromedp.Run(ctx, chromedp.Evaluate(checkScript, &wsStatus)); err == nil {
		authData.WebSocketStatus = wsStatus
		fmt.Printf("\n🔌 WebSocket Status:\n")
		for k, v := range wsStatus {
			fmt.Printf("  %s: %v\n", k, v)
		}
	}

	// Wait a bit more to capture some messages
	fmt.Println("\nWaiting 5 more seconds to capture WebSocket traffic...")
	time.Sleep(5 * time.Second)

	// Save to file
	fmt.Printf("\n💾 Saving captured data to %s...\n", *output)
	data, err := json.MarshalIndent(authData, "", "  ")
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error marshaling data: %v\n", err)
		os.Exit(1)
	}

	if err := os.WriteFile(*output, data, 0644); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing file: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("\n✅ Authentication data captured successfully!")
	fmt.Println("\nSummary:")
	fmt.Printf("  - Cookies: %d\n", len(authData.Cookies))
	fmt.Printf("  - HTTPOnly Cookies: %d\n", len(authData.HTTPOnlyCookies))
	fmt.Printf("  - Network Requests: %d\n", len(authData.NetworkRequests))
	fmt.Printf("  - WebSocket Frames: %d\n", len(authData.WebSocketFrames))
	fmt.Printf("  - WebSocket URL: %s\n", authData.WebSocketURL)
}

// Data structures
type AuthData struct {
	RoomURL           string                   `json:"room_url"`
	CaptureTime       time.Time                `json:"capture_time"`
	Cookies           []Cookie                 `json:"cookies"`
	HTTPOnlyCookies   []Cookie                 `json:"httponly_cookies"`
	RequestHeaders    map[string]interface{}   `json:"request_headers"`
	ResponseHeaders   map[string]interface{}   `json:"response_headers"`
	WebSocketURL      string                   `json:"websocket_url"`
	WebSocketRequestID string                  `json:"websocket_request_id"`
	WebSocketStatus   map[string]interface{}   `json:"websocket_status"`
	WebSocketFrames   []WebSocketFrame         `json:"websocket_frames"`
	NetworkRequests   []NetworkRequest         `json:"network_requests"`
}

type Cookie struct {
	Name     string `json:"name"`
	Value    string `json:"value"`
	Domain   string `json:"domain"`
	Path     string `json:"path"`
	Secure   bool   `json:"secure,omitempty"`
	HTTPOnly bool   `json:"httponly,omitempty"`
	SameSite string `json:"same_site,omitempty"`
}

type WebSocketFrame struct {
	Direction string    `json:"direction"`
	Data      string    `json:"data"`
	Time      time.Time `json:"time"`
}

type NetworkRequest struct {
	URL     string                 `json:"url"`
	Method  string                 `json:"method"`
	Headers map[string]interface{} `json:"headers"`
	Time    time.Time              `json:"time"`
}

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

func contains(s, substr string) bool {
	return len(s) >= len(substr) && (s == substr || 
		(len(s) > len(substr) && (
			s[:len(substr)] == substr || 
			findSubstring(s, substr))))
}

func findSubstring(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] + "..."
}

func getWebSocketHookScript() string {
	return `
		(function() {
			if (window.__KOSMI_WS_HOOK_INSTALLED__) return;
			
			const OriginalWebSocket = window.WebSocket;
			window.__KOSMI_MESSAGE_QUEUE__ = [];
			window.__KOSMI_WS__ = null;
			
			window.WebSocket = function(url, protocols) {
				const socket = new OriginalWebSocket(url, protocols);
				
				if (url.includes('engine.kosmi.io') || url.includes('gql-ws')) {
					console.log('[Auth Capture] WebSocket created:', url);
					window.__KOSMI_WS__ = socket;
					window.__KOSMI_WS_CONNECTED__ = true;
				}
				
				return socket;
			};
			
			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;
			
			window.__KOSMI_WS_HOOK_INSTALLED__ = true;
		})();
	`
}