package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type config struct {
	listen            string
	profiles          string
	mgmtListen        string
	mgmtAdvertiseHost string
	rewriteIDName     bool
	idNameSeparator   string
	engineName        string
	lc0Path           string
	weights           string
	backend           string
	backendOpts       string
	threads           int
	nncache           int
	minibatchSize     int
	maxPrefetch       int
	authTokenEnv      string
	shutdownWait      time.Duration
	idleExitWait      time.Duration
	engineIdleWait    time.Duration
}

type engineSpec struct {
	name          string
	port          string
	listenAddr    string
	lc0Path       string
	execPathLabel string
	staticArgs    []string
	useStaticArgs bool
	runtime       *engineRuntime
}

type runningEngine struct {
	cmd    *exec.Cmd
	stdin  io.WriteCloser
	lines  chan string
	done   chan struct{}
	waitMu sync.Mutex
	err    error
	initialized atomic.Bool
}

type engineRuntime struct {
	spec        *engineSpec
	cfg         config
	sessionMu   sync.Mutex
	mu          sync.Mutex
	run         *runningEngine
	active      atomic.Bool
	lastInputNS atomic.Int64
}

type sessionInfo struct {
	id          int64
	remote      string
	local       string
	engineName  string
	enginePort  string
	engineInst  string
	connectedAt time.Time
}

type serverState struct {
	activeSessions int64
	lastActivityNS int64
	nextSessionID  int64
	mu             sync.Mutex
	sessions       map[int64]sessionInfo
}

type profileEntry struct {
	name       string
	port       string
	executable string
	args       string
}

var profileConfLine = regexp.MustCompile(`^(\S+)\s+(\S+)\s+(\S+)(?:\s+(.*))?$`)

func main() {
	cfg := parseFlags()
	engines, err := buildEngineSpecs(cfg)
	if err != nil {
		log.Fatalf("engine config failed: %v", err)
	}

	state := &serverState{
		lastActivityNS: time.Now().UnixNano(),
		sessions:       map[int64]sessionInfo{},
	}

	for i := range engines {
		engines[i].runtime = &engineRuntime{
			spec: &engines[i],
			cfg:  cfg,
		}
		engines[i].runtime.lastInputNS.Store(time.Now().UnixNano())
		if cfg.engineIdleWait > 0 {
			go engines[i].runtime.monitorIdle()
		}
	}

	if cfg.idleExitWait > 0 {
		log.Printf("idle auto-exit enabled timeout=%s", cfg.idleExitWait)
		go monitorIdleExit(cfg.idleExitWait, state)
	}

	if strings.TrimSpace(cfg.mgmtListen) != "" {
		go runManagementListener(cfg, state, engines)
	}

	for i := range engines {
		engine := &engines[i]
		ln, err := net.Listen("tcp", engine.listenAddr)
		if err != nil {
			log.Fatalf("failed to listen for engine %s on %s: %v", engine.name, engine.listenAddr, err)
		}
		log.Printf("uci listener active engine=%s addr=%s exec=%s", engine.name, engine.listenAddr, engine.execPathLabel)
		go runUCIAcceptLoop(ln, cfg, state, engine)
	}

	select {}
}

