package ws

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/sandrodach/mindsift/internal/domain"
)

// DocWSClient represents an active WebSocket connection to a specific document room.
type DocWSClient struct {
	hub    *Hub
	room   *DocRoom
	conn   *websocket.Conn
	userID string
	send   chan []byte
}

// DocSSEClient represents an active SSE connection to a specific document room.
type DocSSEClient struct {
	hub    *Hub
	room   *DocRoom
	id     string
	userID string
	send   chan []byte
}

// DocRoom coordinates in-memory real-time updates and debounced persistence for one document.
type DocRoom struct {
	mu                     sync.Mutex
	docID                  string
	hub                    *Hub
	stateBlob              []byte
	pendingUpdates         [][]byte
	persistedUpdatesCount  int
	wsClients              map[*DocWSClient]bool
	sseClients             map[*DocSSEClient]bool
	debounceTimer          *time.Timer
	dirty                  bool
	lastUserID             string
}

func (h *Hub) GetOrCreateDocRoom(docID string) *DocRoom {
	h.mu.Lock()
	defer h.mu.Unlock()

	if h.docRooms == nil {
		h.docRooms = make(map[string]*DocRoom)
	}

	if room, exists := h.docRooms[docID]; exists {
		return room
	}

	room := &DocRoom{
		docID:      docID,
		hub:        h,
		wsClients:  make(map[*DocWSClient]bool),
		sseClients: make(map[*DocSSEClient]bool),
	}

	// Pre-load current state from CRDT store if available
	if h.crdtStore != nil {
		stateBlob, updates, err := h.crdtStore.GetDocumentCRDTState(docID)
		if err == nil {
			room.stateBlob = stateBlob
			room.pendingUpdates = updates
			room.persistedUpdatesCount = len(updates)
		}
	}

	h.docRooms[docID] = room
	return room
}

// ApplyUpdate processes an incremental Yjs binary update in-memory, broadcasts it immediately
// to all other connected clients (WS & SSE), and schedules an asynchronous debounced flush to SQLite.
func (r *DocRoom) ApplyUpdate(update []byte, fromUserID string, senderWS *DocWSClient, senderSSE *DocSSEClient) {
	if len(update) == 0 {
		return
	}

	r.mu.Lock()
	r.pendingUpdates = append(r.pendingUpdates, update)
	r.dirty = true
	if fromUserID != "" {
		r.lastUserID = fromUserID
	}

	b64Update := base64.StdEncoding.EncodeToString(update)
	msg := domain.WSMessage{
		Event: "CRDT_UPDATE",
		Payload: map[string]interface{}{
			"docId":      r.docID,
			"update":     b64Update,
			"fromUserId": fromUserID,
		},
		Timestamp: domain.NowISO(),
	}
	msgBytes, _ := json.Marshal(msg)

	// Broadcast to all WS clients except sender
	for client := range r.wsClients {
		if senderWS != nil && client == senderWS {
			continue
		}
		select {
		case client.send <- msgBytes:
		default:
		}
	}

	// Broadcast to all SSE clients except sender
	for client := range r.sseClients {
		if senderSSE != nil && client == senderSSE {
			continue
		}
		select {
		case client.send <- msgBytes:
		default:
		}
	}

	// Schedule 2.5s debounced persistence
	r.scheduleDebouncedFlushLocked()
	r.mu.Unlock()

	// Also inform legacy room subscribers if any
	r.hub.BroadcastRoom("doc:"+r.docID, "CRDT_UPDATE", map[string]interface{}{
		"docId":      r.docID,
		"update":     b64Update,
		"fromUserId": fromUserID,
	}, nil)
}

// CompactState replaces the incremental pending updates with a full compacted state blob.
func (r *DocRoom) CompactState(stateBlob []byte, fromUserID string) {
	if len(stateBlob) == 0 {
		return
	}

	r.mu.Lock()
	r.stateBlob = stateBlob
	r.pendingUpdates = nil
	r.persistedUpdatesCount = 0
	r.dirty = true
	if fromUserID != "" {
		r.lastUserID = fromUserID
	}
	r.scheduleDebouncedFlushLocked()
	r.mu.Unlock()
}

