package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"

	"github.com/gorilla/websocket"
)

const (
	// Kosmi WebSocket endpoint
	kosmiWSURL = "wss://engine.kosmi.io/gql-ws"
	
	// Your room ID (MUST include @ symbol!)
	roomID = "@hyperspaceout"
	
	// User agent and app version (from AUTH_FINDINGS.md)
	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"
)

// GraphQL-WS Protocol message types
const (
	MessageTypeConnectionInit = "connection_init"
	MessageTypeConnectionAck  = "connection_ack"
	MessageTypeSubscribe      = "subscribe"
	MessageTypeNext           = "next"
	MessageTypeError          = "error"
	MessageTypeComplete       = "complete"
)

// Message represents a graphql-ws protocol message
type Message struct {
	ID      string      `json:"id,omitempty"`
	Type    string      `json:"type"`
	Payload interface{} `json:"payload,omitempty"` // Can be map or array
}

func main() {
	log.Println("🚀 Starting Kosmi GraphQL-WS Proof of Concept")
	log.Printf("📡 WebSocket URL: %s", kosmiWSURL)
	log.Printf("🏠 Room ID: %s", roomID)

	// Step 1: Get anonymous token
	log.Println("\n📝 Step 1: Getting anonymous token...")
	token, err := getAnonymousToken()
	if err != nil {
		log.Fatalf("❌ Failed to get token: %v", err)
	}
	log.Printf("✅ Got token: %s...", token[:50])

	// Set up interrupt handler
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)
	
	log.Println("\n🔌 Step 2: Connecting to WebSocket...")

	// Create WebSocket dialer with headers
	dialer := websocket.Dialer{
		Subprotocols: []string{"graphql-transport-ws"},
		HandshakeTimeout: 10 * time.Second,
	}

	// Add headers that a browser would send
	headers := http.Header{}
	headers.Add("Origin", "https://app.kosmi.io")
	headers.Add("User-Agent", userAgent)

	// Connect to WebSocket
	log.Println("🔌 Establishing WebSocket connection...")
	conn, resp, err := dialer.Dial(kosmiWSURL, headers)
	if err != nil {
		log.Fatalf("❌ Failed to connect: %v (response: %+v)", err, resp)
	}
	defer conn.Close()

	log.Println("✅ WebSocket connected!")
	log.Printf("📋 Response status: %s", resp.Status)
	log.Printf("📋 Subprotocol: %s", conn.Subprotocol())

	// Channel for receiving messages
	done := make(chan struct{})
	
	// Start reading messages
	go func() {
		defer close(done)
		for {
			var msg Message
			err := conn.ReadJSON(&msg)
			if err != nil {
				log.Printf("❌ Read error: %v", err)
				return
			}
			
			log.Printf("📨 Received: %s (id: %s)", msg.Type, msg.ID)
			if msg.Payload != nil {
				payloadJSON, _ := json.MarshalIndent(msg.Payload, "", "  ")
				log.Printf("   Payload: %s", string(payloadJSON))
			}
			
			// Handle different message types
			switch msg.Type {
			case MessageTypeConnectionAck:
				log.Println("✅ Connection acknowledged by server!")
				// First, subscribe to messages (like the browser does)
				go subscribeToMessages(conn)
				
			case MessageTypeNext:
				log.Println("✅ Received 'next' - operation successful!")
				
			case MessageTypeError:
				log.Println("❌ Received 'error' from server!")
				
			case MessageTypeComplete:
				log.Println("✅ Received 'complete' - operation finished!")
				// After joining, send a message
				if msg.ID == "join-room" {
					log.Println("\n⏭️  Room joined, now sending message...")
					go sendMessage(conn)
				}
			}
		}
	}()

	// Step 1: Send connection_init with required payload
	log.Println("📤 Sending connection_init...")
	
	// Base64 encode the user agent (required by Kosmi)
	uaEncoded := base64.StdEncoding.EncodeToString([]byte(userAgent))
	
	initMsg := Message{
		Type: MessageTypeConnectionInit,
		Payload: map[string]interface{}{
			"token": token,       // JWT token from anonLogin
			"ua":    uaEncoded,   // Base64-encoded User-Agent
			"v":     appVersion,  // App version "4364"
			"r":     "",          // Room-related field (empty for anonymous)
		},
	}
	
	if err := conn.WriteJSON(initMsg); err != nil {
		log.Fatalf("❌ Failed to send connection_init: %v", err)
	}
	
	log.Println("   Payload: token, ua (base64), v=4364, r=\"\"")

	// Wait for interrupt or done
	select {
	case <-done:
		log.Println("🛑 Connection closed")
	case <-interrupt:
		log.Println("🛑 Interrupt received, closing connection...")
		
		// Send close message
		err := conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if err != nil {
			log.Printf("❌ Error sending close: %v", err)
		}
		
		select {
		case <-done:
		case <-time.After(time.Second):
		}
	}
}

