package main

import (
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"time"

	"github.com/gorilla/websocket"
)

const (
	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"
	appVersion = "4364"
)

func main() {
	roomURL := flag.String("room", "https://app.kosmi.io/room/@hyperspaceout", "Full Kosmi room URL")
	token := flag.String("token", "", "JWT token (optional, will try to extract from page)")
	flag.Parse()

	fmt.Println("🌐 Testing session-based WebSocket connection")
	fmt.Printf("   Room URL: %s\n\n", *roomURL)

	// Create HTTP client with cookie jar
	jar, err := cookiejar.New(nil)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create cookie jar: %v\n", err)
		os.Exit(1)
	}

	client := &http.Client{
		Jar:     jar,
		Timeout: 30 * time.Second,
	}

	// Step 1: Visit the room page to establish session
	fmt.Println("1️⃣  Visiting room page to establish session...")
	if err := visitRoomPage(client, *roomURL); err != nil {
		fmt.Fprintf(os.Stderr, "❌ Failed to visit room: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("✅ Session established!")

	// Print cookies
	u, _ := url.Parse(*roomURL)
	cookies := client.Jar.Cookies(u)
	fmt.Printf("\n   📋 Cookies received: %d\n", len(cookies))
	for _, c := range cookies {
		fmt.Printf("      - %s=%s\n", c.Name, truncate(c.Value, 50))
	}

	// Step 2: Connect WebSocket with cookies
	fmt.Println("\n2️⃣  Connecting WebSocket with session cookies...")
	
	roomID := extractRoomID(*roomURL)
	conn, err := connectWebSocketWithSession(client.Jar, *token, roomID)
	if err != nil {
		fmt.Fprintf(os.Stderr, "❌ Failed to connect WebSocket: %v\n", err)
		os.Exit(1)
	}
	defer conn.Close()
	fmt.Println("✅ WebSocket connected!")

	// Step 3: Listen for messages
	fmt.Println("\n👂 Listening for messages (press Ctrl+C to exit)...\n")
	
	messageCount := 0
	for {
		var msg map[string]interface{}
		if err := conn.ReadJSON(&msg); err != nil {
			fmt.Fprintf(os.Stderr, "Error reading message: %v\n", err)
			break
		}

		msgType, _ := msg["type"].(string)
		
		switch msgType {
		case "next":
			payload, _ := msg["payload"].(map[string]interface{})
			data, _ := payload["data"].(map[string]interface{})
			
			if newMessage, ok := data["newMessage"].(map[string]interface{}); ok {
				messageCount++
				body, _ := newMessage["body"].(string)
				user, _ := newMessage["user"].(map[string]interface{})
				username, _ := user["displayName"].(string)
				if username == "" {
					username, _ = user["username"].(string)
				}
				timestamp, _ := newMessage["time"].(float64)
				
				t := time.Unix(int64(timestamp), 0)
				fmt.Printf("[%s] %s: %s\n", t.Format("15:04:05"), username, body)
			}
		case "connection_ack":
			fmt.Println("   ✅ Received connection_ack")
		case "complete":
			id, _ := msg["id"].(string)
			fmt.Printf("   [Subscription %s completed]\n", id)
		case "error":
			fmt.Printf("   ⚠️  Error: %+v\n", msg)
		case "ka":
			// Keep-alive, ignore
		default:
			fmt.Printf("   📨 %s\n", msgType)
		}
	}

	fmt.Printf("\n📊 Total messages received: %d\n", messageCount)
}

func visitRoomPage(client *http.Client, roomURL string) error {
	req, err := http.NewRequest("GET", roomURL, nil)
	if err != nil {
		return err
	}

	req.Header.Set("User-Agent", userAgent)
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8")
	req.Header.Set("Accept-Language", "en-US,en;q=0.9")

	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return fmt.Errorf("status %d", resp.StatusCode)
	}

	// Read and discard body (but process Set-Cookie headers)
	io.Copy(io.Discard, resp.Body)
	
	return nil
}

func connectWebSocketWithSession(jar http.CookieJar, token, roomID string) (*websocket.Conn, error) {
	dialer := websocket.Dialer{
		Jar:             jar,
		Subprotocols:    []string{"graphql-ws"},
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	
	headers := http.Header{
		"Origin":     []string{"https://app.kosmi.io"},
		"User-Agent": []string{userAgent},
	}
	
	conn, resp, err := dialer.Dial("wss://engine.kosmi.io/gql-ws", headers)
	if err != nil {
		if resp != nil {
			fmt.Printf("   Response status: %d\n", resp.StatusCode)
			// Print response headers
			fmt.Println("   Response headers:")
			for k, v := range resp.Header {
				fmt.Printf("      %s: %v\n", k, v)
			}
		}
		return nil, err
	}

	// Send connection_init
	// If token not provided, try without it
	uaEncoded := base64.StdEncoding.EncodeToString([]byte(userAgent))
	
	payload := map[string]interface{}{
		"ua": uaEncoded,
		"v":  appVersion,
		"r":  "",
	}
	
	if token != "" {
		payload["token"] = token
	}
	
	initMsg := map[string]interface{}{
		"type":    "connection_init",
		"payload": payload,
	}

	if err := conn.WriteJSON(initMsg); err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to send connection_init: %w", err)
	}

	// Wait for ack
	var ackMsg map[string]interface{}
	if err := conn.ReadJSON(&ackMsg); err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to read ack: %w", err)
	}

	msgType, _ := ackMsg["type"].(string)
	if msgType != "connection_ack" {
		conn.Close()
		return nil, fmt.Errorf("expected connection_ack, got %s", msgType)
	}

	// Subscribe to messages
	query := fmt.Sprintf(`
		subscription {
			newMessage(roomId: "%s") {
				body
				time
				user {
					displayName
					username
				}
			}
		}
	`, roomID)

	subMsg := map[string]interface{}{
		"id":   "newMessage-subscription",
		"type": "subscribe",
		"payload": map[string]interface{}{
			"query":     query,
			"variables": map[string]interface{}{},
		},
	}

	if err := conn.WriteJSON(subMsg); err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to subscribe: %w", err)
	}

	return conn, nil
}

func extractRoomID(roomURL string) string {
	// Extract room ID from URL
	// https://app.kosmi.io/room/@roomname -> @roomname
	// https://app.kosmi.io/room/roomid -> roomid
	parts := make([]string, 0)
	for _, part := range []byte(roomURL) {
		if part == '/' {
			parts = append(parts, "")
		} else if len(parts) > 0 {
			parts[len(parts)-1] += string(part)
		}
	}
	
	if len(parts) > 0 {
		return parts[len(parts)-1]
	}
	return roomURL
}

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