package main

import (
	"context"
	"encoding/binary"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/auth"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/connection"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/hub"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/licensing"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/logger"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/protocol"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/server"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/web"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/wsauth"
)

// MyHandler implements the server.Handler interface
type MyHandler struct {
	log    *logger.Logger
	auth   *auth.Authenticator
	srv    *server.Server
	hub    *hub.Hub
	signer licensing.Signer // CMD_MASTERSETTING signer (local / remote / noop)
}

// OnConnect is called when a client connects
func (h *MyHandler) OnConnect(ctx *connection.Context) {
	// Only log connection established, detailed info logged on login
}

// OnDisconnect is called when a client disconnects
func (h *MyHandler) OnDisconnect(ctx *connection.Context) {
	// Always clean up any sub-context mapping first — the connection may
	// be a screen / terminal sub-conn rather than a main login connection.
	// Both Unbind* calls are no-ops if not tracked. UnbindTerminalConn
	// also fires OnTerminalClosed so the browser sees the session end on
	// unexpected device-side drops.
	h.hub.UnbindScreenConn(ctx)
	h.hub.UnbindTerminalConn(ctx)

	info := ctx.GetInfo()
	// Only treat this disconnect as a device-going-offline event if this
	// ctx is the device's MAIN login connection. Phase 6 added ClientID
	// pinning to sub-conns (via ConnAuth — needed for terminal routing),
	// so a non-empty ClientID alone no longer distinguishes main from
	// sub. Closing a screen / terminal sub-conn must NOT remove the
	// device from the hub.
	if info.ClientID != "" && h.hub.MainConn(info.ClientID) == ctx {
		h.log.ClientEvent("offline", ctx.ID, ctx.GetPeerIP(),
			"clientID", info.ClientID,
			"computer", info.ComputerName,
		)
		// Tear down any active sub-conn sessions BEFORE Unregister so the
		// browser sees screen/terminal close events alongside the
		// device-offline event, instead of frames/output continuing to
		// stream from orphaned sub-conn ctxs until they time out on
		// their own. Both calls no-op if there's no active session.
		h.hub.CloseScreen(info.ClientID)
		h.hub.CloseTerminalSession(info.ClientID)
		h.hub.Unregister(info.ClientID)
	}
}

// OnReceive is called when data is received from a client
func (h *MyHandler) OnReceive(ctx *connection.Context, data []byte) {
	if len(data) == 0 {
		return
	}

	// Terminal-bound sub-conns deliver RAW shell output with no leading
	// command byte — see client/ConPTYManager.cpp:328 (Send2Server with
	// just the buffer). We must short-circuit BEFORE the command switch
	// or the first output byte will be misinterpreted as a token.
	// Exception: a length-1 packet whose byte is TOKEN_TERMINAL_CLOSE
	// is the device's "shell exited" notification, NOT data.
	if devID := h.hub.TerminalDeviceID(ctx); devID != "" {
		if len(data) == 1 && data[0] == protocol.TokenTerminalClose {
			h.log.Info("terminal closed by device=%s conn=%d", devID, ctx.ID)
			h.hub.CloseTerminalSession(devID)
			return
		}
		// Wrap with the 'TRM1' magic the browser uses to demultiplex
		// terminal output from screen frames over the shared WS.
		packet := make([]byte, 4+len(data))
		copy(packet[:4], protocol.TerminalBinaryMagic[:])
		copy(packet[4:], data)
		h.hub.PublishTerminalData(devID, packet)
		return
	}

	cmd := data[0]
	// Handle commands
	switch cmd {
	case protocol.TokenLogin:
		h.handleLogin(ctx, data)
	case protocol.TokenAuth:
		h.handleAuth(ctx, data)
	case protocol.TokenHeartbeat:
		h.handleHeartbeat(ctx, data)
	case protocol.TokenConnAuth:
		h.handleConnAuth(ctx, data)
	case protocol.TokenBitmapInfo:
		h.handleBitmapInfo(ctx, data)
	case protocol.TokenTerminalStart:
		h.handleTerminalStart(ctx, true)
	case protocol.TokenShellStart:
		h.handleTerminalStart(ctx, false)
	case protocol.TokenTerminalClose:
		// Pre-bind close (rare — device gives up before the server
		// finished its half of the handshake). Best-effort cleanup.
		if devID := h.deviceIDOfSubConn(ctx); devID != "" {
			h.log.Info("pre-bind terminal close: device=%s conn=%d", devID, ctx.ID)
			h.hub.CloseTerminalSession(devID)
		}
	case protocol.TokenFirstScreen:
		// TOKEN_FIRSTSCREEN delivers a RAW BGRA baseline frame, not an
		// H264 unit — bytes ≈ width × height × 4. The C++ MFC dialog
		// blits it directly into a DIB; web viewers only consume H264 NAL
		// data, so dropping it here is correct. The first real H264 IDR
		// arrives shortly after via TOKEN_NEXTSCREEN.
	case protocol.TokenNextScreen:
		h.handleScreenFrame(ctx, data, false)
	case protocol.TokenKeyframe:
		// Sent by the client only when frameID % m_GOP == 0; the client's
		// DEFAULT_GOP is 0x7FFFFFFF (effectively infinite), so this token
		// is essentially unused in practice. Treat as a no-op for now —
		// IDRs always arrive in-band via TOKEN_NEXTSCREEN and we catch
		// them via the H264 NAL scan in handleScreenFrame.
	case protocol.CmdCursorImage:
		// Custom cursor bitmaps — relayed in Phase 5+ when the web cursor
		// overlay learns to render arbitrary BGRA images. Drop silently for
		// now; the standard IDC_* index (data[10] of every frame header) is
		// what we actually use right now.
	default:
		// Other commands are not implemented yet
		h.log.Info("Unhandled command %d from client %d", cmd, ctx.ID)
	}
}

