package websocket

import (
	"encoding/json"
	"log"
	"strconv"
	"sync"
	"time"

	"github.com/gofiber/websocket/v2"
	"gorm.io/gorm"
	internalDb "server/internal/db"
)

const (
	// Heartbeat kiểu game: phát hiện mất kết nối trong ~30–45s.
	wsPingInterval = 15 * time.Second
	wsPongWait     = 45 * time.Second
	// Grace khi SV reconnect — tránh nhấp nháy online/offline.
	studentOfflineGrace = 8 * time.Second
)

type SocketMsg struct {
	Event string         `json:"event"`
	Data  map[string]any `json:"data"`
}

type SocketClient struct {
	Conn      *websocket.Conn
	StudentID int64
	ClassID   int64
	StaffID   uint
	Role      string // "student" | "teacher"
	Addr      string
	writeMu   sync.Mutex
}

func (c *SocketClient) WriteJSON(v any) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.Conn.WriteJSON(v)
}

func (c *SocketClient) WriteRaw(msg []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	return c.Conn.WriteMessage(websocket.TextMessage, msg)
}

type offlineGrace struct {
	until   time.Time
	classID int64
}

type WsHub struct {
	mu              sync.RWMutex
	students        map[int64]*SocketClient
	teachers        map[string]*SocketClient
	teachersByStaff map[uint][]string // staffId -> teacher connection addresses
	subscribers     map[int64]map[string]string // studentId -> teacher connection address -> mode ("grid" | "focus")
	grace           map[int64]offlineGrace
	graceTimers     map[int64]*time.Timer
	lastRelay       map[string]time.Time        // studentID:event -> last relay time
}

var Hub = &WsHub{
	students:        make(map[int64]*SocketClient),
	teachers:        make(map[string]*SocketClient),
	teachersByStaff: make(map[uint][]string),
	subscribers:     make(map[int64]map[string]string),
	grace:           make(map[int64]offlineGrace),
	graceTimers:     make(map[int64]*time.Timer),
	lastRelay:       make(map[string]time.Time),
}

func (h *WsHub) IsStudentOnline(studentRkID int64) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	if _, ok := h.students[studentRkID]; ok {
		return true
	}
	g, ok := h.grace[studentRkID]
	return ok && time.Now().Before(g.until)
}

func (h *WsHub) PushExamToStudent(studentRkID int64, data map[string]any) {
	h.mu.RLock()
	client, ok := h.students[studentRkID]
	h.mu.RUnlock()
	if !ok || client == nil {
		return
	}
	_ = client.WriteJSON(SocketMsg{Event: "exam:paper-sent", Data: data})
}

func (h *WsHub) PushChatToStudent(studentRkID int64, data map[string]any) {
	h.mu.RLock()
	client, ok := h.students[studentRkID]
	h.mu.RUnlock()
	if !ok || client == nil {
		return
	}
	_ = client.WriteJSON(SocketMsg{Event: "chat:message", Data: data})
}

func (h *WsHub) PushChatToStaff(staffID uint, data map[string]any) {
	if data == nil {
		data = map[string]any{}
	}
	data["targetStaffId"] = staffID
	msg := SocketMsg{Event: "chat:message", Data: data}
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, t := range h.teachers {
		if t != nil && t.Role == "teacher" {
			_ = t.WriteJSON(msg)
		}
	}
}

