package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Owner Guard: the machine's owner always comes first (docs/DESIGN-TH.md section 5).
//
//	available -> draining -> owner      (owner activity / foreign GPU work)
//	any       -> paused                 (manual: web button or `wxch-agent -pause`)
//
// Entering owner/paused: stop accepting new requests, let running ones finish within a grace period
// (then cancel them), and unload every model so the memory is free for the owner.
// Leaving owner mode needs a quiet period, so a mouse wiggle does not make the machine flap.

type guardMode string

const (
	modeAvailable guardMode = "available"
	modeDraining  guardMode = "draining"
	modeOwner     guardMode = "owner"
	modePaused    guardMode = "paused"
)

type guardConfig struct {
	Enabled     bool
	ActiveIdleS int64 // owner counts as present when the desktop was used less than this many seconds ago
	ReturnIdleS int64 // owner counts as gone after this many quiet seconds
	GraceS      int64 // running requests may finish for this long before being cancelled
}

func defaultGuardConfig() guardConfig {
	return guardConfig{Enabled: true, ActiveIdleS: 30, ReturnIdleS: 300, GraceS: 30}
}

type guardInput struct {
	Now        time.Time
	OwnerIdleS int64    // seconds since keyboard/mouse; -1 unknown (no desktop session)
	ForeignGPU []string // GPU compute processes that are not ours
	Paused     bool     // manual pause
	InFlight   int      // requests currently running on this machine
}

type guardDecision struct {
	Mode          guardMode
	Reason        string
	StopAccepting bool
	CancelAll     bool // cancel requests still running after the grace period
	UnloadAll     bool // free every loaded model (sent once per episode)
}

type ownerGuard struct {
	cfg         guardConfig
	mode        guardMode
	drainStart  time.Time
	lastForeign time.Time
	unloaded    bool
	kind        string // why the machine is not available: desktop | gpu | cooldown | paused ("" when available)
}

func newOwnerGuard(cfg guardConfig) *ownerGuard {
	return &ownerGuard{cfg: cfg, mode: modeAvailable}
}

func (g *ownerGuard) step(in guardInput) guardDecision {
	cfg := g.cfg
	// the foreign-GPU signal needs a quiet period of its own (it has no idle counter)
	if len(in.ForeignGPU) > 0 {
		g.lastForeign = in.Now
	}

	// manual pause beats everything
	if in.Paused {
		return g.pausedStep(in)
	}

	var signal, kind string
	if cfg.Enabled {
		switch {
		case in.OwnerIdleS >= 0 && in.OwnerIdleS < cfg.ActiveIdleS:
			signal = fmt.Sprintf("owner is using the machine (active %ds ago)", in.OwnerIdleS)
			kind = "desktop"
		case len(in.ForeignGPU) > 0:
			signal = "GPU in use by " + strings.Join(in.ForeignGPU, ", ")
			kind = "gpu"
		}
	}

	switch g.mode {
	case modeAvailable:
		if signal == "" {
			g.unloaded = false
			g.kind = ""
			return guardDecision{Mode: modeAvailable}
		}
		g.mode = modeDraining
		g.drainStart = in.Now
		g.kind = kind
		return g.drain(in, signal)

	case modeDraining:
		if kind != "" {
			g.kind = kind
		}
		return g.drain(in, signal)

	default: // owner or paused -> check whether we may come back
		if signal != "" {
			g.mode = modeOwner
			g.kind = kind
			return g.hold(signal)
		}
		quiet := in.OwnerIdleS < 0 || in.OwnerIdleS >= cfg.ReturnIdleS
		foreignQuiet := g.lastForeign.IsZero() || in.Now.Sub(g.lastForeign) >= time.Duration(cfg.ReturnIdleS)*time.Second
		if cfg.Enabled && (!quiet || !foreignQuiet) {
			g.mode = modeOwner
			g.kind = "cooldown"
			return g.hold("waiting for the owner to be away long enough")
		}
		g.mode = modeAvailable
		g.unloaded = false
		g.kind = ""
		return guardDecision{Mode: modeAvailable}
	}
}