// handleConnAuth answers a sub-connection identity handshake. Every sub-conn
// the client opens (screen, terminal, file, ...) sends a 512-byte
// ConnAuthPacket as its very first payload and blocks for up to 10 s waiting
// on our 256-byte ConnAuthAck. Without an OK reply the client closes the
// connection, so a missing ack here means nothing else can proceed.
//
// The handshake includes an HMAC signature field. The reference server
// treats verification failures as soft (logs and still allows commands),
// and the signing primitive lives in a vendored component out of scope
// for this server, so we always reply OK and let TOKEN_BITMAPINFO carry
// the device ID via offset 41 when the screen sub-conn proceeds.
func (h *MyHandler) handleConnAuth(ctx *connection.Context, data []byte) {
	// Pin the parent device's ClientID onto the sub-conn. Without this,
	// later 1-byte tokens (TOKEN_TERMINAL_START / TOKEN_SHELL_START) have
	// no way to identify which device they belong to — they carry no
	// clientID themselves. ConnAuthPacket layout has clientID at offset 1
	// (uint64 LE); see common/commands.h::ConnAuthPacket.
	if len(data) >= protocol.ConnAuthOffClientID+8 {
		clientID := binary.LittleEndian.Uint64(
			data[protocol.ConnAuthOffClientID : protocol.ConnAuthOffClientID+8])
		if clientID != 0 {
			// Sub-conns never go through handleLogin, so their ctx.Info
			// is otherwise empty. We only need ClientID for routing.
			info := ctx.GetInfo()
			info.ClientID = strconv.FormatUint(clientID, 10)
			ctx.SetInfo(info)
		}
	}

	ack := make([]byte, protocol.ConnAuthAckSize)
	ack[0] = protocol.TokenConnAuth
	ack[protocol.ConnAuthAckOffStatus] = protocol.ConnAuthStatusOK
	binary.LittleEndian.PutUint64(
		ack[protocol.ConnAuthAckOffServerTime:protocol.ConnAuthAckOffServerTime+8],
		uint64(time.Now().Unix()))
	if err := h.srv.Send(ctx, ack); err != nil {
		h.log.Error("ConnAuth ack send failed for conn=%d: %v", ctx.ID, err)
	}
}

// deviceIDOfSubConn resolves the parent device of a sub-conn from the
// ClientID pinned by handleConnAuth. Returns "" for the rare case of a
// legacy client that skipped ConnAuth (the Go server's only target is
// modern clients, so this is effectively a paranoia check).
func (h *MyHandler) deviceIDOfSubConn(ctx *connection.Context) string {
	return ctx.GetInfo().ClientID
}

// handleTerminalStart fires when the device's freshly-spawned shell
// sub-conn announces itself. TOKEN_TERMINAL_START (232) means PTY mode
// (Linux/macOS or Windows ConPTY); TOKEN_SHELL_START (128) means the
// legacy Windows cmd-pipe path. Both packets are 1-byte tokens — the
// device identity comes from ConnAuth's pinned ClientID.
//
// After binding we send:
//   - For PTY only: an initial CMD_TERMINAL_RESIZE 80x24 so the shell
//     doesn't render at the PTY default before the browser's first fit.
//     vim/htop look broken otherwise. The browser will follow up with a
//     real term_resize once xterm.js sizes the canvas.
//   - Always: COMMAND_NEXT to unblock the device's read thread (the
//     ConPTYManager ReadThread sits on m_hEventDlgOpen until then —
//     see client/ConPTYManager.cpp:259).
func (h *MyHandler) handleTerminalStart(ctx *connection.Context, isPTY bool) {
	devID := h.deviceIDOfSubConn(ctx)
	if devID == "" {
		h.log.Warn("terminal start with no clientID: conn=%d", ctx.ID)
		ctx.Close()
		return
	}
	if !h.hub.BindTerminalConn(devID, ctx, isPTY) {
		// No pending session — this is a stale sub-conn (e.g. browser
		// gave up and closed term_close already). Drop it.
		h.log.Warn("orphan terminal sub-conn: device=%s conn=%d isPTY=%v",
			devID, ctx.ID, isPTY)
		ctx.Close()
		return
	}

	if isPTY {
		if err := h.srv.Send(ctx, protocol.BuildTerminalResize(80, 24)); err != nil {
			h.log.Error("initial resize send failed: conn=%d: %v", ctx.ID, err)
		}
	}
	if err := h.srv.Send(ctx, []byte{protocol.CommandNext}); err != nil {
		h.log.Error("COMMAND_NEXT send failed on terminal: conn=%d: %v", ctx.ID, err)
	}
	h.log.Info("terminal bound: device=%s conn=%d isPTY=%v", devID, ctx.ID, isPTY)
}