func (h *WsHub) BroadcastStudentViolation(data map[string]any) {
	msg := SocketMsg{Event: "teacher:student-violation", Data: data}
	msgBytes, err := json.Marshal(msg)
	if err != nil {
		return
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, t := range h.teachers {
		if t != nil {
			_ = t.WriteRaw(msgBytes)
		}
	}
}

func (h *WsHub) broadcastPresence(studentID int64, online bool, classID int64) {
	msg := SocketMsg{
		Event: "presence:update",
		Data: map[string]any{
			"studentId": studentID,
			"online":    online,
			"classId":   classID,
		},
	}
	msgBytes, err := json.Marshal(msg)
	if err != nil {
		return
	}
	for _, t := range h.teachers {
		if t != nil {
			_ = t.WriteRaw(msgBytes)
		}
	}
}

// ForceStudentOffline — tắt ngay (không chờ grace), dùng khi SV cố ý đóng app / vi phạm.
func (h *WsHub) ForceStudentOffline(studentID int64) {
	h.mu.Lock()
	defer h.mu.Unlock()

	classID := int64(0)
	if client, ok := h.students[studentID]; ok && client != nil {
		classID = client.ClassID
		delete(h.students, studentID)
	}
	if g, ok := h.grace[studentID]; ok {
		if classID == 0 {
			classID = g.classID
		}
	}
	h.cancelGraceLocked(studentID)

	log.Printf("[WS] Student %d force offline (intentional quit/violation)", studentID)
	h.broadcastPresence(studentID, false, classID)
	if teachers, exists := h.subscribers[studentID]; exists {
		for tAddr := range teachers {
			if t, found := h.teachers[tAddr]; found {
				_ = t.WriteJSON(SocketMsg{
					Event: "teacher:stream-stopped",
					Data:  map[string]any{"studentId": studentID},
				})
			}
		}
	}
}

func (h *WsHub) cancelGraceLocked(studentID int64) {
	if t, ok := h.graceTimers[studentID]; ok {
		t.Stop()
		delete(h.graceTimers, studentID)
	}
	delete(h.grace, studentID)
}

func (h *WsHub) finalizeStudentOffline(studentID int64, classID int64) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if _, online := h.students[studentID]; online {
		return
	}
	if _, ok := h.grace[studentID]; !ok {
		return
	}
	delete(h.grace, studentID)
	delete(h.graceTimers, studentID)

	log.Printf("[WS] Student %d offline after grace", studentID)
	h.broadcastPresence(studentID, false, classID)

	if teachers, exists := h.subscribers[studentID]; exists {
		for tAddr := range teachers {
			if t, found := h.teachers[tAddr]; found {
				_ = t.WriteJSON(SocketMsg{
					Event: "teacher:stream-stopped",
					Data:  map[string]any{"studentId": studentID},
				})
			}
		}
	}
}

func (h *WsHub) GetOnlineStudentIDs(classID int64) []int64 {
	h.mu.RLock()
	defer h.mu.RUnlock()

	seen := map[int64]bool{}
	var ids []int64
	now := time.Now()
	for _, client := range h.students {
		if client.ClassID == classID {
			seen[client.StudentID] = true
			ids = append(ids, client.StudentID)
		}
	}
	for sid, g := range h.grace {
		if now.Before(g.until) && g.classID == classID && !seen[sid] {
			ids = append(ids, sid)
		}
	}
	if ids == nil {
		return []int64{}
	}
	return ids
}

func (h *WsHub) Register(c *SocketClient) {
	h.mu.Lock()
	defer h.mu.Unlock()

	c.Addr = c.Conn.RemoteAddr().String()
	if c.Role == "student" {
		wasInGrace := false
		if _, ok := h.grace[c.StudentID]; ok {
			wasInGrace = true
		}
		h.cancelGraceLocked(c.StudentID)
		h.students[c.StudentID] = c
		log.Printf("[WS] Student %d registered (Address: %s, Class: %d, reconnectGrace=%v)", c.StudentID, c.Addr, c.ClassID, wasInGrace)
		h.broadcastPresence(c.StudentID, true, c.ClassID)
		if _, exists := h.subscribers[c.StudentID]; exists {
			h.updateStudentStreamModeLocked(c.StudentID)
		}
	} else if c.Role == "teacher" {
		h.teachers[c.Addr] = c
		if c.StaffID > 0 {
			h.teachersByStaff[c.StaffID] = append(h.teachersByStaff[c.StaffID], c.Addr)
		}
		log.Printf("[WS] Teacher registered (Address: %s, Staff: %d)", c.Addr, c.StaffID)
	}
}

func (h *WsHub) Unregister(c *SocketClient) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if c.Role == "student" {
		current, exists := h.students[c.StudentID]
		if !exists || current != c {
			log.Printf("[WS] Student %d stale disconnect ignored (replaced by newer connection)", c.StudentID)
			return
		}
		delete(h.students, c.StudentID)
		classID := c.ClassID
		studentID := c.StudentID

		h.cancelGraceLocked(studentID)
		deadline := time.Now().Add(studentOfflineGrace)
		h.grace[studentID] = offlineGrace{until: deadline, classID: classID}
		h.graceTimers[studentID] = time.AfterFunc(studentOfflineGrace, func() {
			h.finalizeStudentOffline(studentID, classID)
		})
		log.Printf("[WS] Student %d disconnected — grace %v before offline", studentID, studentOfflineGrace)
		// Chưa broadcast offline / stream-stopped — chờ grace (reconnect nhanh như game).
	} else if c.Role == "teacher" {
		delete(h.teachers, c.Addr)
		if c.StaffID > 0 {
			list := h.teachersByStaff[c.StaffID]
			next := list[:0]
			for _, addr := range list {
				if addr != c.Addr {
					next = append(next, addr)
				}
			}
			if len(next) == 0 {
				delete(h.teachersByStaff, c.StaffID)
			} else {
				h.teachersByStaff[c.StaffID] = next
			}
		}
		log.Printf("[WS] Teacher %s disconnected", c.Addr)

		for sID, teachersMap := range h.subscribers {
			delete(teachersMap, c.Addr)
			if len(teachersMap) == 0 {
				delete(h.subscribers, sID)
				if student, exists := h.students[sID]; exists {
					_ = student.WriteJSON(SocketMsg{Event: "stop_screenshot_stream"})
					_ = student.WriteJSON(SocketMsg{Event: "stop_webcam_stream"})
				}
			} else {
				h.updateStudentStreamModeLocked(sID)
			}
		}
	}
}