// BroadcastAwareness forwards user awareness state to all other participants in the room.
func (r *DocRoom) BroadcastAwareness(payload json.RawMessage, fromUserID string, senderWS *DocWSClient, senderSSE *DocSSEClient) {
	msg := domain.WSMessage{
		Event:     "CRDT_AWARENESS",
		Payload:   payload,
		Timestamp: domain.NowISO(),
	}
	msgBytes, _ := json.Marshal(msg)

	r.mu.Lock()
	defer r.mu.Unlock()

	for client := range r.wsClients {
		if senderWS != nil && client == senderWS {
			continue
		}
		select {
		case client.send <- msgBytes:
		default:
		}
	}

	for client := range r.sseClients {
		if senderSSE != nil && client == senderSSE {
			continue
		}
		select {
		case client.send <- msgBytes:
		default:
		}
	}
}

func (r *DocRoom) scheduleDebouncedFlushLocked() {
	if r.debounceTimer != nil {
		r.debounceTimer.Stop()
	}
	r.debounceTimer = time.AfterFunc(2500*time.Millisecond, func() {
		r.Flush()
	})
}

// Flush persists in-memory pending updates to SQLite asynchronously in a background goroutine.
func (r *DocRoom) Flush() {
	r.mu.Lock()
	if !r.dirty {
		r.mu.Unlock()
		return
	}

	// Prepare data to flush
	docID := r.docID
	userID := r.lastUserID
	stateBlob := r.stateBlob
	var updatesToFlush [][]byte

	if len(stateBlob) > 0 && len(r.pendingUpdates) == 0 {
		// Compacted state snapshot
		r.dirty = false
		r.mu.Unlock()

		go func() {
			if r.hub.crdtStore != nil {
				_ = r.hub.crdtStore.SaveDocumentCRDTState(docID, stateBlob, userID)
			}
		}()
		return
	}

	if r.persistedUpdatesCount < len(r.pendingUpdates) {
		updatesToFlush = make([][]byte, len(r.pendingUpdates)-r.persistedUpdatesCount)
		copy(updatesToFlush, r.pendingUpdates[r.persistedUpdatesCount:])
		r.persistedUpdatesCount = len(r.pendingUpdates)
	}
	r.dirty = false
	r.mu.Unlock()

	if len(updatesToFlush) > 0 {
		go func() {
			if r.hub.crdtStore != nil {
				for _, u := range updatesToFlush {
					_ = r.hub.crdtStore.AppendDocumentCRDTUpdate(docID, u, userID)
				}
			}
		}()
	}
}

// GetInitialSnapshot returns the base64-encoded state blob and array of incremental updates.
func (r *DocRoom) GetInitialSnapshot() (string, []string) {
	r.mu.Lock()
	defer r.mu.Unlock()

	b64State := ""
	if len(r.stateBlob) > 0 {
		b64State = base64.StdEncoding.EncodeToString(r.stateBlob)
	}

	var b64Updates []string
	for _, u := range r.pendingUpdates {
		if len(u) > 0 {
			b64Updates = append(b64Updates, base64.StdEncoding.EncodeToString(u))
		}
	}
	return b64State, b64Updates
}

// ServeDocWS handles WebSocket connection for GET /api/v1/ws/docs/{id}
func (r *DocRoom) ServeDocWS(w http.ResponseWriter, req *http.Request, userID string) {
	conn, err := upgrader.Upgrade(w, req, nil)
	if err != nil {
		log.Printf("[doc-ws] upgrade error for doc %s: %v", r.docID, err)
		return
	}

	client := &DocWSClient{
		hub:    r.hub,
		room:   r,
		conn:   conn,
		userID: userID,
		send:   make(chan []byte, 256),
	}

	r.mu.Lock()
	r.wsClients[client] = true
	r.mu.Unlock()

	// Send initial snapshot
	state, updates := r.GetInitialSnapshot()
	initMsg := domain.WSMessage{
		Event: "CRDT_INIT",
		Payload: map[string]interface{}{
			"docId":   r.docID,
			"state":   state,
			"updates": updates,
		},
		Timestamp: domain.NowISO(),
	}
	initBytes, _ := json.Marshal(initMsg)
	client.send <- initBytes

	go client.writePump()
	go client.readPump()
}