// retryIn estimates how many seconds until the machine may be available again if the owner stays away
// (rounded up). 0 means unknown (another program is on the GPU, a manual pause) or not applicable.
func (g *ownerGuard) retryIn(in guardInput) int64 {
	if (g.mode != modeOwner && g.mode != modeDraining) || in.Paused || len(in.ForeignGPU) > 0 {
		return 0
	}
	quiet := time.Duration(g.cfg.ReturnIdleS) * time.Second
	var wait time.Duration
	if in.OwnerIdleS >= 0 && in.OwnerIdleS < g.cfg.ReturnIdleS {
		wait = time.Duration(g.cfg.ReturnIdleS-in.OwnerIdleS) * time.Second
	}
	if !g.lastForeign.IsZero() {
		if w := quiet - in.Now.Sub(g.lastForeign); w > wait {
			wait = w
		}
	}
	return int64((wait + time.Second - 1) / time.Second)
}

// drain stops new work and waits for running requests, up to the grace period.
func (g *ownerGuard) drain(in guardInput, reason string) guardDecision {
	graceOver := in.Now.Sub(g.drainStart) >= time.Duration(g.cfg.GraceS)*time.Second
	if in.InFlight == 0 || graceOver {
		g.mode = modeOwner
		d := g.hold(reason)
		d.CancelAll = in.InFlight > 0
		return d
	}
	return guardDecision{Mode: modeDraining, Reason: reason, StopAccepting: true}
}

// hold is the steady state of owner/paused mode; UnloadAll is requested only once.
func (g *ownerGuard) hold(reason string) guardDecision {
	d := guardDecision{Mode: g.mode, Reason: reason, StopAccepting: true}
	if !g.unloaded {
		d.UnloadAll = true
		g.unloaded = true
	}
	return d
}

// pausedStep handles a manual pause: same drain-then-unload as owner activity, mode "paused".
func (g *ownerGuard) pausedStep(in guardInput) guardDecision {
	const reason = "paused by the owner"
	g.kind = "paused"
	if g.mode == modeAvailable {
		g.mode = modeDraining
		g.drainStart = in.Now
	}
	if g.mode == modeDraining {
		graceOver := in.Now.Sub(g.drainStart) >= time.Duration(g.cfg.GraceS)*time.Second
		if in.InFlight > 0 && !graceOver {
			return guardDecision{Mode: modeDraining, Reason: reason, StopAccepting: true}
		}
		g.mode = modePaused
		d := g.hold(reason)
		d.CancelAll = in.InFlight > 0
		return d
	}
	g.mode = modePaused
	return g.hold(reason)
}

// ---- wiring to the running agent ----

var guardState = struct {
	sync.Mutex
	g      *ownerGuard
	last   guardDecision
	paused bool
}{}

func initGuard() {
	guardState.Lock()
	defer guardState.Unlock()
	cfg := defaultGuardConfig()
	nc := loadNodeConfig()
	if nc.OwnerGuardOff {
		cfg.Enabled = false
	}
	if nc.OwnerActiveIdleS > 0 {
		cfg.ActiveIdleS = int64(nc.OwnerActiveIdleS)
	}
	if nc.OwnerReturnIdleS > 0 {
		cfg.ReturnIdleS = int64(nc.OwnerReturnIdleS)
	}
	if nc.OwnerGraceS > 0 {
		cfg.GraceS = int64(nc.OwnerGraceS)
	}
	guardState.g = newOwnerGuard(cfg)
	guardState.last = guardDecision{Mode: modeAvailable}
}

// guardAccepting says whether new inference requests may start now.
func guardAccepting() (bool, string) {
	guardState.Lock()
	defer guardState.Unlock()
	if guardState.g == nil {
		return true, ""
	}
	return !guardState.last.StopAccepting, guardState.last.Reason
}

func guardSnapshot() (guardMode, string, bool) {
	guardState.Lock()
	defer guardState.Unlock()
	return guardState.last.Mode, guardState.last.Reason, guardState.paused
}

// guardWhy is the coarse reason the machine is not available ("" when it is) and the estimated seconds
// until it may be (0 = unknown). Unlike the reason text it names no program, so it is safe to show to other users.
func guardWhy(ownerIdle int64, foreign []string) (string, int64) {
	guardState.Lock()
	defer guardState.Unlock()
	if guardState.g == nil {
		return "", 0
	}
	in := guardInput{Now: time.Now(), OwnerIdleS: ownerIdle, ForeignGPU: foreign, Paused: guardState.paused}
	return guardState.g.kind, guardState.g.retryIn(in)
}

