package main

import (
	"bytes"
	"crypto/sha256"
	"crypto/subtle"
	"crypto/tls"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/racingmars/go3270"
	"golang.org/x/crypto/bcrypt"
)

func ensureSlotCapacity(s *UserMasterSession, slotIdx int) {
	for len(s.Slots) <= slotIdx {
		s.Slots = append(s.Slots, nil)
	}
}

var (
	GlobalConfig     *Config
	ConfiguredSystem []TargetSystem

	configFile  string
	listen      string
	certFile    string
	keyFile     string
	genPassword string
	publicFlag  bool

	serverListener   net.Listener
	shutdownFunc     = func() { os.Exit(0) }
	adminNoticeSleep = 1500 * time.Millisecond

	registryMu   sync.Mutex
	userSessions = make(map[string]*UserMasterSession)

	telnetEOR            = []byte{IAC, EOR}
	telnetDoTimingMark   = []byte{IAC, DO, OPT_TIME_MARK}
	telnetWillTimingMark = []byte{IAC, WILL, OPT_TIME_MARK}
)

func init() {
	flag.StringVar(&configFile, "c", "config.toml", "Path to toml config file")
	flag.StringVar(&listen, "l", "0.0.0.0:2323", "Listen address")
	flag.StringVar(&certFile, "t", "", "Path to TLS certificate file (PEM)")
	flag.StringVar(&keyFile, "k", "", "Path to TLS private key file (PEM)")
	flag.StringVar(&genPassword, "p", "", "Generate bcrypt hash for password and exit")
	flag.BoolVar(&publicFlag, "public", false, "Allow public access")
	flag.BoolVar(&publicFlag, "P", false, "Allow public access")
}

func main() {
	flag.Parse()

	var hasGenPassword, hasListen bool
	flag.CommandLine.Visit(func(f *flag.Flag) {
		switch f.Name {
		case "p":
			hasGenPassword = true
		case "l":
			hasListen = true
		}
	})

	if hasGenPassword {
		if len(strings.TrimSpace(genPassword)) < 3 {
			fmt.Fprintln(os.Stderr, "Error: Password cannot be less than 3 characters")
			os.Exit(1)
		}
		hash, err := bcrypt.GenerateFromPassword([]byte(genPassword), bcrypt.DefaultCost)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error generating bcrypt hash: %v\n", err)
			os.Exit(1)
		}
		fmt.Println(string(hash))
		return
	}

	cfg, err := loadHosts(configFile)
	if err != nil {
		log.Fatalf("Configuration error: %v", err)
	}

	if hasListen {
		cfg.Server.BindAddress = listen
	}

	enabledTLS, certPath, keyPath, err := resolveTLSConfig(certFile, keyFile, &cfg.Server)
	if err != nil {
		log.Fatalf("TLS configuration error: %v", err)
	}

	if publicFlag {
		cfg.Public = true
		cfg.Server.Public = true
	}

	GlobalConfig = cfg
	ConfiguredSystem = cfg.Systems

	applyCodepage(cfg.Server.Codepage)

	log.Printf("Loaded configuration from %s (%d target systems, Codepage: %s)",
		configFile, len(ConfiguredSystem), cfg.Server.Codepage)

	if len(cfg.Users) > 0 {
		if isPublic() {
			log.Printf("Public access enabled (%d users configured)", len(cfg.Users))
		} else {
			log.Printf("User whitelist enabled (%d users configured)", len(cfg.Users))
		}
	} else {
		log.Printf("User whitelist not configured (public access mode)")
	}

	for _, sys := range ConfiguredSystem {
		tlsMode := "Plaintext"
		if sys.TLS {
			tlsMode = "TLS"
		}
		log.Printf("  + %s -> %s (%s)", sys.Name, sys.HostPort, tlsMode)
	}

	bindAddr := cfg.Server.BindAddress
	listener, err := startServerListener(bindAddr, enabledTLS, certPath, keyPath)
	if err != nil {
		log.Fatalf("Failed to bind on %s: %v", bindAddr, err)
	}
	serverListener = listener
	defer listener.Close()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
	go func() {
		sig, ok := <-sigCh
		if !ok {
			return
		}
		log.Printf("Received signal %v: draining sessions and shutting down...", sig)
		for _, sess := range drainAllSessions() {
			terminateSession(sess, "Server shutting down.")
		}
		_ = listener.Close()
	}()
	defer signal.Stop(sigCh)

	if enabledTLS {
		log.Printf("Session Broker online and listening on %s (TLS)", bindAddr)
	} else {
		log.Printf("Session Broker online and listening on %s", bindAddr)
	}

	var tempDelay time.Duration
	for {
		conn, err := listener.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				log.Printf("Listener closed, stopping main loop: %v", err)
				return
			}
			if tempDelay == 0 {
				tempDelay = 5 * time.Millisecond
			} else {
				tempDelay *= 2
			}
			if maxDelay := 1 * time.Second; tempDelay > maxDelay {
				tempDelay = maxDelay
			}
			log.Printf("Accept error: %v; retrying in %v", err, tempDelay)
			time.Sleep(tempDelay)
			continue
		}
		tempDelay = 0
		log.Printf("[%s] Accepted new inbound connection", conn.RemoteAddr())
		go handleNewConnection(conn)
	}
}