func (h *WsHub) updateStudentStreamModeLocked(studentID int64) {
	student, exists := h.students[studentID]
	if !exists || student == nil {
		return
	}

	teachersMap, ok := h.subscribers[studentID]
	if !ok || len(teachersMap) == 0 {
		_ = student.WriteJSON(SocketMsg{Event: "stop_screenshot_stream"})
		_ = student.WriteJSON(SocketMsg{Event: "stop_webcam_stream"})
		return
	}

	// Determine effective mode: if any teacher wants "focus", use focus. Otherwise "grid".
	effectiveMode := "grid"
	for _, mode := range teachersMap {
		if mode == "focus" {
			effectiveMode = "focus"
			break
		}
	}

	intervalMs := 3000 // default for grid
	if effectiveMode == "focus" {
		intervalMs = 500 // 500ms for focus view
	}

	_ = student.WriteJSON(SocketMsg{
		Event: "start_screenshot_stream",
		Data:  map[string]any{"interval": intervalMs},
	})
	_ = student.WriteJSON(SocketMsg{
		Event: "start_webcam_stream",
		Data:  map[string]any{"interval": intervalMs},
	})
	log.Printf("[WS] Updated student %d stream mode to %s (interval: %dms)", studentID, effectiveMode, intervalMs)
}

func (h *WsHub) Subscribe(teacherAddr string, studentID int64, mode string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	teachersMap, exists := h.subscribers[studentID]
	if !exists {
		teachersMap = make(map[string]string)
		h.subscribers[studentID] = teachersMap
	}

	teachersMap[teacherAddr] = mode
	log.Printf("[WS] Teacher %s subscribed to student %d stream in %s mode", teacherAddr, studentID, mode)

	h.updateStudentStreamModeLocked(studentID)
}

func (h *WsHub) Unsubscribe(teacherAddr string, studentID int64) {
	h.mu.Lock()
	defer h.mu.Unlock()

	teachersMap, exists := h.subscribers[studentID]
	if !exists {
		return
	}

	delete(teachersMap, teacherAddr)
	log.Printf("[WS] Teacher %s unsubscribed from student %d stream", teacherAddr, studentID)

	if len(teachersMap) == 0 {
		delete(h.subscribers, studentID)
		log.Printf("[WS] Student %d has no more proctor subscribers. Stopping streams.", studentID)
		if student, exists := h.students[studentID]; exists {
			_ = student.WriteJSON(SocketMsg{Event: "stop_screenshot_stream"})
			_ = student.WriteJSON(SocketMsg{Event: "stop_webcam_stream"})
		}
	} else {
		h.updateStudentStreamModeLocked(studentID)
	}
}