// handleBitmapInfo is the first packet on a freshly-arrived screen
// sub-connection. Packet layout (after the command byte at data[0]):
//
//	[BITMAPINFOHEADER:40][clientID:8 uint64 LE][dlgID:8 uint64 LE][...]
//
// So clientID lives at data[41..49] and dlgID at data[49..57]. We use
// clientID (= MasterID) to bind this sub-context to its parent device.
func (h *MyHandler) handleBitmapInfo(ctx *connection.Context, data []byte) {
	if len(data) < 49 {
		h.log.Warn("TOKEN_BITMAPINFO from conn %d too short (%d bytes)", ctx.ID, len(data))
		return
	}
	clientID := uint64(data[41]) | uint64(data[42])<<8 | uint64(data[43])<<16 | uint64(data[44])<<24 |
		uint64(data[45])<<32 | uint64(data[46])<<40 | uint64(data[47])<<48 | uint64(data[48])<<56
	deviceID := strconv.FormatUint(clientID, 10)

	if !h.hub.BindScreenConn(deviceID, ctx) {
		// Device not registered — main login hasn't happened (or device just
		// went offline). Drop the orphan sub-conn rather than leak it.
		h.log.Warn("orphan screen sub-conn %d for unknown device %s; closing", ctx.ID, deviceID)
		ctx.Close()
		return
	}

	// BITMAPINFOHEADER starts at data[1]. biWidth at offset 4, biHeight at
	// offset 8 (both int32 LE). biHeight may be negative for top-down DIBs.
	width := int(int32(binary.LittleEndian.Uint32(data[5:9])))
	height := int(int32(binary.LittleEndian.Uint32(data[9:13])))
	if height < 0 {
		height = -height
	}

	h.log.Info("screen sub-conn bound: conn=%d device=%s resolution=%dx%d",
		ctx.ID, deviceID, width, height)
	h.hub.PublishResolution(deviceID, width, height)

	// Notify the client its "dialog is open" so it stops blocking in
	// Manager::WaitForDialogOpen (client/Manager.cpp:259). Without this
	// the client waits a full 8 s timeout before it begins streaming
	// real H264 frames via TOKEN_NEXTSCREEN. 32-byte packet matches the
	// C++ CScreenSpyDlg::SendNext layout:
	//   [0]=COMMAND_NEXT [1..9]=dlgID uint64 [9..13]=capabilities uint32
	//   [13..17]=scrollInterval int32 [17..32]=zero reserved
	// We don't need scroll-detect / a real dlgID, so leave them zero.
	nextCmd := make([]byte, 32)
	nextCmd[0] = protocol.CommandNext
	if err := h.srv.Send(ctx, nextCmd); err != nil {
		h.log.Error("COMMAND_NEXT send failed for conn=%d: %v", ctx.ID, err)
	}
}

// handleScreenFrame relays one TOKEN_FIRSTSCREEN / TOKEN_NEXTSCREEN packet
// to all browsers watching this device. The on-the-wire packet starts with
// the token byte then a small fixed header (algorithm, cursor pos, cursor
// index) before the H.264 NAL payload. The browser-facing WS packet uses
// the C++-compatible layout: [deviceID:4 LE][frameType:1][dataLen:4 LE][H264:N].
//
// alwaysKey=true is used for TOKEN_FIRSTSCREEN (always IDR by construction);
// TOKEN_NEXTSCREEN is keyframe iff the NAL stream contains a 5/7/8 unit.
func (h *MyHandler) handleScreenFrame(ctx *connection.Context, data []byte, alwaysKey bool) {
	deviceID := h.hub.ScreenDeviceID(ctx)
	if deviceID == "" {
		return // not a bound screen sub-conn — drop
	}
	// data[0] is the token; the 11-byte header sits at data[1..12].
	const skip = 1 + protocol.ScreenFrameHeaderLen
	if len(data) <= skip {
		return
	}
	// Cursor index lives at the last byte of the small per-frame header
	// (offset 1 + 1 + 8 = 10). Publish before the heavy frame work so the
	// browser sees cursor updates even if we end up dropping frames later.
	h.hub.PublishCursor(deviceID, data[10])

	h264 := data[skip:]
	isKey := alwaysKey || protocol.IsH264Keyframe(h264)

	// Build the WS packet exactly as the C++ ScreenSpyDlg does — the front-end
	// decoder reads these offsets directly.
	id64, _ := strconv.ParseUint(deviceID, 10, 64)
	idLow := uint32(id64)
	frameType := byte(0)
	if isKey {
		frameType = 1
	}
	dataLen := uint32(len(h264))

	packet := make([]byte, 9+len(h264))
	binary.LittleEndian.PutUint32(packet[0:4], idLow)
	packet[4] = frameType
	binary.LittleEndian.PutUint32(packet[5:9], dataLen)
	copy(packet[9:], h264)

	h.hub.PublishScreenFrame(deviceID, packet, isKey)
}