func resolveTLSConfig(certFlag, keyFlag string, cfgServer *ServerConfig) (bool, string, string, error) {
	cert := strings.TrimSpace(certFlag)
	key := strings.TrimSpace(keyFlag)

	tlsEnabled := key != "" || cert != ""

	if tlsEnabled && (key == "" || cert == "") {
		return false, "", "", fmt.Errorf("-t and -k require both of each other")
	}

	if !tlsEnabled && cfgServer != nil && cfgServer.TLS {
		tlsEnabled = true
		cert = strings.TrimSpace(cfgServer.CertFile)
		key = strings.TrimSpace(cfgServer.KeyFile)
		if cert == "" || key == "" {
			return false, "", "", fmt.Errorf("both certificate and key are required when TLS is enabled")
		}
	}

	return tlsEnabled, cert, key, nil
}

func startServerListener(bindAddr string, enabledTLS bool, certPath, keyPath string) (net.Listener, error) {
	if enabledTLS {
		cert, err := tls.LoadX509KeyPair(certPath, keyPath)
		if err != nil {
			return nil, fmt.Errorf("failed to load TLS certificate/key (%s, %s): %w", certPath, keyPath, err)
		}
		tlsConfig := &tls.Config{
			Certificates: []tls.Certificate{cert},
		}
		return tls.Listen("tcp4", bindAddr, tlsConfig)
	}
	return net.Listen("tcp4", bindAddr)
}

func computeSessionKey(username, password string) string {
	h := sha256.New()
	h.Write([]byte(strings.TrimSpace(username)))
	h.Write([]byte(":"))
	h.Write([]byte(password))
	return hex.EncodeToString(h.Sum(nil))
}

func isPublic() bool {
	return publicFlag || (GlobalConfig != nil && (GlobalConfig.Public || GlobalConfig.Server.Public))
}

func authenticateAndAuthorize(username, password string) AuthResult {
	if GlobalConfig == nil || len(GlobalConfig.Users) == 0 {
		return AuthResult{OK: true, Reason: "", IsAdmin: false}
	}

	trimmedUser := strings.TrimSpace(username)
	var matchedUser *UserConfig
	for i := range GlobalConfig.Users {
		if strings.TrimSpace(GlobalConfig.Users[i].Username) == trimmedUser {
			matchedUser = &GlobalConfig.Users[i]
			break
		}
	}

	if matchedUser != nil {
		if hash := matchedUser.GetHash(); hash != "" {
			if strings.HasPrefix(hash, "$2") {
				if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)); err == nil {
					return AuthResult{OK: true, Reason: "", IsAdmin: matchedUser.Admin}
				}
			} else if subtle.ConstantTimeCompare([]byte(hash), []byte(password)) == 1 {
				log.Printf("[SECURITY WARNING] User '%s' authenticated with plaintext password in config. Use -p to hash it.", username)
				return AuthResult{OK: true, Reason: "", IsAdmin: matchedUser.Admin}
			}
		}
	}

	if isPublic() {
		return AuthResult{OK: true, Reason: "", IsAdmin: false}
	}
	return AuthResult{OK: false, Reason: fmt.Sprintf("user '%s' not authorized", username)}
}

func applyCodepage(cp string) {
	switch strings.ToLower(strings.TrimSpace(cp)) {
	case "1047", "cp1047", "ibm1047":
		go3270.SetCodepage(go3270.Codepage1047())
	case "bracket", "brackets", "cpbracket":
		go3270.SetCodepage(go3270.CodepageBracket())
	case "500", "cp500":
		go3270.SetCodepage(go3270.Codepage500())
	case "273", "cp273":
		go3270.SetCodepage(go3270.Codepage273())
	case "1140", "cp1140":
		go3270.SetCodepage(go3270.Codepage1140())
	default:
		go3270.SetCodepage(go3270.Codepage037())
	}
}