func parseFlags() config {
	var cfg config
	defaultIdleExit := strings.TrimSpace(os.Getenv("UCI_IDLE_EXIT_TIMEOUT"))
	if defaultIdleExit == "" {
		defaultIdleExit = "0"
	}
	defaultEngineIdle := strings.TrimSpace(os.Getenv("LC0_ENGINE_IDLE_TIMEOUT"))
	if defaultEngineIdle == "" {
		defaultEngineIdle = "1h"
	}
	defaultMgmtListen := strings.TrimSpace(os.Getenv("UCI_MGMT_LISTEN"))
	if defaultMgmtListen == "" {
		defaultMgmtListen = "127.0.0.1:53350"
	}
	defaultMgmtAdvertiseHost := strings.TrimSpace(os.Getenv("MGMT_ADVERTISE_HOST"))
	if defaultMgmtAdvertiseHost == "" {
		defaultMgmtAdvertiseHost = strings.TrimSpace(os.Getenv("TS_HOSTNAME"))
	}
	defaultEngineName := strings.TrimSpace(os.Getenv("UCI_ENGINE_NAME"))
	if defaultEngineName == "" {
		defaultEngineName = "lc0-cloud"
	}
	defaultRewriteIDName := parseBoolEnv(os.Getenv("UCI_REWRITE_ID_NAME"), false)
	defaultIDNameSeparator := strings.TrimSpace(os.Getenv("UCI_ID_NAME_SEPARATOR"))
	if defaultIDNameSeparator == "" {
		defaultIDNameSeparator = "##"
	}

	var idleExitRaw string
	var engineIdleRaw string
	flag.StringVar(&cfg.listen, "listen", "127.0.0.1:5555", "Default TCP listen address (host:port) for single-engine mode")
	flag.StringVar(&cfg.profiles, "profiles", strings.TrimSpace(os.Getenv("UCI_ENGINE_PROFILES")), "Path to engine_profiles.conf-style engine profile file")
	flag.StringVar(&cfg.mgmtListen, "mgmt-listen", defaultMgmtListen, "Management listen address for legacy-compatible scan/query; empty disables")
	flag.StringVar(&cfg.mgmtAdvertiseHost, "mgmt-advertise-host", defaultMgmtAdvertiseHost, "Host/IP announced in management scan responses")
	flag.BoolVar(&cfg.rewriteIDName, "rewrite-id-name", defaultRewriteIDName, "Rewrite engine 'id name' responses to legacy-style identity")
	flag.StringVar(&cfg.idNameSeparator, "id-name-separator", defaultIDNameSeparator, "Separator used when rewriting 'id name' (default ##)")
	flag.StringVar(&cfg.engineName, "engine-name", defaultEngineName, "Default engine name in single-engine mode")
	flag.StringVar(&cfg.lc0Path, "lc0", "/usr/local/bin/lc0", "Default path to lc0 binary")
	flag.StringVar(&cfg.weights, "weights", "/data/weights.pb.gz", "Default path to weights file, empty or <autodiscover> enables lookup")
	flag.StringVar(&cfg.backend, "backend", "cuda-fp16", "Default lc0 backend (e.g. cuda-fp16, cuda, cuda-auto, demux)")
	flag.StringVar(&cfg.backendOpts, "backend-opts", "", "Default lc0 backend options")
	flag.IntVar(&cfg.threads, "threads", 2, "Default lc0 threads")
	flag.IntVar(&cfg.nncache, "nncache", 200000, "Default lc0 nncache size")
	flag.IntVar(&cfg.minibatchSize, "minibatch-size", 0, "Default lc0 minibatch size (0 lets lc0 choose)")
	flag.IntVar(&cfg.maxPrefetch, "max-prefetch", 32, "Default lc0 max prefetch")
	flag.StringVar(&cfg.authTokenEnv, "auth-token-env", "UCI_AUTH_TOKEN", "Env var containing optional auth token")
	flag.StringVar(&idleExitRaw, "idle-exit-timeout", defaultIdleExit, "Exit bridge when no active sessions for this duration (e.g. 15m, 1h, or seconds). 0 disables")
	flag.StringVar(&engineIdleRaw, "engine-idle-timeout", defaultEngineIdle, "Stop persistent lc0 process after no engine input for this duration while idle; 0 disables")
	flag.Parse()
	if strings.TrimSpace(cfg.idNameSeparator) == "" {
		cfg.idNameSeparator = "##"
	}

	cfg.shutdownWait = parseShutdownTimeout(os.Getenv("LC0_SHUTDOWN_TIMEOUT"))
	cfg.idleExitWait = parseIdleTimeout(idleExitRaw)
	cfg.engineIdleWait = parseIdleTimeout(engineIdleRaw)
	return cfg
}

func parseBoolEnv(raw string, fallback bool) bool {
	v := strings.ToLower(strings.TrimSpace(raw))
	switch v {
	case "1", "true", "yes", "y", "on":
		return true
	case "0", "false", "no", "n", "off":
		return false
	case "":
		return fallback
	default:
		return fallback
	}
}