// handleLogin handles client login (TOKEN_LOGIN = 102)
func (h *MyHandler) handleLogin(ctx *connection.Context, data []byte) {
	info, err := protocol.ParseLoginInfo(data)
	if err != nil {
		h.log.Error("Failed to parse login info from client %d: %v", ctx.ID, err)
		return
	}

	// The device's unique ID lives in reserved field 16 (RES_CLIENT_ID) as a
	// decimal string of a uint64 — the same number the device later puts at
	// offset 41 of TOKEN_BITMAPINFO. Using szMasterID here is WRONG: it is a
	// compile-time MASTER_HASH constant shared by every binary built from
	// the same source, so all clients would collide in the hub.
	clientID := info.GetReservedField(protocol.ResFieldClientID)
	if clientID == "" || clientID == "0" {
		// Legacy fallback (very old clients that don't fill RES_CLIENT_ID).
		// MasterID is still preferable to a per-connection number because it
		// at least stays stable across reconnects of the same binary.
		clientID = info.MasterID
	}
	if clientID == "" {
		clientID = fmt.Sprintf("conn-%d", ctx.ID)
	}

	// Store client info
	reserved := info.ParseReserved()
	clientInfo := connection.ClientInfo{
		ClientID:     clientID,
		ComputerName: info.PCName,
		OS:           info.OsVerInfo,
		Version:      info.ModuleVersion,
		HasCamera:    info.WebCamExist,
		InstallTime:  info.StartTime,
	}

	// Parse additional info from reserved field
	if len(reserved) > protocol.ResFieldClientType {
		clientInfo.ClientType = info.GetReservedField(protocol.ResFieldClientType)
	}
	if len(reserved) > 2 {
		clientInfo.CPU = info.GetReservedField(2)
	}
	if len(reserved) > protocol.ResFieldFilePath {
		clientInfo.FilePath = info.GetReservedField(protocol.ResFieldFilePath)
	}
	if len(reserved) > protocol.ResFieldClientPubIP {
		clientInfo.IP = info.GetReservedField(protocol.ResFieldClientPubIP)
	}

	ctx.SetInfo(clientInfo)
	ctx.IsLoggedIn.Store(true)

	h.log.ClientEvent("online", ctx.ID, ctx.GetPeerIP(),
		"clientID", clientID,
		"computer", info.PCName,
		"os", info.OsVerInfo,
		"version", info.ModuleVersion,
		"path", clientInfo.FilePath,
	)

	// PCName carries "ComputerName/Group"; ModuleVersion carries "Version-Capability".
	// strings.Cut returns the full string as the head when the separator is
	// absent, which gives us the natural "no group / no capability" fallback.
	name, group, _ := strings.Cut(info.PCName, "/")
	version, capability, _ := strings.Cut(info.ModuleVersion, "-")

	// Client-reported geo string (RES_CLIENT_LOC).
	location := ""
	if len(reserved) > protocol.ResFieldClientLoc {
		location = info.GetReservedField(protocol.ResFieldClientLoc)
	}
	// RES_RESOLUTION is already formatted by the client as "N:W*H"
	// (see client/LoginServer.cpp:414). Pass through unchanged so the web
	// UI's device card renders it next to the version label.
	resolution := ""
	if len(reserved) > protocol.ResFieldResolution {
		resolution = info.GetReservedField(protocol.ResFieldResolution)
	}

	// Register with hub so the web side can list this device. Sub-connections
	// (screen / terminal etc.) reuse the MasterID and will overwrite this entry
	// harmlessly, but only the main login carries enough info to be useful here.
	h.hub.Register(&hub.Device{
		ID:          clientID,
		Name:        name,
		Group:       group,
		Version:     version,
		Capability:  capability,
		OS:          info.OsVerInfo,
		CPU:         clientInfo.CPU,
		FilePath:    clientInfo.FilePath,
		InstallTime: info.StartTime,
		Location:    location,
		Resolution:  resolution,
		PeerIP:      ctx.GetPeerIP(),
		PublicIP:    clientInfo.IP,
		ConnectedAt: time.Now(),
	}, ctx)

	// Push CMD_MASTERSETTING with a signature over "StartTime|ClientID".
	// The client's private FileUpload init verifies this before allowing
	// screen / file operations — without it the binary aborts itself.
	h.sendMasterSetting(ctx, info.StartTime, clientID)
}