func loadHosts(configPath string) (*Config, error) {
	data, err := os.ReadFile(configPath)
	if err != nil {
		return nil, fmt.Errorf("could not read %s: %w", configPath, err)
	}

	cfg := Config{
		Server: ServerConfig{
			BindAddress: listen,
			Codepage:    "037",
		},
		Login: LoginConfig{
			Title:          "SESSMAN3270 SESSION BROKER",
			Subtitle:       "Enter your Username and Password to begin or resume a session.",
			PromptUsername: "Username ===>",
			PromptPassword: "Password ===>",
			Instructions:   "[Enter] Submit   [Clear] Reset ",
		},
		Menu: MenuConfig{
			Title:        "SESSMAN3270 CENTRAL SESSION MANAGER",
			PromptSelect: "Selection ===>",
			Footer:       "[PA3] Return to Menu  [PF3] Disconnect",
		},
	}

	if _, err := toml.Decode(string(data), &cfg); err != nil {
		return nil, fmt.Errorf("failed to parse TOML from %s: %w", configPath, err)
	}

	if len(cfg.Systems) == 0 {
		return nil, fmt.Errorf("no systems defined in %s", configPath)
	}

	return &cfg, nil
}

func cleanupSessionIfEmpty(s *UserMasterSession) {
	if s == nil || s.SessionKey == "" {
		return
	}
	registryMu.Lock()
	defer registryMu.Unlock()

	s.Mu.Lock()
	shouldDelete := s.IsDisconnected && s.countActiveSlotsLocked() == 0
	key := s.SessionKey
	s.Mu.Unlock()

	if shouldDelete && userSessions[key] == s {
		delete(userSessions, key)
	}
}

func handleNewConnection(conn net.Conn) {
	if tlsConn, ok := conn.(*tls.Conn); ok {
		_ = tlsConn.SetDeadline(time.Now().Add(5 * time.Second))
		if err := tlsConn.Handshake(); err != nil {
			log.Printf("[%s] TLS handshake failed: %v", conn.RemoteAddr(), err)
			_ = conn.Close()
			return
		}
		_ = tlsConn.SetDeadline(time.Time{})
	}

	devInfo, err := go3270.NegotiateTelnet(conn)
	if err != nil {
		log.Printf("[%s] Telnet negotiation failed: %v", conn.RemoteAddr(), err)
		_ = conn.Close()
		return
	}
	termType := ResolveTerminalType(devInfo.TerminalType(), devInfo)
	log.Printf("[%s] Telnet negotiation completed (Terminal: %s)", conn.RemoteAddr(), termType)

	creds, err := promptLogin(conn, devInfo)
	if err != nil {
		log.Printf("[%s] Client disconnected during login: %v", conn.RemoteAddr(), err)
		_ = conn.Close()
		return
	}

	sessionKey := creds.SessionKey
	shortTag := sessionKey[:8]

	authRes := authenticateAndAuthorize(creds.Username, creds.Password)
	if !authRes.OK {
		log.Printf("[%s] Authentication failed for user '%s': %s", conn.RemoteAddr(), creds.Username, authRes.Reason)
		send3270Message(conn, fmt.Sprintf("Authentication failed: %s.", authRes.Reason), devInfo.Codepage())
		time.Sleep(adminNoticeSleep)
		_ = conn.Close()
		return
	}

	isAdmin := authRes.IsAdmin

	registryMu.Lock()
	userSession, exists := userSessions[sessionKey]
	if !exists {
		userSession = &UserMasterSession{
			Username:   creds.Username,
			SessionKey: sessionKey,
			DisplayTag: shortTag,
			Slots:      make([]*SessionSlot, 0),
		}
		userSessions[sessionKey] = userSession
	}
	registryMu.Unlock()

	userSession.Mu.Lock()
	oldConn := userSession.ClientConn
	userSession.ClientConn = conn
	userSession.Username = creds.Username
	userSession.TerminalType = termType
	userSession.DevInfo = devInfo
	userSession.IsDisconnected = false
	userSession.ActiveSlotIdx = -1
	userSession.IsAdmin = isAdmin
	activeSlots := userSession.countActiveSlotsLocked()
	userSession.Mu.Unlock()

	if !exists {
		log.Printf("[%s] '%s' authenticated -> allocated new session %s (Admin: %v)", conn.RemoteAddr(), creds.Username, shortTag, isAdmin)
	} else {
		log.Printf("[%s] '%s' authenticated -> resumed existing session %s (%d active slots, Admin: %v)",
			conn.RemoteAddr(), creds.Username, shortTag, activeSlots, isAdmin)
	}

	if oldConn != nil {
		log.Printf("[%s] '%s' Closing previous client connection socket", shortTag, creds.Username)
		_ = oldConn.Close()
	}

	if activeSlots > 0 {
		handleResumeSession(userSession, conn)
		return
	}

	runMenuLoop(userSession, conn)
}