func getAnonymousToken() (string, error) {
	// GraphQL mutation for anonymous login
	mutation := map[string]interface{}{
		"query": `mutation { anonLogin { token } }`,
	}
	
	jsonBody, err := json.Marshal(mutation)
	if err != nil {
		return "", err
	}
	
	log.Printf("   Sending: %s", string(jsonBody))
	
	// Create HTTP client
	client := &http.Client{Timeout: 10 * time.Second}
	
	// Create request with body
	req, err := http.NewRequest("POST", "https://engine.kosmi.io/", nil)
	if err != nil {
		return "", err
	}
	
	// Set headers and body
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Referer", "https://app.kosmi.io/")
	req.Header.Set("User-Agent", userAgent)
	req.Body = io.NopCloser(bytes.NewReader(jsonBody))
	req.ContentLength = int64(len(jsonBody))
	
	// Send request
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	
	log.Printf("   Response status: %d", resp.StatusCode)
	
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("HTTP %d", resp.StatusCode)
	}
	
	// Parse response
	var result map[string]interface{}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}
	
	// Extract token from response
	if data, ok := result["data"].(map[string]interface{}); ok {
		if anonLogin, ok := data["anonLogin"].(map[string]interface{}); ok {
			if token, ok := anonLogin["token"].(string); ok {
				return token, nil
			}
		}
	}
	
	resultJSON, _ := json.MarshalIndent(result, "", "  ")
	return "", fmt.Errorf("no token in response: %s", string(resultJSON))
}

func subscribeToMessages(conn *websocket.Conn) {
	// Wait a bit
	time.Sleep(1 * time.Second)
	
	log.Println("\n📨 Step 3: Subscribing to message feed (like browser does)...")
	log.Printf("   Room ID: %s", roomID)
	
	// Subscribe to new messages (copied from browser monitoring)
	subscription := `subscription OnNewMessage($roomId: String!) {
		newMessage(roomId: $roomId) {
			id
			body
			time
			user {
				id
				displayName
				username
				avatarUrl
			}
		}
	}`
	
	msg := Message{
		ID:   "subscribe-messages",
		Type: MessageTypeSubscribe,
		Payload: map[string]interface{}{
			"query": subscription,
			"variables": map[string]interface{}{
				"roomId": roomID,
			},
		},
	}
	
	msgJSON, _ := json.MarshalIndent(msg, "", "  ")
	log.Printf("📋 Sending subscription:\n%s", string(msgJSON))
	
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("❌ Failed to subscribe: %v", err)
		return
	}
	
	log.Println("✅ Subscription sent!")
	log.Println("   Note: Subscriptions stay open and don't send 'complete'")
	log.Println("   Proceeding to next step...")
	
	// Wait a bit, then proceed to join room
	time.Sleep(2 * time.Second)
	go joinRoomWithCorrectID(conn)
}

func joinRoomWithCorrectID(conn *websocket.Conn) {
	// Wait a bit
	time.Sleep(1 * time.Second)
	
	log.Println("\n🚪 Step 4: Joining room with correct ID format...")
	log.Printf("   Room ID: %s", roomID)
	
	// Join room mutation with correct ID format
	mutation := `mutation JoinRoom($id: String!) {
		joinRoom(id: $id) {
			ok
		}
	}`
	
	msg := Message{
		ID:   "join-room",
		Type: MessageTypeSubscribe,
		Payload: map[string]interface{}{
			"query": mutation,
			"variables": map[string]interface{}{
				"id": roomID,
			},
		},
	}
	
	msgJSON, _ := json.MarshalIndent(msg, "", "  ")
	log.Printf("📋 Sending joinRoom mutation:\n%s", string(msgJSON))
	
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("❌ Failed to join room: %v", err)
		return
	}
	
	log.Println("✅ Join room request sent! Waiting for response...")
}