// sendMasterSetting builds the 1001-byte CMD_MASTERSETTING reply and ships it
// down the main TCP connection. Most fields stay zeroed — only Signature
// matters today. The signer is one of:
//   - LocalSigner: HMAC directly with master key (operator's own deployment)
//   - RemoteSigner: HTTPS POST to operator's License Server (customer deployment)
//   - NoOpSigner: returns empty signature (free tier; client refuses screen/file ops)
//
// On signer error (License Server unreachable + no cache hit), we still ship
// a zeroed signature so the packet is well-formed; the client will retry on
// next reconnect.
func (h *MyHandler) sendMasterSetting(ctx *connection.Context, startTime, clientID string) {
	buf := make([]byte, 1+protocol.MasterSettingsSize)
	buf[0] = protocol.CmdMasterSetting

	// ReportInterval (int32 LE at struct offset 0, +1 for the cmd byte).
	// Sending 0 makes the client drop the active-window field of its
	// heartbeat, which kills the web UI's live activeWindow updates.
	binary.LittleEndian.PutUint32(
		buf[1:5],
		uint32(protocol.DefaultReportIntervalSec))

	sig, err := h.signer.Sign(startTime, clientID)
	if err != nil {
		h.log.Error("signer (%s) failed for clientID=%s: %v — sending zeroed signature",
			h.signer.Mode(), clientID, err)
	} else if sig == "" {
		// NoOpSigner path, or LocalSigner with empty master key — same effect.
		// Log only once per process via the startup banner; don't spam here.
	} else {
		const sigOffset = 1 + protocol.MasterSettingsOffSignature
		copy(buf[sigOffset:sigOffset+protocol.MasterSettingsSignatureLen], []byte(sig))
	}

	if err := h.srv.Send(ctx, buf); err != nil {
		h.log.Error("CMD_MASTERSETTING send failed for conn=%d: %v", ctx.ID, err)
	}
}

// handleAuth handles authorization request (TOKEN_AUTH = 100)
func (h *MyHandler) handleAuth(ctx *connection.Context, data []byte) {
	result := h.auth.Authenticate(data)
	info := ctx.GetInfo()

	authType := "V1"
	if result.IsV2 {
		authType = "V2"
	}

	if result.Valid {
		h.log.Info("Auth %s success: clientID=%s computer=%s ip=%s sn=%s passcode=%s",
			authType, info.ClientID, info.ComputerName, ctx.GetPeerIP(), result.SN, result.Passcode)
	} else {
		h.log.Warn("Auth %s failed: clientID=%s computer=%s ip=%s sn=%s passcode=%s",
			authType, info.ClientID, info.ComputerName, ctx.GetPeerIP(), result.SN, result.Passcode)
	}

	// Build and send response
	resp := h.auth.BuildResponse(result)
	if err := h.srv.Send(ctx, resp); err != nil {
		h.log.Error("Failed to send auth response to client %d: %v", ctx.ID, err)
	}
}