func dialUpstream(target TargetSystem) (net.Conn, error) {
	dialer := &net.Dialer{Timeout: 10 * time.Second}

	if target.TLS {
		host, _, err := net.SplitHostPort(target.HostPort)
		if err != nil {
			host = target.HostPort
		}
		tlsConfig := &tls.Config{
			InsecureSkipVerify: target.InsecureSkipVerify,
			ServerName:         host,
		}
		return tls.DialWithDialer(dialer, "tcp", target.HostPort, tlsConfig)
	}

	return dialer.Dial("tcp", target.HostPort)
}

func sendQueryReply(s *UserMasterSession, upConn net.Conn, isTN3270E, isBound bool) {
	s.Mu.Lock()
	termType := s.TerminalType
	if termType == "" {
		termType = ResolveTerminalType(termType, s.DevInfo)
	}
	rows, cols, _ := GetTerminalMatrix(termType, s.DevInfo)
	s.Mu.Unlock()
	isColor := strings.Contains(termType, "3279")
	reply := buildQueryReply(rows, cols, isColor, isTN3270E, isBound)
	_, _ = upConn.Write(reply)
}

func ensureSlotConnected(s *UserMasterSession, slotIdx int, background bool, clientConn net.Conn) *SessionSlot {
	s.Mu.Lock()
	if s.ClientConn != clientConn {
		s.Mu.Unlock()
		return nil
	}

	username, displayTag := s.Username, s.DisplayTag

	if slotIdx < len(s.Slots) && s.Slots[slotIdx] != nil {
		slot := s.Slots[slotIdx]
		if background {
			log.Printf("[%s] '%s' Slot [%d] %s already active in background", displayTag, username, slotIdx+1, slot.System.Name)
		} else {
			log.Printf("[%s] '%s' Resuming existing slot [%d] %s", displayTag, username, slotIdx+1, slot.System.Name)
		}
		s.Mu.Unlock()
		return slot
	}

	if slotIdx < 0 || slotIdx >= len(ConfiguredSystem) {
		s.Mu.Unlock()
		return nil
	}

	termType := ResolveTerminalType(s.TerminalType, s.DevInfo)
	_, _, maxBuf := GetTerminalMatrix(termType, s.DevInfo)
	cp := s.codepageLocked()
	s.Mu.Unlock()

	target := ConfiguredSystem[slotIdx]
	mode := "TCP"
	if target.TLS {
		mode = "TLS"
	}
	if background {
		log.Printf("[%s] '%s' Background dialing %s (%s, %s, Terminal: %s)...", displayTag, username, target.Name, target.HostPort, mode, termType)
	} else {
		log.Printf("[%s] '%s' Dialing %s (%s, %s, Terminal: %s)...", displayTag, username, target.Name, target.HostPort, mode, termType)
	}

	upConn, err := dialUpstream(target)
	if err != nil {
		log.Printf("[%s] '%s' Failed to connect to %s (%s): %v", displayTag, username, target.Name, target.HostPort, err)
		send3270Message(clientConn, fmt.Sprintf("Error dialing %s over %s: %v", target.Name, mode, err), cp)
		time.Sleep(adminNoticeSleep)
		return nil
	}

	initData, isTN3270E, isBound, err := negotiateUpstreamTelnet(upConn, termType)
	if err != nil {
		_ = upConn.Close()
		log.Printf("[%s] '%s' Upstream telnet negotiation failed for %s: %v", displayTag, username, target.Name, err)
		send3270Message(clientConn, fmt.Sprintf("Error negotiating with %s: %v", target.Name, err), cp)
		time.Sleep(adminNoticeSleep)
		return nil
	}

	if background {
		log.Printf("[%s] '%s' Backgrounded connection established to %s (%s, TN3270E: %v, Terminal: %s)", displayTag, username, target.Name, target.HostPort, isTN3270E, termType)
		if is3270Query(initData) {
			sendQueryReply(s, upConn, isTN3270E, isBound)
		}
	} else {
		log.Printf("[%s] '%s' Upstream connection to %s (%s)", displayTag, username, target.Name, target.HostPort)
	}

	s.Mu.Lock()
	defer s.Mu.Unlock()
	if s.ClientConn != clientConn {
		_ = upConn.Close()
		return nil
	}

	ensureSlotCapacity(s, slotIdx)

	if s.Slots[slotIdx] != nil {
		_ = upConn.Close()
		return s.Slots[slotIdx]
	}

	slot := &SessionSlot{
		System:        target,
		Upstream:      upConn,
		LastScreen:    initData,
		IsTN3270E:     isTN3270E,
		IsBound:       isBound,
		MaxBufferSize: maxBuf,
	}
	s.Slots[slotIdx] = slot
	go upstreamReader(s, slot, slotIdx)
	return slot
}

