package web

import (
	"encoding/json"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/hub"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/protocol"
)

// dispatch routes one inbound message to its handler. The `raw` payload is
// passed through so handlers can re-parse to their own shape.
//
// Phase 3 implements: get_salt, login, get_devices, ping, disconnect.
// Phase 4/5/6 commands (connect, mouse, key, term_*, etc.) get a friendly
// "not yet implemented" reply so the browser UI doesn't hang silently.
func (h *wsHub) dispatch(c *wsClient, cmd string, raw []byte) {
	switch cmd {
	case "get_salt":
		h.handleGetSalt(c, raw)
	case "login":
		h.handleLogin(c, raw)
	case "get_devices":
		h.handleGetDevices(c, raw)
	case "ping":
		// no-op heartbeat; the read itself was the keep-alive signal
	case "disconnect":
		h.handleDisconnect(c, raw)

	case "connect":
		h.handleConnect(c, raw)
	case "rdp_reset":
		// silently ignored — UI uses this as a fire-and-forget
	case "mouse", "key":
		// silently ignored — no remote screen yet
	case "term_open":
		h.replyNotImplemented(c, "term_closed", "Web terminal not yet implemented on Go server")
	case "term_input", "term_resize", "term_close":
		// silently ignored — no terminal session

	// Admin operations (Phase 7).
	case "create_user":
		h.replyNotImplemented(c, "create_user_result", "User management not yet implemented")
	case "delete_user":
		h.replyNotImplemented(c, "delete_user_result", "User management not yet implemented")
	case "list_users":
		h.replyNotImplemented(c, "list_users_result", "User management not yet implemented")
	case "get_groups":
		c.queue([]byte(`{"cmd":"groups","ok":true,"groups":[]}`))
	}
}

func (h *wsHub) replyNotImplemented(c *wsClient, replyCmd, msg string) {
	c.queue(mustJSON(map[string]any{
		"cmd": replyCmd,
		"ok":  false,
		"msg": msg,
	}))
}

// ----- handlers ------------------------------------------------------------

func (h *wsHub) handleGetSalt(c *wsClient, raw []byte) {
	var in struct {
		Username string `json:"username"`
	}
	_ = json.Unmarshal(raw, &in)

	salt, ok := h.auth.GetSalt(in.Username)
	// Do not leak which usernames exist: always return ok=true with a salt.
	// For unknown users hand back the empty salt (matches admin convention)
	// so the timing/shape of the response is uniform.
	if !ok {
		salt = ""
	}
	c.queue(mustJSON(map[string]any{
		"cmd":  "salt",
		"ok":   true,
		"salt": salt,
	}))
}

func (h *wsHub) handleLogin(c *wsClient, raw []byte) {
	var in struct {
		Username string `json:"username"`
		Response string `json:"response"`
		Nonce    string `json:"nonce"`
	}
	if err := json.Unmarshal(raw, &in); err != nil {
		c.queue(mustJSON(map[string]any{"cmd": "login_result", "ok": false, "msg": "Invalid request"}))
		return
	}

	// Bind the response to the challenge we issued at connect time so that
	// replays from a different connection can't reuse a captured response.
	if in.Nonce == "" || in.Nonce != c.nonce {
		c.queue(mustJSON(map[string]any{"cmd": "login_result", "ok": false, "msg": "Invalid challenge"}))
		return
	}

	token, role, err := h.auth.VerifyLogin(in.Username, in.Response, in.Nonce)
	if err != nil {
		// Burn the challenge on failure too — forces a new round on retry.
		c.nonce = ""
		c.queue(mustJSON(map[string]any{"cmd": "login_result", "ok": false, "msg": "Invalid credentials"}))
		return
	}
	c.nonce = ""
	c.token = token
	c.role = role
	h.log.Info("ws login: user=%s role=%s addr=%s", in.Username, role, c.addr)

	c.queue(mustJSON(map[string]any{
		"cmd":   "login_result",
		"ok":    true,
		"token": token,
		"role":  role,
	}))
}

// handleConnect kicks off a screen-sharing session for the browser. We send
// COMMAND_SCREEN_SPY to the device's main TCP connection; the device then
// opens a new sub-connection (TOKEN_BITMAPINFO) which the TCP side binds to
// the device via hub.BindScreenConn. Frame relay to the browser is handled
// in Phase 4.2 once frames actually arrive.
//
// Reply semantics: returning connect_result.ok=true (without width/height)
// triggers the browser's "Waiting for video..." spinner. We can't deliver
// width/height here because we don't yet know them — they show up in the
// first TOKEN_BITMAPINFO from the device.
// handleDisconnect detaches this client from any device it was watching and
// — if no other authenticated client is still watching — closes the device's
// screen sub-connection. Closing the TCP sub-conn is the signal the C++
// device firmware uses to stop screen capture, so this is how we ask the
// device to free its encoder.
func (h *wsHub) handleDisconnect(c *wsClient, _ []byte) {
	// Mirror handleConnect: take h.mu so event-handler readers
	// (OnResolutionChange/OnScreenFrame) get a consistent view of c.watching.
	h.mu.Lock()
	prev := c.watching
	c.watching = ""
	h.mu.Unlock()
	c.queue([]byte(`{"cmd":"disconnect_result","ok":true}`))
	if prev != "" && h.countWatchers(prev) == 0 {
		h.devices.CloseScreen(prev)
	}
}