// handleHeartbeat handles heartbeat from client (TOKEN_HEARTBEAT = 101)
// Heartbeat structure (after command byte):
//   - offset 0: Time (8 bytes, uint64)
//   - offset 8: ActiveWnd (512 bytes)
//   - offset 520: Ping (4 bytes, int)
//   - offset 524: HasSoftware (4 bytes, int)
//   - offset 528: SN (20 bytes)
//   - offset 548: Passcode (44 bytes)
//   - offset 592: PwdHmac (8 bytes, uint64)
//
// HeartbeatACK structure:
//   - offset 0: Time (8 bytes, uint64)
//   - offset 8: Authorized (1 byte, char)
//   - offset 9: Reserved (23 bytes)
func (h *MyHandler) handleHeartbeat(ctx *connection.Context, data []byte) {

	// Parse Time from heartbeat request (offset 1, 8 bytes)
	var hbTime uint64
	if len(data) >= 9 {
		hbTime = uint64(data[1]) | uint64(data[2])<<8 | uint64(data[3])<<16 | uint64(data[4])<<24 |
			uint64(data[5])<<32 | uint64(data[6])<<40 | uint64(data[7])<<48 | uint64(data[8])<<56
	}

	// Forward live fields (ActiveWnd + Ping) to the hub so the web UI can
	// display current latency and foreground window per device. Skip until
	// login has happened — the hub is keyed by MasterID, which only exists
	// post-login.
	if info := ctx.GetInfo(); info.ClientID != "" {
		var rtt int32
		var activeWindow string
		// ActiveWnd at data[9..521] is a 512-byte NUL-padded string. Encoding
		// depends on the client: new clients advertise CLIENT_CAP_UTF8 (bit
		// 0x0002 in the moduleVersion hex tail) and ship UTF-8 directly;
		// legacy Windows clients still use CP936 (GBK). Decoding UTF-8 bytes
		// as GBK turns Chinese characters into mojibake — see the matching
		// C++ guard at server/2015Remote/WebService.cpp:1530.
		if len(data) >= 9+512 {
			activeWindow = protocol.DecodeClientString(
				data[9:9+512],
				h.hub.Capability(info.ClientID),
				info.ClientType,
			)
		}
		// Ping at data[521..525] is a little-endian int32.
		if len(data) >= 525 {
			rtt = int32(uint32(data[521]) | uint32(data[522])<<8 |
				uint32(data[523])<<16 | uint32(data[524])<<24)
		}
		h.hub.UpdateLive(info.ClientID, int(rtt), activeWindow)
	}

	// Authenticate heartbeat if it contains authorization info
	// data[1:] skips the command byte to get the raw Heartbeat structure
	var authorized byte = 0
	if len(data) > 1 {
		authResult := h.auth.AuthenticateHeartbeat(data[1:])
		if authResult.Authorized {
			authorized = 2 // Auth by admin
			// Log authorization success (only log once per connection to avoid spam)
			if !ctx.IsAuthorized.Load() {
				ctx.IsAuthorized.Store(true)
				info := ctx.GetInfo()
				if authResult.IsV2 {
					// V2 authorization
					h.log.Info("Heartbeat auth V2 success: clientID=%s computer=%s ip=%s sn=%s passcode=%s",
						info.ClientID, info.ComputerName, ctx.GetPeerIP(), authResult.SN, authResult.Passcode)
				} else {
					// V1 authorization
					h.log.Info("Heartbeat auth V1 success: clientID=%s computer=%s ip=%s sn=%s passcode=%s pwdHmac=%d",
						info.ClientID, info.ComputerName, ctx.GetPeerIP(), authResult.SN, authResult.Passcode, authResult.PwdHmac)
				}
			}
		}
	}

	// Build HeartbeatACK response: CMD_HEARTBEAT_ACK(1) + HeartbeatACK(32)
	resp := make([]byte, 33)
	resp[0] = protocol.CommandHeartbeat // CMD_HEARTBEAT_ACK = 216
	// Time at offset 1 (8 bytes, little-endian)
	resp[1] = byte(hbTime)
	resp[2] = byte(hbTime >> 8)
	resp[3] = byte(hbTime >> 16)
	resp[4] = byte(hbTime >> 24)
	resp[5] = byte(hbTime >> 32)
	resp[6] = byte(hbTime >> 40)
	resp[7] = byte(hbTime >> 48)
	resp[8] = byte(hbTime >> 56)
	// Authorized at offset 9 (1 byte)
	resp[9] = authorized
	// Reserved[23] at offset 10 is already zero

	if err := h.srv.Send(ctx, resp); err != nil {
		h.log.Error("Failed to send heartbeat ACK to client %d: %v", ctx.ID, err)
	}
}

// parsePorts parses a semicolon-separated port string and returns port numbers
func parsePorts(portStr string) ([]int, error) {
	var ports []int
	parts := strings.Split(portStr, ";")
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		port, err := strconv.Atoi(p)
		if err != nil {
			return nil, fmt.Errorf("invalid port %q: %v", p, err)
		}
		if port < 1 || port > 65535 {
			return nil, fmt.Errorf("port %d out of range (1-65535)", port)
		}
		ports = append(ports, port)
	}
	if len(ports) == 0 {
		return nil, fmt.Errorf("no valid ports specified")
	}
	return ports, nil
}