func connectOrResumeSlot(s *UserMasterSession, slotIdx int, clientConn net.Conn) {
	slot := ensureSlotConnected(s, slotIdx, false, clientConn)
	if slot == nil {
		return
	}

	s.Mu.Lock()
	if s.ClientConn != clientConn {
		s.Mu.Unlock()
		return
	}
	s.ActiveSlotIdx = slotIdx
	s.IsDisconnected = false
	s.Mu.Unlock()

	defer func() {
		s.Mu.Lock()
		if s.ClientConn == clientConn {
			s.ActiveSlotIdx = -1
		}
		s.Mu.Unlock()
	}()

	slot.Mu.Lock()
	slot.ScreenChanged = false
	slot.IsHeld = false
	if len(slot.LastScreen) > 0 {
		_, _ = clientConn.Write(slot.LastScreen)
	}
	slot.Mu.Unlock()

	bridgeClientToUpstream(s, slot, clientConn)
}

func cleanCloseUpstream(up net.Conn, isTN3270E bool) {
	if up == nil {
		return
	}
	defer up.Close()

	_ = up.SetWriteDeadline(time.Now().Add(250 * time.Millisecond))

	// 1. If TN3270E, send RFC 2355 UNBIND data message (reason 0x01: normal end of session)
	if isTN3270E {
		unbindMsg := []byte{TN3270E_DATA_UNBIND, 0x00, 0x00, 0x00, 0x00, TN3270E_UNBIND_NORMAL, IAC, EOR}
		_, _ = up.Write(unbindMsg)
	}

	// 2. Send RFC 854 Telnet option unwinding
	unwind := make([]byte, 0, 18)
	unwind = append(unwind,
		IAC, WONT, OPT_BINARY,
		IAC, DONT, OPT_BINARY,
		IAC, WONT, OPT_EOR,
		IAC, DONT, OPT_EOR,
	)
	if isTN3270E {
		unwind = append(unwind, IAC, WONT, OPT_TN3270E, IAC, DONT, OPT_TN3270E)
	}
	_, _ = up.Write(unwind)
}

func killSlot(s *UserMasterSession, slotIdx int) {
	s.Mu.Lock()
	if slotIdx < 0 || slotIdx >= len(s.Slots) || s.Slots[slotIdx] == nil {
		s.Mu.Unlock()
		return
	}
	slot := s.Slots[slotIdx]
	s.Slots[slotIdx] = nil
	if s.ActiveSlotIdx == slotIdx {
		s.ActiveSlotIdx = -1
	}
	sysName := slot.System.Name
	username, displayTag := s.Username, s.DisplayTag
	s.Mu.Unlock()

	slot.Mu.Lock()
	up := slot.Upstream
	isTN3270E := slot.IsTN3270E
	slot.Upstream = nil
	slot.Mu.Unlock()

	cleanCloseUpstream(up, isTN3270E)
	cleanupSessionIfEmpty(s)

	log.Printf("[%s] '%s' Killed system connection [%d] %s (TN3270E: %v)", displayTag, username, slotIdx+1, sysName, isTN3270E)
}