func buildEngineSpecs(cfg config) ([]engineSpec, error) {
	if strings.TrimSpace(cfg.profiles) == "" {
		return []engineSpec{
			{
				name:          cfg.engineName,
				port:          listenPort(cfg.listen),
				listenAddr:    cfg.listen,
				lc0Path:       cfg.lc0Path,
				execPathLabel: cfg.lc0Path,
				useStaticArgs: false,
			},
		}, nil
	}

	entries, err := loadProfiles(cfg.profiles)
	if err != nil {
		return nil, err
	}
	listenHost := listenHost(cfg.listen)
	baseDir := filepath.Dir(cfg.profiles)
	out := make([]engineSpec, 0, len(entries))

	for _, p := range entries {
		addr := net.JoinHostPort(listenHost, p.port)
		path := resolveProfileExecutable(baseDir, p)
		staticArgs, useStatic := parseProfileArgs(p.args)
		out = append(out, engineSpec{
			name:          p.name,
			port:          p.port,
			listenAddr:    addr,
			lc0Path:       path,
			execPathLabel: path,
			staticArgs:    staticArgs,
			useStaticArgs: useStatic,
		})
	}

	return out, nil
}

func loadProfiles(path string) ([]profileEntry, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open profiles %s: %w", path, err)
	}
	defer f.Close()

	var out []profileEntry
	s := bufio.NewScanner(f)
	lineNo := 0
	seen := map[string]bool{}

	for s.Scan() {
		lineNo++
		line := strings.TrimSpace(s.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		m := profileConfLine.FindStringSubmatch(line)
		if m == nil {
			return nil, fmt.Errorf("invalid profile line %d: %q", lineNo, line)
		}

		portNum, err := strconv.Atoi(m[2])
		if err != nil || portNum < 1 || portNum > 65535 {
			return nil, fmt.Errorf("invalid port on line %d: %q", lineNo, m[2])
		}
		if seen[m[2]] {
			return nil, fmt.Errorf("duplicate port on line %d: %q", lineNo, m[2])
		}
		seen[m[2]] = true

		args := strings.TrimSpace(m[4])
		out = append(out, profileEntry{
			name:       m[1],
			port:       m[2],
			executable: m[3],
			args:       args,
		})
	}
	if err := s.Err(); err != nil {
		return nil, fmt.Errorf("read profiles %s: %w", path, err)
	}
	if len(out) == 0 {
		return nil, fmt.Errorf("no usable profiles in %s", path)
	}

	return out, nil
}