func (h *WsHub) RelayFrame(studentID int64, event string, data map[string]any) {
	h.mu.Lock()
	teachersMap, exists := h.subscribers[studentID]
	if !exists || len(teachersMap) == 0 {
		h.mu.Unlock()
		return
	}

	// Compute effective mode/interval
	effectiveMode := "grid"
	for _, mode := range teachersMap {
		if mode == "focus" {
			effectiveMode = "focus"
			break
		}
	}

	minInterval := 3000 * time.Millisecond
	if effectiveMode == "focus" {
		minInterval = 500 * time.Millisecond
	}

	relayKey := strconv.FormatInt(studentID, 10) + ":" + event
	lastTime := h.lastRelay[relayKey]
	now := time.Now()
	if now.Sub(lastTime) < minInterval {
		h.mu.Unlock()
		return
	}
	h.lastRelay[relayKey] = now

	addrs := make([]string, 0, len(teachersMap))
	for addr := range teachersMap {
		addrs = append(addrs, addr)
	}
	h.mu.Unlock()

	relayEvent := "teacher:screenshot-stream-frame"
	if event == "webcam_stream_frame" {
		relayEvent = "teacher:webcam-stream-frame"
	}

	msg := SocketMsg{
		Event: relayEvent,
		Data: map[string]any{
			"studentId":   studentID,
			"imageBuffer": data["imageBuffer"],
		},
	}
	msgBytes, err := json.Marshal(msg)
	if err != nil {
		return
	}

	var dead []string
	h.mu.RLock()
	for _, addr := range addrs {
		if t, found := h.teachers[addr]; found {
			if err := t.WriteRaw(msgBytes); err != nil {
				log.Printf("[WS] Relay to teacher %s failed: %v", addr, err)
				dead = append(dead, addr)
			}
		}
	}
	h.mu.RUnlock()

	for _, addr := range dead {
		if t, ok := func() (*SocketClient, bool) {
			h.mu.RLock()
			defer h.mu.RUnlock()
			t, ok := h.teachers[addr]
			return t, ok
		}(); ok && t != nil {
			_ = t.Conn.Close()
		}
	}
}

func WebSocketHandler(db *gorm.DB) func(*websocket.Conn) {
	return func(c *websocket.Conn) {
		role := c.Query("role", "student")
		studentIDStr := c.Query("studentId", "0")
		classIDStr := c.Query("classId", "0")
		staffIDStr := c.Query("staffId", "0")

		studentID, _ := strconv.ParseInt(studentIDStr, 10, 64)
		classID, _ := strconv.ParseInt(classIDStr, 10, 64)
		staffID64, _ := strconv.ParseUint(staffIDStr, 10, 64)

		if role == "student" && db != nil {
			resolvedClassID := internalDb.FindActiveClassForStudent(db, studentID)
			if resolvedClassID > 0 {
				classID = resolvedClassID
			}
		}

		client := &SocketClient{
			Conn:      c,
			StudentID: studentID,
			ClassID:   classID,
			StaffID:   uint(staffID64),
			Role:      role,
		}

		Hub.Register(client)
		defer func() {
			Hub.Unregister(client)
			c.Close()
		}()

		c.SetReadDeadline(time.Now().Add(wsPongWait))
		c.SetPongHandler(func(string) error {
			return c.SetReadDeadline(time.Now().Add(wsPongWait))
		})

		pingDone := make(chan struct{})
		defer close(pingDone)
		go func() {
			ticker := time.NewTicker(wsPingInterval)
			defer ticker.Stop()
			for {
				select {
				case <-pingDone:
					return
				case <-ticker.C:
					client.writeMu.Lock()
					err := c.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(5*time.Second))
					client.writeMu.Unlock()
					if err != nil {
						return
					}
				}
			}
		}()

		for {
			_, msgBytes, err := c.ReadMessage()
			if err != nil {
				break
			}
			_ = c.SetReadDeadline(time.Now().Add(wsPongWait))

			var msg SocketMsg
			if err := json.Unmarshal(msgBytes, &msg); err != nil {
				continue
			}

			switch msg.Event {
			case "client:ping":
				_ = client.WriteJSON(SocketMsg{Event: "client:pong", Data: map[string]any{"t": time.Now().UnixMilli()}})

			case "screenshot_stream_frame", "webcam_stream_frame":
				Hub.RelayFrame(client.StudentID, msg.Event, msg.Data)

			case "teacher:subscribe":
				if client.Role == "teacher" {
					if sIDVal, ok := msg.Data["studentId"]; ok {
						var sID int64
						switch v := sIDVal.(type) {
						case float64:
							sID = int64(v)
						case string:
							sID, _ = strconv.ParseInt(v, 10, 64)
						}
						mode, _ := msg.Data["mode"].(string)
						if mode == "" {
							mode = "focus"
						}
						if sID > 0 {
							Hub.Subscribe(client.Addr, sID, mode)
						}
					}
				}

			case "teacher:unsubscribe":
				if client.Role == "teacher" {
					if sIDVal, ok := msg.Data["studentId"]; ok {
						var sID int64
						switch v := sIDVal.(type) {
						case float64:
							sID = int64(v)
						case string:
							sID, _ = strconv.ParseInt(v, 10, 64)
						}
						if sID > 0 {
							Hub.Unsubscribe(client.Addr, sID)
						}
					}
				}
			}
		}
	}
}