func bridgeClientToUpstream(s *UserMasterSession, slot *SessionSlot, clientConn net.Conn) {
	if clientConn == nil {
		return
	}
	buf := make([]byte, 4096)

	for {
		n, err := clientConn.Read(buf)
		if err != nil {
			s.Mu.Lock()
			username, displayTag := s.Username, s.DisplayTag
			if s.ClientConn == clientConn {
				log.Printf("[%s] '%s' disconnected during active session [%s].", displayTag, username, slot.System.Name)
				s.IsDisconnected = true
				s.ActiveSlotIdx = -1
				s.Mu.Unlock()
				cleanupSessionIfEmpty(s)
			} else {
				log.Printf("[%s] '%s' Transferring previous session [%s] to new terminal.", displayTag, username, slot.System.Name)
				s.Mu.Unlock()
			}
			return
		}

		packet := buf[:n]
		aid := extractAID(packet)
		if aid == go3270.AIDPA3 {
			username, displayTag := s.Identity()
			log.Printf("[%s] '%s' sent PA3. Backgrounding [%s] and returning to menu.", displayTag, username, slot.System.Name)
			return
		}

		slot.Mu.Lock()
		upConn := slot.Upstream
		isTN3270E := slot.IsTN3270E
		isBound := slot.IsBound
		seq := slot.SeqNumber
		if isTN3270E && isBound {
			slot.SeqNumber = (seq + 1) & TN3270E_MAX_SEQ
		}
		slot.Mu.Unlock()

		if upConn == nil {
			username, displayTag := s.Identity()
			log.Printf("[%s] '%s' Upstream host [%s] has disconnected.", displayTag, username, slot.System.Name)
			send3270Message(clientConn, fmt.Sprintf("Upstream host %s sent disconnect.", slot.System.Name), s.Codepage())
			time.Sleep(adminNoticeSleep)
			return
		}

		toSend := packet
		if isTN3270E {
			if isBound {
				seqHi := byte(seq >> 8)
				seqLo := byte(seq & 0xFF)
				if seqLo == IAC {
					toSend = make([]byte, 0, 6+len(packet))
					toSend = append(toSend, TN3270E_DATA_3270, 0x00, 0x00, seqHi, IAC, IAC)
				} else {
					toSend = make([]byte, 0, 5+len(packet))
					toSend = append(toSend, TN3270E_DATA_3270, 0x00, 0x00, seqHi, seqLo)
				}
			} else {
				toSend = make([]byte, 0, 5+len(packet))
				toSend = append(toSend, TN3270E_DATA_SSCP_LU, 0x00, 0x00, 0x00, 0x00)
			}
			toSend = append(toSend, packet...)
		}

		_, err = upConn.Write(toSend)
		if err != nil {
			username, displayTag := s.Identity()
			log.Printf("[%s] '%s' Upstream write error on [%s]: %v", displayTag, username, slot.System.Name, err)
			return
		}
	}
}