func resolveProfileExecutable(baseDir string, p profileEntry) string {
	exe := strings.TrimSpace(p.executable)
	if filepath.IsAbs(exe) {
		return exe
	}
	if strings.ContainsAny(exe, `/\`) {
		return filepath.Join(baseDir, exe)
	}
	return filepath.Join(baseDir, p.name, exe)
}

func parseProfileArgs(raw string) ([]string, bool) {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return nil, false
	}

	if strings.Contains(raw, ":") {
		parts := strings.Split(raw, ":")
		out := make([]string, 0, len(parts))
		allFlags := true
		for _, p := range parts {
			p = strings.TrimSpace(p)
			if p != "" {
				out = append(out, p)
				if !strings.HasPrefix(p, "-") {
					allFlags = false
				}
			}
		}
		if len(out) > 1 && allFlags {
			return out, true
		}
	}

	fields := strings.Fields(raw)
	if len(fields) == 0 {
		return nil, false
	}
	return fields, true
}

func listenHost(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return "127.0.0.1"
	}
	if host == "" {
		return "0.0.0.0"
	}
	return host
}

func parseShutdownTimeout(raw string) time.Duration {
	if strings.TrimSpace(raw) == "" {
		return 3 * time.Second
	}
	if d, err := time.ParseDuration(raw); err == nil {
		return d
	}
	if s, err := strconv.Atoi(raw); err == nil && s > 0 {
		return time.Duration(s) * time.Second
	}
	log.Printf("invalid LC0_SHUTDOWN_TIMEOUT=%q, defaulting to 3s", raw)
	return 3 * time.Second
}

func parseIdleTimeout(raw string) time.Duration {
	v := strings.TrimSpace(raw)
	if v == "" || v == "0" {
		return 0
	}
	if d, err := time.ParseDuration(v); err == nil {
		if d < 0 {
			log.Fatalf("invalid idle timeout %q: must be >= 0", raw)
		}
		return d
	}
	if s, err := strconv.Atoi(v); err == nil {
		if s < 0 {
			log.Fatalf("invalid idle timeout %q: must be >= 0", raw)
		}
		return time.Duration(s) * time.Second
	}
	log.Fatalf("invalid idle timeout %q: use duration format (e.g. 15m) or integer seconds", raw)
	return 0
}

func monitorIdleExit(timeout time.Duration, state *serverState) {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		if atomic.LoadInt64(&state.activeSessions) > 0 {
			continue
		}
		last := atomic.LoadInt64(&state.lastActivityNS)
		if last == 0 {
			continue
		}
		idleFor := time.Since(time.Unix(0, last))
		if idleFor >= timeout {
			log.Printf("no active sessions for %s (threshold %s), exiting bridge", idleFor.Round(time.Second), timeout)
			os.Exit(0)
		}
	}
}

func runUCIAcceptLoop(ln net.Listener, cfg config, state *serverState, engine *engineSpec) {
	defer ln.Close()

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("accept error engine=%s addr=%s err=%v", engine.name, engine.listenAddr, err)
			continue
		}
		atomic.AddInt64(&state.activeSessions, 1)
		atomic.StoreInt64(&state.lastActivityNS, time.Now().UnixNano())
		go handleConn(conn, cfg, state, engine)
	}
}

func handleConn(conn net.Conn, cfg config, state *serverState, engine *engineSpec) {
	defer conn.Close()
	remote := conn.RemoteAddr().String()
	local := conn.LocalAddr().String()
	sessionID := atomic.AddInt64(&state.nextSessionID, 1)

	registerSession(state, sessionInfo{
		id:          sessionID,
		remote:      remote,
		local:       local,
		engineName:  engine.name,
		enginePort:  engine.port,
		engineInst:  fmt.Sprintf("session-%d", sessionID),
		connectedAt: time.Now(),
	})

	log.Printf("connection start engine=%s remote=%s", engine.name, remote)
	defer func() {
		atomic.AddInt64(&state.activeSessions, -1)
		atomic.StoreInt64(&state.lastActivityNS, time.Now().UnixNano())
		unregisterSession(state, sessionID)
		log.Printf("connection end engine=%s remote=%s", engine.name, remote)
	}()

	token := strings.TrimSpace(os.Getenv(cfg.authTokenEnv))
	reader := bufio.NewReader(conn)
	if token != "" {
		if err := doAuth(conn, reader, token); err != nil {
			log.Printf("auth failed engine=%s remote=%s err=%v", engine.name, remote, err)
			writeSessionError(conn, "unauthorized")
			return
		}
		log.Printf("auth success engine=%s remote=%s", engine.name, remote)
	}

	engine.runtime.sessionMu.Lock()
	engine.runtime.active.Store(true)
	defer func() {
		engine.runtime.active.Store(false)
		engine.runtime.sessionMu.Unlock()
	}()

	args, err := resolveEngineArgs(engine, cfg)
	if err != nil {
		log.Printf("engine args resolution failed engine=%s remote=%s err=%v", engine.name, remote, err)
		writeSessionError(conn, fmt.Sprintf("engine args resolution failed: %v", err))
		return
	}

	run, err := engine.runtime.ensureRunning(args, remote)
	if err != nil {
		log.Printf("failed to start engine=%s remote=%s cmd=%s args=%v err=%v", engine.name, remote, engine.lc0Path, args, err)
		writeSessionError(conn, fmt.Sprintf("failed to start engine: %v", err))
		return
	}

	if run.initialized.Load() {
		if err := engine.runtime.resetToReady(run, remote); err != nil {
			log.Printf("engine reset before session failed engine=%s remote=%s err=%v", engine.name, remote, err)
			writeSessionError(conn, fmt.Sprintf("engine reset before session failed: %v", err))
			_ = engine.runtime.stopProcess("failed pre-session reset")
			return
		}
	} else {
		log.Printf("skipping pre-session reset for fresh engine=%s remote=%s until first uci handshake", engine.name, remote)
	}

	stopForward := make(chan struct{})
	inDone := make(chan error, 1)
	go func() {
		inDone <- engine.runtime.forwardClientInput(reader, run)
	}()

	outDone := make(chan error, 1)
	go func() {
		outDone <- forwardEngineOutput(conn, run.lines, cfg, engine, conn, stopForward)
	}()

	var sessionErr error
	outputStopped := false
	select {
	case sessionErr = <-inDone:
		close(stopForward)
		_ = conn.Close()
		if sessionErr != nil && !errors.Is(sessionErr, net.ErrClosed) {
			log.Printf("input copy ended engine=%s remote=%s err=%v", engine.name, remote, sessionErr)
		}
	case sessionErr = <-outDone:
		outputStopped = true
		close(stopForward)
		_ = conn.Close()
		if sessionErr != nil && !errors.Is(sessionErr, net.ErrClosed) {
			log.Printf("output copy ended engine=%s remote=%s err=%v", engine.name, remote, sessionErr)
		}
	case <-run.done:
		close(stopForward)
		_ = conn.Close()
		sessionErr = run.exitErr()
	}

	if !outputStopped {
		<-outDone
	}

	if run.initialized.Load() {
		if err := engine.runtime.resetToReady(run, remote); err != nil {
			log.Printf("engine reset after session failed engine=%s remote=%s err=%v", engine.name, remote, err)
			_ = engine.runtime.stopProcess("failed post-session reset")
		}
	} else {
		log.Printf("skipping post-session reset for engine=%s remote=%s because uci handshake never completed", engine.name, remote)
	}

	if sessionErr != nil && !errors.Is(sessionErr, io.EOF) {
		log.Printf("session ended engine=%s remote=%s err=%v", engine.name, remote, sessionErr)
	}
}

func resolveEngineArgs(engine *engineSpec, cfg config) ([]string, error) {
	if engine.useStaticArgs {
		args := make([]string, len(engine.staticArgs))
		copy(args, engine.staticArgs)
		return args, nil
	}
	weights, err := resolveWeights(cfg.weights)
	if err != nil {
		return nil, err
	}
	return buildDefaultLc0Args(weights, cfg), nil
}

func forwardEngineOutput(dst io.Writer, src <-chan string, cfg config, engine *engineSpec, conn net.Conn, stop <-chan struct{}) error {
	writer := bufio.NewWriter(dst)

	for {
		select {
		case <-stop:
			return nil
		case line, ok := <-src:
			if !ok {
				return nil
			}
			line = maybeRewriteIDNameLine(line, cfg, engine, conn)
			if _, werr := writer.WriteString(line); werr != nil {
				return werr
			}
			if ferr := writer.Flush(); ferr != nil {
				return ferr
			}
		}
	}
}

func maybeRewriteIDNameLine(line string, cfg config, engine *engineSpec, conn net.Conn) string {
	if !cfg.rewriteIDName {
		return line
	}

	trimmed := strings.TrimRight(line, "\r\n")
	if !strings.HasPrefix(trimmed, "id name ") {
		return line
	}

	suffix := line[len(trimmed):]
	origName := strings.TrimPrefix(trimmed, "id name ")
	host := scanAdvertiseHost(cfg, conn)
	rewritten := fmt.Sprintf("id name %s%s_%s_%s%s%s", legacyHeaderPrefix(), host, engine.port, engine.name, cfg.idNameSeparator, origName)
	return rewritten + suffix
}

func (r *runningEngine) setExitErr(err error) {
	r.waitMu.Lock()
	defer r.waitMu.Unlock()
	r.err = err
}

func (r *runningEngine) exitErr() error {
	r.waitMu.Lock()
	defer r.waitMu.Unlock()
	return r.err
}

func (rt *engineRuntime) monitorIdle() {
	ticker := time.NewTicker(idleCheckInterval(rt.cfg.engineIdleWait))
	defer ticker.Stop()

	for range ticker.C {
		if rt.cfg.engineIdleWait <= 0 || rt.active.Load() {
			continue
		}
		last := rt.lastInputNS.Load()
		if last == 0 {
			continue
		}
		idleFor := time.Since(time.Unix(0, last))
		if idleFor < rt.cfg.engineIdleWait {
			continue
		}

		rt.sessionMu.Lock()
		if rt.active.Load() {
			rt.sessionMu.Unlock()
			continue
		}
		last = rt.lastInputNS.Load()
		idleFor = time.Since(time.Unix(0, last))
		if idleFor >= rt.cfg.engineIdleWait {
			if err := rt.stopProcess(fmt.Sprintf("engine idle timeout after %s", idleFor.Round(time.Second))); err != nil {
				log.Printf("engine idle shutdown failed engine=%s err=%v", rt.spec.name, err)
			}
		}
		rt.sessionMu.Unlock()
	}
}

func idleCheckInterval(timeout time.Duration) time.Duration {
	if timeout <= 0 {
		return 5 * time.Second
	}
	interval := timeout / 4
	if interval < time.Second {
		return time.Second
	}
	if interval > 5*time.Second {
		return 5 * time.Second
	}
	return interval
}

func (rt *engineRuntime) ensureRunning(args []string, remote string) (*runningEngine, error) {
	rt.mu.Lock()
	if rt.run != nil {
		run := rt.run
		rt.mu.Unlock()
		return run, nil
	}
	rt.mu.Unlock()

	cmd := exec.Command(rt.spec.lc0Path, args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("stdin pipe failed: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdout pipe failed: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("stderr pipe failed: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	run := &runningEngine{
		cmd:   cmd,
		stdin: stdin,
		lines: make(chan string, 1024),
		done:  make(chan struct{}),
	}

	rt.mu.Lock()
	rt.run = run
	rt.mu.Unlock()
	rt.lastInputNS.Store(time.Now().UnixNano())

	log.Printf("engine started engine=%s remote=%s pid=%d cmd=%s args=%v", rt.spec.name, remote, cmd.Process.Pid, rt.spec.lc0Path, args)

	go func() {
		s := bufio.NewScanner(stdout)
		buf := make([]byte, 0, 64*1024)
		s.Buffer(buf, 1024*1024)
		for s.Scan() {
			line := s.Text()
			if strings.TrimSpace(line) == "uciok" {
				run.initialized.Store(true)
			}
			run.lines <- line + "\n"
		}
		if err := s.Err(); err != nil {
			log.Printf("stdout scan error engine=%s err=%v", rt.spec.name, err)
		}
		close(run.lines)
	}()

	go func() {
		s := bufio.NewScanner(stderr)
		buf := make([]byte, 0, 64*1024)
		s.Buffer(buf, 1024*1024)
		for s.Scan() {
			log.Printf("engine stderr engine=%s: %s", rt.spec.name, s.Text())
		}
		if err := s.Err(); err != nil {
			log.Printf("stderr scan error engine=%s err=%v", rt.spec.name, err)
		}
	}()

	go func() {
		err := cmd.Wait()
		run.setExitErr(err)
		close(run.done)

		rt.mu.Lock()
		if rt.run == run {
			rt.run = nil
		}
		rt.mu.Unlock()

		if err != nil {
			log.Printf("engine exited engine=%s pid=%d err=%v", rt.spec.name, cmd.Process.Pid, err)
		} else {
			log.Printf("engine exited cleanly engine=%s pid=%d", rt.spec.name, cmd.Process.Pid)
		}
	}()

	return run, nil
}

func (rt *engineRuntime) forwardClientInput(reader *bufio.Reader, run *runningEngine) error {
	for {
		line, err := reader.ReadString('\n')
		if len(line) > 0 {
			trimmed := strings.TrimSpace(line)
			if trimmed != "" {
				rt.lastInputNS.Store(time.Now().UnixNano())
			}
			if trimmed == "quit" {
				return nil
			}
			if _, werr := io.WriteString(run.stdin, line); werr != nil {
				return werr
			}
		}
		if err != nil {
			if errors.Is(err, io.EOF) {
				return nil
			}
			return err
		}
	}
}

func (rt *engineRuntime) resetToReady(run *runningEngine, remote string) error {
	select {
	case <-run.done:
		return run.exitErr()
	default:
	}

	rt.drainLines(run)
	if _, err := io.WriteString(run.stdin, "stop\nisready\n"); err != nil {
		return err
	}

	timer := time.NewTimer(30 * time.Second)
	defer timer.Stop()

	for {
		select {
		case <-run.done:
			return run.exitErr()
		case <-timer.C:
			return fmt.Errorf("timed out waiting for readyok engine=%s remote=%s", rt.spec.name, remote)
		case line, ok := <-run.lines:
			if !ok {
				return run.exitErr()
			}
			if strings.TrimSpace(line) == "readyok" {
				return nil
			}
		}
	}
}

func (rt *engineRuntime) drainLines(run *runningEngine) {
	for {
		select {
		case _, ok := <-run.lines:
			if !ok {
				return
			}
		default:
			return
		}
	}
}

func (rt *engineRuntime) stopProcess(reason string) error {
	rt.mu.Lock()
	run := rt.run
	rt.mu.Unlock()
	if run == nil {
		return nil
	}

	log.Printf("stopping engine process engine=%s reason=%s", rt.spec.name, reason)
	select {
	case <-run.done:
		return run.exitErr()
	default:
	}

	if _, err := io.WriteString(run.stdin, "quit\n"); err != nil && !errors.Is(err, os.ErrClosed) {
		log.Printf("engine quit write failed engine=%s err=%v", rt.spec.name, err)
	}

	return waitForProcessExit(run, rt.cfg.shutdownWait, rt.spec.name)
}

func registerSession(state *serverState, s sessionInfo) {
	state.mu.Lock()
	defer state.mu.Unlock()
	state.sessions[s.id] = s
}

func unregisterSession(state *serverState, id int64) {
	state.mu.Lock()
	defer state.mu.Unlock()
	delete(state.sessions, id)
}

func snapshotSessions(state *serverState) []sessionInfo {
	state.mu.Lock()
	defer state.mu.Unlock()

	out := make([]sessionInfo, 0, len(state.sessions))
	for _, s := range state.sessions {
		out = append(out, s)
	}
	sort.Slice(out, func(i, j int) bool {
		return out[i].id < out[j].id
	})
	return out
}

func runManagementListener(cfg config, state *serverState, engines []engineSpec) {
	ln, err := net.Listen("tcp", cfg.mgmtListen)
	if err != nil {
		log.Printf("management listener disabled; failed to listen on %s: %v", cfg.mgmtListen, err)
		return
	}
	defer ln.Close()
	log.Printf("management listener active on %s", cfg.mgmtListen)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("management accept error: %v", err)
			continue
		}
		go handleManagementConn(conn, cfg, state, engines)
	}
}

func handleManagementConn(conn net.Conn, cfg config, state *serverState, engines []engineSpec) {
	defer conn.Close()
	remote := conn.RemoteAddr().String()
	buf := make([]byte, 1024)

	_ = conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	n, err := conn.Read(buf)
	_ = conn.SetReadDeadline(time.Time{})
	if err != nil {
		if errors.Is(err, io.EOF) {
			return
		}
		log.Printf("management read failed remote=%s err=%v", remote, err)
		return
	}

	cmd := strings.ToLower(strings.TrimSpace(string(buf[:n])))
	switch {
	case strings.HasPrefix(cmd, "scan"):
		serveScan(conn, cfg, engines)
		log.Printf("management scan served remote=%s", remote)
	case cmd == "query":
		serveQuery(conn, cfg, state, engines)
		log.Printf("management query served remote=%s", remote)
	default:
		_, _ = io.WriteString(conn, "unsupported\n")
		log.Printf("management unknown command remote=%s cmd=%q", remote, cmd)
	}
}

func serveScan(conn net.Conn, cfg config, engines []engineSpec) {
	host := scanAdvertiseHost(cfg, conn)
	for _, e := range engines {
		line := fmt.Sprintf("%s%s_%s_%s\n", legacyHeaderPrefix(), host, e.port, e.name)
		_, _ = io.WriteString(conn, line)
	}
	_, _ = io.WriteString(conn, "scanisdone\n")
}

func serveQuery(conn net.Conn, _ config, state *serverState, engines []engineSpec) {
	hostName, err := os.Hostname()
	if err != nil || strings.TrimSpace(hostName) == "" {
		hostName = "UNKNOWN_SERVER"
	}

	sessions := snapshotSessions(state)
	var b strings.Builder
	fmt.Fprintf(&b, "\n===== Engine Table Entries from Server (%s) OS-ARCH (%s)  =====\n", hostName, legacyOSArch())

	for i, e := range engines {
		if i > 0 {
			fmt.Fprintf(&b, "\n")
		}
		fmt.Fprintf(&b, "Engine(%s) TCP Port(%s)\n", e.name, e.port)
		fmt.Fprintf(&b, "   Executable On Server(%s)\n", e.execPathLabel)
		fmt.Fprintf(&b, "   Connected Users:\n")
		for _, s := range sessions {
			if s.engineName != e.name || s.enginePort != e.port {
				continue
			}
			fmt.Fprintf(&b, "      * Client IP[%s] Socket(%d) Server IP[%s] Engine Instance(%s)\n",
				hostFromAddr(s.remote), s.id, hostFromAddr(s.local), s.engineInst)
		}
	}

	fmt.Fprintf(&b, "================================<<<querydone\n\n")
	_, _ = io.WriteString(conn, b.String())
}

func scanAdvertiseHost(cfg config, conn net.Conn) string {
	if v := strings.TrimSpace(cfg.mgmtAdvertiseHost); v != "" {
		return v
	}
	host := hostFromAddr(conn.LocalAddr().String())
	if host == "" {
		return "127.0.0.1"
	}
	return host
}

func hostFromAddr(addr string) string {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return addr
	}
	return host
}

func listenPort(listenAddr string) string {
	_, port, err := net.SplitHostPort(listenAddr)
	if err != nil {
		return listenAddr
	}
	return port
}

func legacyHeaderPrefix() string {
	switch {
	case runtime.GOOS == "windows":
		return "JRE_X64WIN_"
	case runtime.GOARCH == "arm64":
		return "JRE_XAVIER_"
	default:
		return "JRE_X64LNX_"
	}
}

func legacyOSArch() string {
	switch {
	case runtime.GOOS == "windows":
		return "Windows X86-64"
	case runtime.GOARCH == "arm64":
		return "Xavier ARM64"
	default:
		return "Linux X86-64"
	}
}

func doAuth(conn net.Conn, reader *bufio.Reader, token string) error {
	_ = conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	line, err := reader.ReadString('\n')
	_ = conn.SetReadDeadline(time.Time{})
	if err != nil {
		return fmt.Errorf("failed reading auth line: %w", err)
	}
	line = strings.TrimSpace(line)
	if line != "AUTH "+token {
		return fmt.Errorf("invalid auth line")
	}
	return nil
}

func buildDefaultLc0Args(weights string, cfg config) []string {
	args := []string{
		"-w", weights,
		"--backend=" + cfg.backend,
		fmt.Sprintf("--threads=%d", cfg.threads),
		fmt.Sprintf("--nncache=%d", cfg.nncache),
		fmt.Sprintf("--minibatch-size=%d", cfg.minibatchSize),
		fmt.Sprintf("--max-prefetch=%d", cfg.maxPrefetch),
	}
	if strings.TrimSpace(cfg.backendOpts) != "" {
		args = append(args, "--backend-opts="+cfg.backendOpts)
	}
	return args
}

func resolveWeights(flagValue string) (string, error) {
	v := strings.TrimSpace(flagValue)
	if v != "" && v != "<autodiscover>" {
		if _, err := os.Stat(v); err != nil {
			return "", fmt.Errorf("weights not found at %s: %w", v, err)
		}
		return v, nil
	}

	candidates, err := discoverWeights()
	if err != nil {
		return "", err
	}
	if len(candidates) == 0 {
		return "", errors.New("no weights found (set --weights or mount under /data)")
	}
	log.Printf("autodiscovered weights at %s", candidates[0])
	return candidates[0], nil
}

func discoverWeights() ([]string, error) {
	searchRoots := []string{"/data", "/workspace", "."}
	exts := map[string]bool{".pb": true, ".gz": true}
	seen := map[string]bool{}
	var out []string

	for _, root := range searchRoots {
		info, err := os.Stat(root)
		if err != nil || !info.IsDir() {
			continue
		}
		err = filepath.WalkDir(root, func(path string, d os.DirEntry, walkErr error) error {
			if walkErr != nil {
				return nil
			}
			if d.IsDir() {
				name := d.Name()
				if strings.HasPrefix(name, ".") && path != root {
					return filepath.SkipDir
				}
				return nil
			}
			name := strings.ToLower(d.Name())
			if !(strings.HasSuffix(name, ".pb") || strings.HasSuffix(name, ".pb.gz")) {
				return nil
			}
			ext := filepath.Ext(name)
			if !exts[ext] {
				return nil
			}
			if !seen[path] {
				seen[path] = true
				out = append(out, path)
			}
			return nil
		})
		if err != nil {
			return nil, err
		}
	}

	return out, nil
}

func waitForProcessExit(run *runningEngine, timeout time.Duration, engineName string) error {
	select {
	case <-run.done:
		return run.exitErr()
	case <-time.After(timeout):
		if run.cmd.Process != nil {
			log.Printf("engine shutdown timeout engine=%s timeout=%s, killing pid=%d", engineName, timeout, run.cmd.Process.Pid)
			if err := run.cmd.Process.Kill(); err != nil {
				log.Printf("failed to kill engine=%s err=%v", engineName, err)
			}
		}
		<-run.done
		return run.exitErr()
	}
}

func writeSessionError(conn net.Conn, message string) {
	trimmed := strings.TrimSpace(message)
	if trimmed == "" {
		trimmed = "session failed"
	}
	trimmed = strings.ReplaceAll(trimmed, "\n", " ")
	trimmed = strings.ReplaceAll(trimmed, "\r", " ")
	_, _ = io.WriteString(conn, "error "+trimmed+"\n")
}