// splitCSV splits a comma-separated env-var value into trimmed, non-empty
// entries. Returns nil for an empty input so callers can keep the natural
// "no value → no restriction" semantics with a single nil check.
func splitCSV(s string) []string {
	if s == "" {
		return nil
	}
	parts := strings.Split(s, ",")
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p != "" {
			out = append(out, p)
		}
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

func main() {
	// Parse command line flags
	portStr := flag.String("port", "6543", "Server listen ports (semicolon-separated, e.g. 6543;6544;6545)")
	flag.StringVar(portStr, "p", "6543", "Server listen ports (shorthand)")
	httpPort := flag.Int("http-port", 8080, "HTTP server port for web UI (0 to disable)")
	noConsole := flag.Bool("no-console", false, "Disable console output (for daemon mode)")
	flag.Parse()

	// Parse ports
	ports, err := parsePorts(*portStr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	// Create logger with file output
	logCfg := logger.DefaultConfig()
	logCfg.Level = logger.LevelDebug
	logCfg.Console = !*noConsole
	logCfg.File = "logs/server.log"
	logCfg.MaxSize = 100   // 100 MB
	logCfg.MaxBackups = 10 // keep 10 old files
	logCfg.MaxAge = 30     // 30 days
	logCfg.Compress = true

	log := logger.New(logCfg)

	// Track env vars where we fell back to a built-in default. Printed once
	// at the end of startup so the operator sees what's in effect — vars the
	// operator explicitly set are NOT listed (they already know their value).
	type defaultedEnv struct{ name, value string }
	var defaultsUsed []defaultedEnv
	rememberDefault := func(name, value string) {
		defaultsUsed = append(defaultsUsed, defaultedEnv{name, value})
	}

	// Create auth config
	authCfg := auth.DefaultConfig()
	// PwdHash can be set from environment or config file
	authCfg.PwdHash = os.Getenv("YAMA_PWDHASH")
	if authCfg.PwdHash == "" {
		// Default placeholder - should be configured in production
		authCfg.PwdHash = "61f04dd637a74ee34493fc1025de2c131022536da751c29e3ff4e9024d8eec43"
		rememberDefault("YAMA_PWDHASH", authCfg.PwdHash)
	}
	authCfg.SuperPass = os.Getenv("YAMA_PWD")

	// Create authenticator (shared by all servers)
	authenticator := auth.New(authCfg)

	// Shared device registry — every TCP handler reports devices into it,
	// the HTTP server reads from it.
	deviceHub := hub.New()

	// Build the CMD_MASTERSETTING signer based on env vars:
	//   - YAMA_SIGN_PASSWORD set → LocalSigner (operator's own deployment;
	//     HMAC master key lives here)
	//   - YAMA_LICENSE_SERVER + YAMA_LICENSE_TOKEN set → RemoteSigner
	//     (customer deployment; never sees the master key, fetches signatures
	//     from operator's License Server with 24h cache)
	//   - neither → NoOpSigner (free tier; client refuses screen/file ops
	//     but device list still works)
	signer, mode, err := licensing.NewFromEnv(log)
	if err != nil {
		log.Fatal("Failed to initialize signer: %v", err)
	}
	// signer.Close() is called explicitly after the License Server HTTP is
	// drained at shutdown — sequencing matters because an in-flight
	// handleSign on the HTTP path needs a live signer to complete.
	switch mode {
	case licensing.ModeLocal:
		log.Info("Signer mode: LOCAL (operator deployment, master key held in-process)")
	case licensing.ModeRemote:
		log.Info("Signer mode: REMOTE (customer deployment, %s=%s)",
			licensing.EnvLicenseServer, os.Getenv(licensing.EnvLicenseServer))
	case licensing.ModeNoOp:
		log.Warn("Signer mode: NOOP (no licensing configured; screen/file features disabled on clients)")
	}

	// If the operator also wants this LocalSigner deployment to serve as
	// License Server for RemoteSigner customers, YAMA_LICENSE_PUBLIC_KEY +
	// YAMA_LICENSE_HTTP_ADDR enable it. Off by default.
	licSrv, licAddr, err := licensing.LicenseServerFromEnv(signer, log)
	if err != nil {
		log.Fatal("Failed to initialize License Server: %v", err)
	}

	// Web user authenticator. Bootstrap admin from env var YAMA_WEB_ADMIN_PASS;
	// if unset, fall back to YAMA_PWD (same secret the TCP authorization uses)
	// so a single password env var is enough to bring up the whole stack.
	// If neither is set we use a hard-coded "admin" default so the web UI is
	// usable out of the box — the startup banner surfaces this so operators
	// know to override it in any non-dev deployment.
	const defaultWebAdminPass = "admin"
	webAuth := wsauth.New()
	adminPass := os.Getenv("YAMA_WEB_ADMIN_PASS")
	if adminPass == "" {
		adminPass = os.Getenv("YAMA_PWD")
	}
	usingDefaultWebPass := false
	if adminPass == "" {
		adminPass = defaultWebAdminPass
		rememberDefault("YAMA_WEB_ADMIN_PASS", defaultWebAdminPass)
		usingDefaultWebPass = true
	}
	webAuth.AddAdminFromPlainPassword("admin", adminPass)
	log.Info("Web admin user configured")

	// Persistent users live in users.json next to the binary's working dir
	// — same default the C++ WebService uses (m_ConfigDir + "users.json").
	// Loading is best-effort: a missing file means "no extra users yet".
	usersFile := os.Getenv("YAMA_USERS_FILE")
	if usersFile == "" {
		usersFile = "users.json"
		rememberDefault("YAMA_USERS_FILE", usersFile)
	}
	if err := webAuth.SetUsersFile(usersFile); err != nil {
		log.Warn("Failed to load %s: %v (continuing with admin only)", usersFile, err)
	}

	// Create servers for each port
	var servers []*server.Server
	for _, port := range ports {
		config := server.DefaultConfig()
		config.Port = port
		config.MaxConnections = 9999

		srv := server.New(config)
		srv.SetLogger(log.WithPrefix(fmt.Sprintf("Server:%d", port)))

		// Create handler for this server
		handler := &MyHandler{
			log:    log.WithPrefix(fmt.Sprintf("Handler:%d", port)),
			auth:   authenticator,
			srv:    srv,
			hub:    deviceHub,
			signer: signer,
		}
		srv.SetHandler(handler)

		servers = append(servers, srv)
	}

	// Wire the hub's outbound sender once all TCP servers exist. Any server's
	// Send method will do — the per-connection encoder uses ctx-local state
	// and is independent of which server originally accepted the connection.
	if len(servers) > 0 {
		s := servers[0]
		deviceHub.SetSender(func(ctx *connection.Context, data []byte) error {
			return s.Send(ctx, data)
		})
	}

	// Start all TCP servers
	for _, srv := range servers {
		if err := srv.Start(); err != nil {
			log.Fatal("Failed to start server: %v", err)
		}
	}

	// Web-UI hardening knobs for public-HTTPS deployment.
	//
	// YAMA_WEB_ALLOWED_ORIGINS: comma-separated Origin allowlist (e.g.
	//   "https://yama.example.com,https://yama-mobile.example.com").
	//   Empty (default) → only same-origin WS upgrades accepted, which
	//   is correct when the web UI and WS endpoint share a host.
	//
	// Login rate limits are hard-coded at sensible defaults for the
	// small-user web UI: 10 attempts / minute per IP, 5 / 15 min per
	// username. The handler also injects a 250 ms delay on every failure
	// so online brute force is impractical even within budget.
	allowedOrigins := splitCSV(os.Getenv("YAMA_WEB_ALLOWED_ORIGINS"))
	trustProxy := os.Getenv("YAMA_WEB_TRUST_PROXY") == "1"
	if trustProxy {
		log.Info("Trusting X-Forwarded-For for client IP — make sure a reverse proxy is in front")
	}
	webCfg := web.Config{
		AllowedOrigins:    allowedOrigins,
		LoginIPLimit:      wsauth.NewRateLimiter(10, time.Minute),
		LoginUserLimit:    wsauth.NewRateLimiter(5, 15*time.Minute),
		TrustForwardedFor: trustProxy,
	}

	// Start HTTP server for web UI. Hub gives it read-only access to the
	// device registry; the authenticator owns user accounts and session tokens.
	httpSrv := web.New(*httpPort, log.WithPrefix("Web"), deviceHub, webAuth).
		WithConfig(webCfg)
	if err := httpSrv.Start(); err != nil {
		log.Fatal("Failed to start HTTP server: %v", err)
	}

	// Optionally serve the License Server HTTP endpoints in this process.
	// Only active when the operator set YAMA_LICENSE_PUBLIC_KEY +
	// YAMA_LICENSE_HTTP_ADDR. We serve plain HTTP — operator should put
	// nginx/Caddy in front for TLS. (RemoteSigner customers connect to
	// the public URL configured via YAMA_LICENSE_SERVER on their side.)
	//
	// Timeouts cap Slowloris / FD-exhaustion attacks. Values are generous
	// enough for a slow public link (TLS-terminating proxy in front, real
	// customer round-trips at ~hundreds of ms) but tight enough that a
	// trickle-byte attacker cannot pin a connection indefinitely.
	var licenseHTTP *http.Server
	if licSrv != nil {
		licenseHTTP = &http.Server{
			Addr:              licAddr,
			Handler:           licSrv.Handler(),
			ReadHeaderTimeout: 5 * time.Second,
			ReadTimeout:       15 * time.Second,
			WriteTimeout:      15 * time.Second,
			IdleTimeout:       60 * time.Second,
			MaxHeaderBytes:    16 << 10,
		}
		go func() {
			log.Info("License Server listening on %s (POST /license/sign, /license/heartbeat)", licAddr)
			if err := licenseHTTP.ListenAndServe(); err != nil && err != http.ErrServerClosed {
				log.Error("License Server stopped: %v", err)
			}
		}()
	}

	fmt.Printf("Server started on port(s): %v\n", ports)
	if *httpPort != 0 {
		fmt.Printf("Web UI on http://localhost:%d/\n", *httpPort)
		if usingDefaultWebPass {
			fmt.Printf("    Default login: admin / %s   (set YAMA_WEB_ADMIN_PASS to override)\n",
				defaultWebAdminPass)
		}
	}
	if licenseHTTP != nil {
		fmt.Printf("License Server on http://%s/license/{sign,heartbeat}\n", licAddr)
	}
	if len(defaultsUsed) > 0 {
		fmt.Println()
		fmt.Println("[!] Using built-in defaults (set the env var to override):")
		for _, d := range defaultsUsed {
			fmt.Printf("    %s = %s\n", d.name, d.value)
		}
	}
	fmt.Println("Logs are written to: logs/server.log")
	fmt.Println("Press Ctrl+C to stop...")

	// Wait for interrupt signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	fmt.Println("\nShutting down...")
	// Order matters: drain License Server HTTP first so no handleSign is
	// mid-flight; THEN close the signer (which may release HTTP keepalives
	// in RemoteSigner mode, or be a no-op for LocalSigner/NoOp).
	if licenseHTTP != nil {
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		_ = licenseHTTP.Shutdown(shutdownCtx)
		cancel()
	}
	if err := signer.Close(); err != nil {
		log.Warn("signer Close: %v", err)
	}
	httpSrv.Stop()
	for _, srv := range servers {
		srv.Stop()
	}
	fmt.Println("Server stopped")
}