func upstreamReader(s *UserMasterSession, slot *SessionSlot, slotIdx int) {
	buf := make([]byte, 8192)
	var recBuf []byte

	for {
		slot.Mu.Lock()
		upConn := slot.Upstream
		slot.Mu.Unlock()
		if upConn == nil {
			return
		}

		n, err := upConn.Read(buf)
		if err != nil {
			username, displayTag := s.Identity()
			log.Printf("[%s] '%s' Upstream host %s severed connection. Killing [%d].", displayTag, username, slot.System.Name, slotIdx+1)
			slot.Mu.Lock()
			if slot.Upstream != nil {
				_ = slot.Upstream.Close()
				slot.Upstream = nil
			}
			slot.Mu.Unlock()

			s.Mu.Lock()
			if slotIdx < len(s.Slots) && s.Slots[slotIdx] == slot {
				s.Slots[slotIdx] = nil
				if s.ActiveSlotIdx == slotIdx {
					s.ActiveSlotIdx = -1
				}
			}
			s.Mu.Unlock()
			cleanupSessionIfEmpty(s)
			return
		}

		data := buf[:n]

		// Handle upstream IAC DO TIMING-MARK (FF FD 06) keep-alive locally
		if bytes.Contains(data, telnetDoTimingMark) {
			_, _ = upConn.Write(telnetWillTimingMark)
			data = bytes.ReplaceAll(data, telnetDoTimingMark, nil)
			if len(data) == 0 {
				continue
			}
		}

		s.Mu.Lock()
		isActive := !s.IsDisconnected && (s.ActiveSlotIdx >= 0 && s.ActiveSlotIdx < len(s.Slots) && s.Slots[s.ActiveSlotIdx] == slot)
		client := s.ClientConn
		username, displayTag := s.Username, s.DisplayTag
		_, _, defaultMaxBuf := GetTerminalMatrix(s.TerminalType, s.DevInfo)
		slot.Mu.Lock()
		isTN3270E := slot.IsTN3270E
		slot.Mu.Unlock()
		s.Mu.Unlock()

		if !isTN3270E && isActive && client != nil {
			_, _ = client.Write(data)
		}

		// Cap buffer at 64KB to avoid unbounded memory use from constant ticks
		if len(recBuf) > 65536 {
			recBuf = recBuf[len(recBuf)-16384:]
		}
		recBuf = append(recBuf, data...)
		for {
			idx := indexTelnetEOR(recBuf)
			if idx == -1 {
				break
			}
			record := recBuf[:idx+2]
			recBuf = recBuf[idx+2:]

			screenData := record
			if isTN3270E && len(record) >= 7 { // 5-byte header + data + IAC EOR (2 bytes)
				dataType := record[0]
				respFlag := record[2]
				seqHi := record[3]
				seqLo := record[4]

				if dataType == TN3270E_DATA_UNBIND {
					reason := byte(0)
					if len(record) > 7 {
						reason = record[5]
					}
					log.Printf("[%s] '%s' Upstream host %s sent UNBIND (reason: 0x%02x). Session returned to unbound/SSCP-LU state.", displayTag, username, slot.System.Name, reason)
					slot.Mu.Lock()
					slot.IsBound = false
					slot.Mu.Unlock()
					continue
				}

				if respFlag == TN3270E_RESP_ALWAYS {
					ack := []byte{TN3270E_DATA_RESP, 0x00, 0x00, seqHi, seqLo, 0x00, IAC, EOR}
					_, _ = upConn.Write(ack)
				}

				switch dataType {
				case TN3270E_DATA_BIND:
					slot.Mu.Lock()
					slot.IsBound = true
					slot.Mu.Unlock()
					continue
				case TN3270E_DATA_3270:
					slot.Mu.Lock()
					slot.IsBound = true
					slot.Mu.Unlock()
					screenData = record[5:]
				case TN3270E_DATA_SSCP_LU:
					slot.Mu.Lock()
					slot.IsBound = false
					slot.Mu.Unlock()
					screenData = record[5:]
				default:
					continue
				}
			}

			if isTN3270E && isActive && client != nil {
				_, _ = client.Write(screenData)
			}

			// If background session's host sends a 3270 Query command, respond
			// with the attached terminal's characteristics locally
			if !isActive && is3270Query(screenData) {
				slot.Mu.Lock()
				isBound := slot.IsBound
				slot.Mu.Unlock()
				sendQueryReply(s, upConn, isTN3270E, isBound)
			}

			if len(screenData) > 2 {
				firstByte := screenData[0]
				isEraseWrite := firstByte == CMD_ERASE_WRITE_SNA || firstByte == CMD_ERASE_WRITE ||
					firstByte == CMD_ERASE_WRITE_ALT_SNA || firstByte == CMD_ERASE_WRITE_ALT ||
					firstByte == CMD_ERASE_ALL_UNPROTECTED

				slot.Mu.Lock()
				if slot.MaxBufferSize <= 0 {
					slot.MaxBufferSize = defaultMaxBuf
				}
				maxBuf := slot.MaxBufferSize
				if slot.IsHeld {
					// Allow screen erase/resets so base state is accurate, but
					// drop/suppress incremental write ticks and do not set ScreenChanged
					if isEraseWrite {
						slot.LastScreen = bytes.Clone(screenData)
					}
					slot.Mu.Unlock()
					continue
				}

				if isEraseWrite || len(slot.LastScreen) == 0 {
					slot.LastScreen = bytes.Clone(screenData)
				} else if len(slot.LastScreen)+len(screenData) <= maxBuf {
					slot.LastScreen = append(slot.LastScreen, screenData...)
				} else {
					// Exceeds maxBuf: retain initial screen baseline and append newest write update,
					// trimming older intermediate writes so total buffer remains strictly <= maxBuf
					baseLen := min(maxBuf/2, len(slot.LastScreen))
					eorIdx := lastIndexTelnetEOR(slot.LastScreen[:baseLen])
					if eorIdx != -1 {
						baseLen = eorIdx + 2
					}
					newBuf := make([]byte, 0, baseLen+len(screenData))
					newBuf = append(newBuf, slot.LastScreen[:baseLen]...)
					newBuf = append(newBuf, screenData...)
					slot.LastScreen = newBuf
				}
				if !isActive {
					slot.ScreenChanged = true
				}
				slot.Mu.Unlock()
			}
		}
	}
}