// ServeDocSSE handles SSE stream fallback for GET /api/v1/sync/docs/{id}/stream
func (r *DocRoom) ServeDocSSE(w http.ResponseWriter, req *http.Request, userID string) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache, no-transform")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	client := &DocSSEClient{
		hub:    r.hub,
		room:   r,
		id:     domain.NewID("sse-doc"),
		userID: userID,
		send:   make(chan []byte, 256),
	}

	r.mu.Lock()
	r.sseClients[client] = true
	r.mu.Unlock()

	defer func() {
		r.mu.Lock()
		delete(r.sseClients, client)
		clientCount := len(r.wsClients) + len(r.sseClients)
		r.mu.Unlock()

		if clientCount == 0 {
			r.Flush()
		}
	}()

	// Send initial connected comment and CRDT_INIT
	fmt.Fprintf(w, ": connected (transport: sse, docId: %s)\n\n", r.docID)
	state, updates := r.GetInitialSnapshot()
	initMsg, _ := json.Marshal(domain.WSMessage{
		Event: "CRDT_INIT",
		Payload: map[string]interface{}{
			"docId":   r.docID,
			"state":   state,
			"updates": updates,
		},
		Timestamp: domain.NowISO(),
	})
	fmt.Fprintf(w, "data: %s\n\n", initMsg)
	flusher.Flush()

	ticker := time.NewTicker(15 * time.Second)
	defer ticker.Stop()

	notify := req.Context().Done()

	for {
		select {
		case <-notify:
			return
		case msg, ok := <-client.send:
			if !ok {
				return
			}
			fmt.Fprintf(w, "data: %s\n\n", msg)
			flusher.Flush()
		case <-ticker.C:
			// Ping to keep proxy connection alive
			fmt.Fprintf(w, ": ping\n\n")
			flusher.Flush()
		}
	}
}

func (c *DocWSClient) readPump() {
	defer func() {
		c.room.mu.Lock()
		delete(c.room.wsClients, c)
		clientCount := len(c.room.wsClients) + len(c.room.sseClients)
		c.room.mu.Unlock()

		_ = c.conn.Close()

		// Broadcast awareness leave
		c.room.BroadcastAwareness([]byte(fmt.Sprintf(`{"userId":"%s","activeBlockId":null}`, c.userID)), c.userID, nil, nil)

		if clientCount == 0 {
			c.room.Flush()
		}
	}()

	c.conn.SetReadLimit(8 * 1024 * 1024) // 8MB limit for binary/text updates
	_ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error {
		_ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	for {
		msgType, message, err := c.conn.ReadMessage()
		if err != nil {
			break
		}

		if msgType == websocket.BinaryMessage {
			// Direct binary Yjs update
			c.room.ApplyUpdate(message, c.userID, c, nil)
			continue
		}

		// Text JSON message
		var incoming struct {
			Event   string          `json:"event"`
			Payload json.RawMessage `json:"payload"`
		}
		if json.Unmarshal(message, &incoming) == nil {
			switch incoming.Event {
			case "CRDT_UPDATE":
				var p struct {
					Update string `json:"update"`
				}
				if json.Unmarshal(incoming.Payload, &p) == nil && p.Update != "" {
					if raw, err := base64.StdEncoding.DecodeString(p.Update); err == nil && len(raw) > 0 {
						c.room.ApplyUpdate(raw, c.userID, c, nil)
					}
				}
			case "CRDT_COMPACT":
				var p struct {
					State string `json:"state"`
				}
				if json.Unmarshal(incoming.Payload, &p) == nil && p.State != "" {
					if raw, err := base64.StdEncoding.DecodeString(p.State); err == nil && len(raw) > 0 {
						c.room.CompactState(raw, c.userID)
					}
				}
			case "CRDT_AWARENESS":
				c.room.BroadcastAwareness(incoming.Payload, c.userID, c, nil)
			}
		}
	}
}

func (c *DocWSClient) writePump() {
	ticker := time.NewTicker(30 * time.Second)
	defer func() {
		ticker.Stop()
		_ = c.conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			_ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				_ = c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
				return
			}
		case <-ticker.C:
			_ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}