// guardEvaluate runs one step with fresh inputs and carries out the decision.
func guardEvaluate(ownerIdle int64, foreign []string) {
	guardState.Lock()
	if guardState.g == nil {
		guardState.Unlock()
		return
	}
	in := guardInput{Now: time.Now(), OwnerIdleS: ownerIdle, ForeignGPU: foreign, Paused: pausedFlag(), InFlight: thePool.totalInflight() + jobsRunning()}
	d := guardState.g.step(in)
	guardState.paused = in.Paused
	changed := d.Mode != guardState.last.Mode
	guardState.last = d
	guardState.Unlock()

	if changed {
		logf("owner guard: %s (%s)", d.Mode, d.Reason)
	}
	if d.CancelAll {
		cancelAllInfer()
		cancelAllJobs("the owner is using the machine")
	}
	if d.UnloadAll {
		go unloadAllModels()
	}
}

func unloadAllModels() {
	b := ollamaBackend{}
	loaded, err := b.loaded()
	if err != nil {
		return
	}
	for name := range loaded {
		b.unload(name)
	}
}

// isOurRuntime recognises the GPU worker processes of the runtimes we run ourselves:
// Ollama ("ollama"), LM Studio ("llama-server", "llmster"), vLLM ("VLLM::EngineCore", "VLLM::Worker").
// Treating them as the owner's work would make the guard unload the very model that is being served.
func isOurRuntime(name, cgroup string) bool {
	l := strings.ToLower(filepath.Base(name))
	switch {
	case strings.Contains(l, "ollama"),
		strings.Contains(l, "llmster"),
		strings.Contains(l, "llama") && strings.Contains(l, "server"),
		strings.HasPrefix(l, "vllm"),
		strings.Contains(l, "enginecore"):
		return true
	}
	cg := strings.ToLower(cgroup)
	return strings.Contains(cg, "ollama.service") || strings.Contains(cg, "lmstudio")
}

// otherGPUProcs returns the names of GPU compute processes that are NOT ours.
func otherGPUProcs(procs []gpuProc, cgroupOf func(pid int) string) []string {
	var out []string
	for _, p := range procs {
		cg := ""
		if cgroupOf != nil {
			cg = cgroupOf(p.PID)
		}
		if isOurRuntime(p.Name, cg) || inOurContainer(cg) || inOurHostEngine(p.PID) {
			continue
		}
		out = append(out, filepath.Base(p.Name))
	}
	return out
}

// procCgroup reads the cgroup of a process ("" if unreadable).
func procCgroup(pid int) string {
	b, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/cgroup")
	if err != nil {
		return ""
	}
	return string(b)
}

// inOurContainer says whether a GPU process lives in a container that we started ourselves (the AI module engines, the
// ComfyUI lab container). That is the machine serving other people, not its owner at work, so it must not make the guard
// stop the chat. The container ids come from the runner (it is the one that can ask Docker); without a runner nothing
// matches and the old behaviour stays.
func inOurContainer(cgroup string) bool {
	cg := strings.ToLower(cgroup)
	if !strings.Contains(cg, "docker") && !strings.Contains(cg, "containerd") && !strings.Contains(cg, "libpod") {
		return false
	}
	for _, id := range ourContainerIDs() {
		if id != "" && strings.Contains(cg, id) {
			return true
		}
	}
	return false
}

var ours = struct {
	sync.Mutex
	ids []string
	at  time.Time
}{}

// ourContainerIDsFn is replaced by tests.
var ourContainerIDsFn = func() []string {
	cfg := loadNodeConfig()
	if _, err := os.Stat(runnerSocket(cfg)); err != nil {
		return nil
	}
	d, err := (runnerControl{sock: runnerSocket(cfg)}).call(runnerRequest{Op: "ours"}, 10*time.Second)
	if err != nil {
		return nil
	}
	var ids []string
	json.Unmarshal(d, &ids)
	return ids
}

// ourContainerIDs asks the runner at most every few seconds (the guard looks at the GPU every tick).
func ourContainerIDs() []string {
	ours.Lock()
	defer ours.Unlock()
	if time.Since(ours.at) < 5*time.Second {
		return ours.ids
	}
	ours.ids, ours.at = ourContainerIDsFn(), time.Now()
	return ours.ids
}