func findSession(targetID string) *UserMasterSession {
	cleanID := strings.TrimSpace(targetID)
	cleanTag := strings.TrimPrefix(cleanID, "#")

	registryMu.Lock()
	defer registryMu.Unlock()

	// 1. Exact match on DisplayTag (hex, case-insensitive)
	for _, sess := range userSessions {
		_, displayTag := sess.Identity()
		if strings.EqualFold(displayTag, cleanTag) {
			return sess
		}
	}

	// 2. Direct map lookup on SessionKey (hex, case-insensitive)
	if sess, ok := userSessions[strings.ToLower(cleanID)]; ok {
		return sess
	}

	// 3. Exact match on Username (case-sensitive)
	for _, sess := range userSessions {
		username, _ := sess.Identity()
		if username == cleanID {
			return sess
		}
	}

	// 4. Prefix match on SessionKey (at least 4 chars)
	if len(cleanTag) >= 4 {
		cleanTagLower := strings.ToLower(cleanTag)
		for _, sess := range userSessions {
			if strings.HasPrefix(strings.ToLower(sess.SessionKey), cleanTagLower) {
				return sess
			}
		}
	}

	return nil
}

func terminateSession(sess *UserMasterSession, msg string) {
	sess.Mu.Lock()
	type upInfo struct {
		conn      net.Conn
		isTN3270E bool
	}
	var upstreams []upInfo
	for _, slot := range sess.Slots {
		if slot != nil {
			slot.Mu.Lock()
			if slot.Upstream != nil {
				upstreams = append(upstreams, upInfo{conn: slot.Upstream, isTN3270E: slot.IsTN3270E})
				slot.Upstream = nil
			}
			slot.Mu.Unlock()
		}
	}
	sess.Slots = nil
	sess.ActiveSlotIdx = -1
	sess.IsDisconnected = true
	c := sess.ClientConn
	sess.ClientConn = nil
	cp := sess.codepageLocked()
	sess.Mu.Unlock()

	if c != nil {
		if msg != "" {
			send3270Message(c, msg, cp)
			time.Sleep(adminNoticeSleep / 5)
		}
		_ = c.Close()
	}

	for _, up := range upstreams {
		cleanCloseUpstream(up.conn, up.isTN3270E)
	}
}

func handleAdminKill(adminSession *UserMasterSession, targetID string) bool {
	adminSession.Mu.Lock()
	adminUser := adminSession.Username
	adminConn := adminSession.ClientConn
	adminCP := adminSession.codepageLocked()
	adminSession.Mu.Unlock()

	target := findSession(targetID)
	if target == nil {
		log.Printf("[ADMIN] %s attempted to kill session '%s' (not found)", adminUser, targetID)
		send3270Message(adminConn, fmt.Sprintf("Session '%s' not found.", targetID), adminCP)
		time.Sleep(adminNoticeSleep)
		return false
	}

	targetUser, targetTag := target.Identity()
	log.Printf("[ADMIN] %s killing session '%s' (User: %s)", adminUser, targetTag, targetUser)

	registryMu.Lock()
	delete(userSessions, target.SessionKey)
	registryMu.Unlock()

	if target == adminSession {
		terminateSession(target, "Session terminated.")
		return true
	}

	terminateSession(target, "Session terminated by administrator.")
	send3270Message(adminConn, fmt.Sprintf("Session '%s' (%s) terminated.", targetTag, targetUser), adminCP)
	time.Sleep(adminNoticeSleep)
	return false
}

func drainAllSessions() []*UserMasterSession {
	registryMu.Lock()
	defer registryMu.Unlock()

	allSessions := make([]*UserMasterSession, 0, len(userSessions))
	for _, sess := range userSessions {
		allSessions = append(allSessions, sess)
	}
	userSessions = make(map[string]*UserMasterSession)
	return allSessions
}

func handleAdminKillAll(adminSession *UserMasterSession) {
	adminUser, _ := adminSession.Identity()
	log.Printf("[ADMIN] %s executed 'killall'", adminUser)

	allSessions := drainAllSessions()
	for _, sess := range allSessions {
		terminateSession(sess, "Session terminated by administrator.")
	}
}

func handleAdminBye(adminSession *UserMasterSession) {
	adminUser, _ := adminSession.Identity()
	log.Printf("[ADMIN] %s executed 'bye': shutting down server", adminUser)

	allSessions := drainAllSessions()
	for _, sess := range allSessions {
		terminateSession(sess, "Server shutting down.")
	}

	if serverListener != nil {
		_ = serverListener.Close()
	}

	if shutdownFunc != nil {
		shutdownFunc()
	}
}