func joinRoom(conn *websocket.Conn) {
	// Wait a bit to ensure connection is fully established
	time.Sleep(1 * time.Second)
	
	log.Println("\n🚪 Step 3: Getting room info and joining...")
	
	// First, query for the room to get its actual ID
	query := `query GetRoom($name: String!) {
		roomByName(name: $name) {
			id
			name
			displayName
		}
	}`
	
	msg := Message{
		ID:   "get-room",
		Type: MessageTypeSubscribe,
		Payload: map[string]interface{}{
			"query": query,
			"variables": map[string]interface{}{
				"name": roomID,
			},
		},
	}
	
	msgJSON, _ := json.MarshalIndent(msg, "", "  ")
	log.Printf("📋 Querying for room by name:\n%s", string(msgJSON))
	
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("❌ Failed to query room: %v", err)
		return
	}
	
	log.Println("✅ Room query sent! Waiting for response...")
}

func actuallyJoinRoom(conn *websocket.Conn) {
	// Wait a bit
	time.Sleep(1 * time.Second)
	
	log.Println("\n🚪 Step 3b: Actually joining room with ID...")
	
	// Join room mutation - we'll use the room name for now
	// The server might accept the name directly
	mutation := `mutation JoinRoom($id: String!) {
		joinRoom(id: $id) {
			ok
		}
	}`
	
	msg := Message{
		ID:   "join-room",
		Type: MessageTypeSubscribe,
		Payload: map[string]interface{}{
			"query": mutation,
			"variables": map[string]interface{}{
				"id": roomID,
			},
		},
	}
	
	msgJSON, _ := json.MarshalIndent(msg, "", "  ")
	log.Printf("📋 Sending joinRoom mutation:\n%s", string(msgJSON))
	
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("❌ Failed to join room: %v", err)
		return
	}
	
	log.Println("✅ Join room request sent! Waiting for response...")
}

func introspectSendMessage(conn *websocket.Conn) {
	// Wait a bit
	time.Sleep(1 * time.Second)
	
	log.Println("\n🔍 Step 4: Introspecting room and message mutations...")
	
	// Targeted introspection for room/message related mutations only
	introspectionQuery := `query IntrospectRoomMutations {
		mutations: __type(name: "RootMutationType") {
			fields {
				name
				args {
					name
					type {
						name
						kind
						ofType {
							name
							kind
						}
					}
				}
				type {
					name
					kind
					fields {
						name
						type {
							name
							kind
						}
					}
				}
			}
		}
		queries: __type(name: "RootQueryType") {
			fields {
				name
				args {
					name
					type {
						name
						kind
						ofType {
							name
							kind
						}
					}
				}
				type {
					name
					kind
				}
			}
		}
	}`
	
	msg := Message{
		ID:   "introspect-sendMessage",
		Type: MessageTypeSubscribe,
		Payload: map[string]interface{}{
			"query": introspectionQuery,
		},
	}
	
	msgJSON, _ := json.MarshalIndent(msg, "", "  ")
	log.Printf("📋 Sending introspection query:\n%s", string(msgJSON))
	
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("❌ Failed to send introspection: %v", err)
		return
	}
	
	log.Println("✅ Introspection sent! Waiting for response...")
}

func sendMessage(conn *websocket.Conn) {
	// Wait a bit
	time.Sleep(2 * time.Second)
	
	log.Println("\n📤 Step 5: Attempting to send 'Hello World' message...")
	log.Printf("   Using room ID: %s", roomID)
	
	// Fixed mutation based on introspection and browser monitoring:
	// - roomId is String!, not ID!
	// - roomId MUST include @ symbol (e.g., "@hyperspaceout")
	// - sendMessage returns Success type with field "ok: Boolean"
	mutation := `mutation SendMessage($body: String!, $roomId: String!) {
		sendMessage(body: $body, roomId: $roomId) {
			ok
		}
	}`
	
	// Create subscribe message (yes, mutations use "subscribe" in graphql-ws!)
	msg := Message{
		ID:   fmt.Sprintf("send-message-%d", time.Now().Unix()),
		Type: MessageTypeSubscribe,
		Payload: map[string]interface{}{
			"query": mutation,
			"variables": map[string]interface{}{
				"body":   "Hello World from GraphQL-WS! 🎉",
				"roomId": roomID,
			},
		},
	}
	
	msgJSON, _ := json.MarshalIndent(msg, "", "  ")
	log.Printf("📋 Sending message:\n%s", string(msgJSON))
	
	if err := conn.WriteJSON(msg); err != nil {
		log.Printf("❌ Failed to send message: %v", err)
		return
	}
	
	log.Println("✅ Message sent! Waiting for response...")
	log.Println("   Check the Kosmi chat to see if the message appeared!")
}