// countWatchers returns how many authenticated clients still have their
// `watching` field pointing at deviceID. Called from disconnect paths.
func (h *wsHub) countWatchers(deviceID string) int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	n := 0
	for c := range h.clients {
		if c.watching == deviceID {
			n++
		}
	}
	return n
}

func (h *wsHub) handleConnect(c *wsClient, raw []byte) {
	if !h.requireAuth(c, raw, "connect_result") {
		return
	}
	var in struct {
		ID string `json:"id"`
	}
	if err := json.Unmarshal(raw, &in); err != nil || in.ID == "" {
		c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": false, "msg": "Bad request"}))
		return
	}

	// If a screen session is already live for this device (another browser
	// is already watching), reuse it: hand the new viewer the current
	// resolution and the most recent IDR keyframe so its decoder can start
	// rendering immediately, without waiting for the next IDR (~15 s).
	cache := h.devices.ScreenState(in.ID)
	if cache.Active {
		c.queue(mustJSON(map[string]any{
			"cmd": "connect_result", "ok": true,
			"width": cache.Width, "height": cache.Height,
		}))
		if len(cache.Keyframe) > 0 {
			c.queueBinary(cache.Keyframe)
		}
		h.mu.Lock()
		c.watching = in.ID
		h.mu.Unlock()
		return
	}

	// No active session — kick the device to start capturing. We send the
	// same 32-byte COMMAND_SCREEN_SPY payload the C++ WebService sends:
	//   [0]=COMMAND_SCREEN_SPY, [1]=0 (GDI), [2]=ALGORITHM_H264, [3]=1 (multi-screen),
	//   [4..31]=0.
	cmd := make([]byte, 32)
	cmd[0] = protocol.CommandScreenSpy
	cmd[2] = protocol.AlgorithmH264
	cmd[3] = 1

	// CRITICAL: bind c.watching BEFORE asking the device to start capturing.
	// On fast reconnects the device's screen sub-conn handshake completes in
	// <100 ms, so TOKEN_BITMAPINFO and even the first H264 frame can arrive
	// before this handler finishes — and the resolution_changed / frame
	// broadcasts in wsHub filter on c.watching. With the assignment after
	// SendToDevice the new viewer silently misses the very first IDR and
	// resolution_changed, leaving the page stuck on "Waiting for video".
	//
	// The write needs to share the lock event handlers use to read c.watching
	// (they iterate h.clients under h.mu.RLock). Without that the write is a
	// data race; on a fast reconnect the reader goroutine can keep observing
	// the previous value ("") long enough to drop the first resolution_changed
	// and the first IDR, which produces the exact "every other quick reconnect
	// goes black" symptom — the C++ server avoids it because it does the same
	// state mutation under std::mutex and reaps the memory-barrier as a bonus.
	h.mu.Lock()
	c.watching = in.ID
	h.mu.Unlock()

	if err := h.devices.SendToDevice(in.ID, cmd); err != nil {
		// Roll back the watching flag if we never managed to kick capture.
		h.mu.Lock()
		c.watching = ""
		h.mu.Unlock()
		msg := "Device offline"
		if err != hub.ErrDeviceOffline {
			msg = "Failed to start screen capture"
			h.log.Error("SendToDevice(%s): %v", in.ID, err)
		}
		c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": false, "msg": msg}))
		return
	}
	h.log.Info("[timing] COMMAND_SCREEN_SPY sent to device=%s (cold start)", in.ID)

	c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": true}))
}

func (h *wsHub) handleGetDevices(c *wsClient, raw []byte) {
	if !h.requireAuth(c, raw, "device_list") {
		return
	}
	devices := h.devices.ListDevices()
	c.queue(mustJSON(map[string]any{
		"cmd":     "device_list",
		"ok":      true,
		"devices": devices,
	}))
}

// requireAuth validates the token embedded in raw against the authenticator's
// session store (not against c.token). Tokens live independently of WS
// connections — the browser may reconnect after a visibility/network blip and
// resume with the same token, so we must not tie validity to one WS lifetime.
// On the first authenticated message we cache the token/role on the wsClient
// so broadcasts know to deliver to this connection.
func (h *wsHub) requireAuth(c *wsClient, raw []byte, replyCmd string) bool {
	var in struct {
		Token string `json:"token"`
	}
	_ = json.Unmarshal(raw, &in)
	if in.Token == "" {
		c.queue(mustJSON(map[string]any{"cmd": replyCmd, "ok": false}))
		return false
	}
	sess, err := h.auth.ValidateToken(in.Token)
	if err != nil {
		c.queue(mustJSON(map[string]any{"cmd": replyCmd, "ok": false}))
		return false
	}
	if c.token == "" {
		c.token = in.Token
		c.role = sess.Role
	}
	return true
}

func mustJSON(v any) []byte {
	b, err := json.Marshal(v)
	if err != nil {
		// All callers pass simple map[string]any with primitive values;
		// marshal can't realistically fail. If it does, return a safe fallback.
		return []byte(`{"cmd":"error","msg":"internal encode error"}`)
	}
	return b
}